Skip to content

Commit 0309bc0

Browse files
cdesiniotiscdesiniotis
authored
Merge pull request #64 from varunrsekar/new-mig-profiles
support for new mig profiles
2 parents ae5e6af + fe2f325 commit 0309bc0

2 files changed

Lines changed: 158 additions & 25 deletions

File tree

‎pkg/nvlib/device/mig_profile.go‎

Lines changed: 82 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,9 @@ import (
2828

2929
const (
3030
// AttributeMediaExtensions holds the string representation for the media extension MIG profile attribute.
31-
AttributeMediaExtensions = "me"
31+
AttributeMediaExtensions = "me"
32+
AttributeMediaExtensionsAll = "me.all"
33+
AttributeGraphics = "gfx"
3234
)
3335

3436
// MigProfile represents a specific MIG profile.
@@ -46,6 +48,7 @@ type MigProfileInfo struct {
4648
G int
4749
GB int
4850
Attributes []string
51+
NegAttributes []string
4952
GIProfileID int
5053
CIProfileID int
5154
CIEngProfileID int
@@ -59,14 +62,21 @@ func (d *devicelib) NewMigProfile(giProfileID, ciProfileID, ciEngProfileID int,
5962
switch giProfileID {
6063
case nvml.GPU_INSTANCE_PROFILE_1_SLICE,
6164
nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1,
62-
nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV2:
65+
nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV2,
66+
nvml.GPU_INSTANCE_PROFILE_1_SLICE_GFX,
67+
nvml.GPU_INSTANCE_PROFILE_1_SLICE_NO_ME,
68+
nvml.GPU_INSTANCE_PROFILE_1_SLICE_ALL_ME:
6369
giSlices = 1
6470
case nvml.GPU_INSTANCE_PROFILE_2_SLICE,
65-
nvml.GPU_INSTANCE_PROFILE_2_SLICE_REV1:
71+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_REV1,
72+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_GFX,
73+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_NO_ME,
74+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_ALL_ME:
6675
giSlices = 2
6776
case nvml.GPU_INSTANCE_PROFILE_3_SLICE:
6877
giSlices = 3
69-
case nvml.GPU_INSTANCE_PROFILE_4_SLICE:
78+
case nvml.GPU_INSTANCE_PROFILE_4_SLICE,
79+
nvml.GPU_INSTANCE_PROFILE_4_SLICE_GFX:
7080
giSlices = 4
7181
case nvml.GPU_INSTANCE_PROFILE_6_SLICE:
7282
giSlices = 6
@@ -104,13 +114,27 @@ func (d *devicelib) NewMigProfile(giProfileID, ciProfileID, ciEngProfileID int,
104114
case nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1,
105115
nvml.GPU_INSTANCE_PROFILE_2_SLICE_REV1:
106116
attrs = append(attrs, AttributeMediaExtensions)
117+
case nvml.GPU_INSTANCE_PROFILE_1_SLICE_ALL_ME,
118+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_ALL_ME:
119+
attrs = append(attrs, AttributeMediaExtensionsAll)
120+
case nvml.GPU_INSTANCE_PROFILE_1_SLICE_GFX,
121+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_GFX,
122+
nvml.GPU_INSTANCE_PROFILE_4_SLICE_GFX:
123+
attrs = append(attrs, AttributeGraphics)
124+
}
125+
var negAttrs []string
126+
switch giProfileID {
127+
case nvml.GPU_INSTANCE_PROFILE_1_SLICE_NO_ME,
128+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_NO_ME:
129+
negAttrs = append(negAttrs, AttributeMediaExtensions)
107130
}
108131

109132
p := &MigProfileInfo{
110133
C: ciSlices,
111134
G: giSlices,
112135
GB: int(getMigMemorySizeGB(deviceMemorySizeBytes, migMemorySizeMB)),
113136
Attributes: attrs,
137+
NegAttributes: negAttrs,
114138
GIProfileID: giProfileID,
115139
CIProfileID: ciProfileID,
116140
CIEngProfileID: ciEngProfileID,
@@ -121,7 +145,7 @@ func (d *devicelib) NewMigProfile(giProfileID, ciProfileID, ciEngProfileID int,
121145

122146
// AssertValidMigProfileFormat checks if the string is in the proper format to represent a MIG profile.
123147
func (d *devicelib) AssertValidMigProfileFormat(profile string) error {
124-
_, _, _, _, err := parseMigProfile(profile)
148+
_, err := parseMigProfile(profile)
125149
return err
126150
}
127151

@@ -147,6 +171,9 @@ func (p MigProfileInfo) String() string {
147171
if len(p.Attributes) > 0 {
148172
suffix = "+" + strings.Join(p.Attributes, ",")
149173
}
174+
if len(p.NegAttributes) > 0 {
175+
suffix = "-" + strings.Join(p.NegAttributes, ",")
176+
}
150177
if p.C == p.G {
151178
return fmt.Sprintf("%dg.%dgb%s", p.G, p.GB, suffix)
152179
}
@@ -184,59 +211,89 @@ func (p MigProfileInfo) Equals(other MigProfile) bool {
184211

185212
// Matches checks if a MigProfile matches the string passed in.
186213
func (p MigProfileInfo) Matches(profile string) bool {
187-
c, g, gb, attrs, err := parseMigProfile(profile)
214+
migProfileInfo, err := parseMigProfile(profile)
188215
if err != nil {
189216
return false
190217
}
191-
if c != p.C {
218+
if migProfileInfo.C != p.C {
219+
return false
220+
}
221+
if migProfileInfo.G != p.G {
192222
return false
193223
}
194-
if g != p.G {
224+
if migProfileInfo.GB != p.GB {
195225
return false
196226
}
197-
if gb != p.GB {
227+
if !matchAttributes(migProfileInfo.Attributes, p.Attributes) {
198228
return false
199229
}
200-
if len(attrs) != len(p.Attributes) {
230+
if !matchAttributes(migProfileInfo.NegAttributes, p.NegAttributes) {
201231
return false
202232
}
203-
sort.Strings(attrs)
204-
sort.Strings(p.Attributes)
205-
for i, a := range p.Attributes {
206-
if a != attrs[i] {
233+
return true
234+
}
235+
236+
func matchAttributes(attrs1, attrs2 []string) bool {
237+
if len(attrs1) != len(attrs2) {
238+
return false
239+
}
240+
sort.Strings(attrs1)
241+
sort.Strings(attrs2)
242+
for i, a := range attrs2 {
243+
if a != attrs1[i] {
207244
return false
208245
}
209246
}
210247
return true
211248
}
212249

213-
func parseMigProfile(profile string) (int, int, int, []string, error) {
250+
func parseMigProfile(profile string) (*MigProfileInfo, error) {
214251
// If we are handed the empty string, we cannot parse it.
215252
if profile == "" {
216-
return -1, -1, -1, nil, fmt.Errorf("profile is the empty string")
253+
return nil, fmt.Errorf("profile is the empty string")
217254
}
218255

219-
// Split by + to separate out attributes.
256+
// Split by +/- to separate out attributes.
220257
split := strings.SplitN(profile, "+", 2)
258+
negsplit := strings.SplitN(profile, "-", 2)
259+
// Make sure we don't get both positive and negative attributes.
260+
if len(split) == 2 && len(negsplit) == 2 {
261+
return nil, fmt.Errorf("profile '%v' contains both '+/-' attributes", profile)
262+
}
263+
264+
if len(split) == 1 {
265+
split = negsplit
266+
}
221267

222268
// Check to make sure the c, g, and gb values match.
223269
c, g, gb, err := parseMigProfileFields(split[0])
224270
if err != nil {
225-
return -1, -1, -1, nil, fmt.Errorf("cannot parse fields of '%v': %v", profile, err)
271+
return nil, fmt.Errorf("cannot parse fields of '%v': %v", profile, err)
226272
}
227273

274+
migProfileInfo := &MigProfileInfo{
275+
C: c,
276+
G: g,
277+
GB: gb,
278+
}
228279
// If we have no attributes we are done.
229280
if len(split) == 1 {
230-
return c, g, gb, nil, nil
281+
return migProfileInfo, nil
231282
}
232283

233284
// Make sure we have the same set of attributes.
234285
attrs, err := parseMigProfileAttributes(split[1])
235286
if err != nil {
236-
return -1, -1, -1, nil, fmt.Errorf("cannot parse attributes of '%v': %v", profile, err)
287+
return nil, fmt.Errorf("cannot parse attributes of '%v': %v", profile, err)
288+
}
289+
290+
if len(negsplit) == 2 {
291+
migProfileInfo.NegAttributes = attrs
292+
return migProfileInfo, nil
237293
}
238294

239-
return c, g, gb, attrs, nil
295+
migProfileInfo.Attributes = attrs
296+
return migProfileInfo, nil
240297
}
241298

242299
func parseMigProfileField(s string, field string) (int, error) {
@@ -310,8 +367,11 @@ func parseMigProfileAttributes(s string) ([]string, error) {
310367
if a[0] >= '0' && a[0] <= '9' {
311368
return nil, fmt.Errorf("attribute begins with a number")
312369
}
370+
if a[0] == '.' || a[len(a)-1] == '.' {
371+
return nil, fmt.Errorf("attribute begins/ends with a dot")
372+
}
313373
for _, c := range a {
314-
if (c < 'a' || c > 'z') && (c < 'A' || c > 'Z') && (c < '0' || c > '9') {
374+
if (c < 'a' || c > 'z') && (c < 'A' || c > 'Z') && (c < '0' || c > '9') && c != '.' {
315375
return nil, fmt.Errorf("non alpha-numeric character or digit in attribute")
316376
}
317377
}

‎pkg/nvlib/device/mig_profile_test.go‎

Lines changed: 76 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -48,16 +48,23 @@ func newMockDeviceLib() Interface {
4848
info := nvml.GpuInstanceProfileInfo{}
4949
switch Profile {
5050
case nvml.GPU_INSTANCE_PROFILE_1_SLICE,
51-
nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1:
51+
nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV1,
52+
nvml.GPU_INSTANCE_PROFILE_1_SLICE_GFX,
53+
nvml.GPU_INSTANCE_PROFILE_1_SLICE_NO_ME,
54+
nvml.GPU_INSTANCE_PROFILE_1_SLICE_ALL_ME:
5255
info.MemorySizeMB = 5 * 1024
5356
case nvml.GPU_INSTANCE_PROFILE_1_SLICE_REV2:
5457
info.MemorySizeMB = 10 * 1024
5558
case nvml.GPU_INSTANCE_PROFILE_2_SLICE,
56-
nvml.GPU_INSTANCE_PROFILE_2_SLICE_REV1:
59+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_REV1,
60+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_GFX,
61+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_NO_ME,
62+
nvml.GPU_INSTANCE_PROFILE_2_SLICE_ALL_ME:
5763
info.MemorySizeMB = 10 * 1024
5864
case nvml.GPU_INSTANCE_PROFILE_3_SLICE:
5965
info.MemorySizeMB = 20 * 1024
60-
case nvml.GPU_INSTANCE_PROFILE_4_SLICE:
66+
case nvml.GPU_INSTANCE_PROFILE_4_SLICE,
67+
nvml.GPU_INSTANCE_PROFILE_4_SLICE_GFX:
6168
info.MemorySizeMB = 20 * 1024
6269
case nvml.GPU_INSTANCE_PROFILE_7_SLICE:
6370
info.MemorySizeMB = 40 * 1024
@@ -119,6 +126,36 @@ func TestParseMigProfile(t *testing.T) {
119126
true,
120127
true,
121128
},
129+
{
130+
"Valid 1c.1g.5gb-me",
131+
"1c.1g.5gb-me",
132+
true,
133+
true,
134+
},
135+
{
136+
"Valid 1g.5gb+me.all",
137+
"1g.5gb+me.all",
138+
true,
139+
true,
140+
},
141+
{
142+
"Valid 1c.1g.5gb+me.all",
143+
"1c.1g.5gb+me.all",
144+
true,
145+
true,
146+
},
147+
{
148+
"Valid 1g.5gb+gfx",
149+
"1g.5gb+gfx",
150+
true,
151+
true,
152+
},
153+
{
154+
"Valid 1c.1g.5gb+gfx",
155+
"1c.1g.5gb+gfx",
156+
true,
157+
true,
158+
},
122159
{
123160
"Invalid 0g.0gb",
124161
"0g.0gb",
@@ -347,6 +384,30 @@ func TestParseMigProfile(t *testing.T) {
347384
false,
348385
false,
349386
},
387+
{
388+
"Invalid 1c.1g.5gb+.all",
389+
"1c.1g.5gb+.all",
390+
false,
391+
false,
392+
},
393+
{
394+
"Invalid 1c.1g.5gb+all.",
395+
"1c.1g.5gb+all.",
396+
false,
397+
false,
398+
},
399+
{
400+
"Invalid 1c.1g.5gb+me-me",
401+
"1c.1g.5gb+me-me",
402+
false,
403+
false,
404+
},
405+
{
406+
"Invalid 1c.1g.5gb-me+me",
407+
"1c.1g.5gb-me+me",
408+
false,
409+
false,
410+
},
350411
}
351412

352413
d := newMockDeviceLib()
@@ -411,6 +472,18 @@ func TestParseMigProfileEquals(t *testing.T) {
411472
"1g.5gb+me",
412473
false,
413474
},
475+
{
476+
"Not equal neg attributes",
477+
"1g.5gb",
478+
"1g.5gb-me",
479+
false,
480+
},
481+
{
482+
"Not equal +/- attributes",
483+
"1g.5gb+me",
484+
"1g.5gb-me",
485+
false,
486+
},
414487
}
415488

416489
d := newMockDeviceLib()

0 commit comments

Comments
 (0)