Skip to content

Commit 099d6f1

Browse files
evakravievakravi
committed
fix: variants -> regional_variants, properties -> regional_properties
1 parent 9f912e4 commit 099d6f1

File tree

3 files changed

+34
-26
lines changed

3 files changed

+34
-26
lines changed

‎src/sagemaker/jumpstart/types.py

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -314,7 +314,7 @@ class JumpStartInstanceTypeVariants(JumpStartDataHolderType):
314314
"""Data class for JumpStart instance type variants."""
315315

316316
__slots__ = [
317-
"aliases",
317+
"regional_aliases",
318318
"variants",
319319
]
320320

@@ -336,7 +336,7 @@ def from_json(self, json_obj: Optional[Dict[str, Any]]) -> None:
336336
if json_obj is None:
337337
return
338338

339-
self.aliases: dict = json_obj["aliases"]
339+
self.regional_aliases: dict = json_obj["regional_aliases"]
340340
self.variants: dict = json_obj["variants"]
341341

342342
def to_json(self) -> Dict[str, Any]:
@@ -352,15 +352,15 @@ def get_image_uri(self, instance_type: str, region: str) -> Optional[str]:
352352

353353
image_uri_alias: Optional[str] = None
354354
if instance_type in self.variants:
355-
image_uri_alias = self.variants[instance_type]["properties"].get("image_uri")
355+
image_uri_alias = self.variants[instance_type]["regional_properties"].get("image_uri")
356356
else:
357357
instance_type_family = get_instance_type_family(instance_type)
358358

359359
if instance_type_family in {"", None}:
360360
return None
361361

362362
image_uri_alias = (
363-
self.variants[instance_type_family]["properties"].get("image_uri")
363+
self.variants[instance_type_family]["regional_properties"].get("image_uri")
364364
if instance_type_family in self.variants
365365
else None
366366
)
@@ -371,9 +371,9 @@ def get_image_uri(self, instance_type: str, region: str) -> Optional[str]:
371371
if not image_uri_alias.startswith("$"):
372372
raise TypeError("All image uris should map to an alias and start with '$'.")
373373

374-
if region not in self.aliases:
374+
if region not in self.regional_aliases:
375375
return None
376-
alias_value = self.aliases[region][image_uri_alias[1:]]
376+
alias_value = self.regional_aliases[region][image_uri_alias[1:]]
377377
return alias_value
378378

379379

‎tests/unit/sagemaker/jumpstart/constants.py

Lines changed: 13 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -27,25 +27,29 @@
2727
"py_version": "py3",
2828
},
2929
"hosting_instance_type_variants": {
30-
"aliases": {
30+
"regional_aliases": {
3131
"us-west-2": {
3232
"gpu_image_uri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
3333
"huggingface-pytorch-inference:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04",
3434
"cpu_image_uri": "867930986793.dkr.us-west-2.amazonaws.com/cpu-blah",
3535
}
3636
},
3737
"variants": {
38-
"p2": {"properties": {"image_uri": "$gpu_image_uri"}},
39-
"p3": {"properties": {"image_uri": "$gpu_image_uri"}},
40-
"p4": {"properties": {"image_uri": "$gpu_image_uri"}},
41-
"g4dn": {"properties": {"image_uri": "$gpu_image_uri"}},
42-
"m2": {"properties": {"image_uri": "$cpu_image_uri"}},
43-
"c2": {"properties": {"image_uri": "$cpu_image_uri"}},
38+
"p2": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
39+
"p3": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
40+
"p4": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
41+
"g4dn": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
42+
"m2": {"regional_properties": {"image_uri": "$cpu_image_uri"}},
43+
"c2": {"regional_properties": {"image_uri": "$cpu_image_uri"}},
4444
"ml.g5.48xlarge": {
45-
"properties": {"environment_variables": {"TENSOR_PARALLEL_DEGREE": "8"}}
45+
"regional_properties": {
46+
"environment_variables": {"TENSOR_PARALLEL_DEGREE": "8"}
47+
}
4648
},
4749
"ml.g5.12xlarge": {
48-
"properties": {"environment_variables": {"TENSOR_PARALLEL_DEGREE": "4"}}
50+
"regional_properties": {
51+
"environment_variables": {"TENSOR_PARALLEL_DEGREE": "4"}
52+
}
4953
},
5054
},
5155
},

‎tests/unit/sagemaker/jumpstart/test_types.py

Lines changed: 15 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ def test_jumpstart_model_specs():
161161
def test_jumpstart_instance_variants():
162162
variant = JumpStartInstanceTypeVariants(
163163
{
164-
"aliases": {
164+
"regional_aliases": {
165165
"us-west-2": {
166166
"gpu_image_uri": "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
167167
"huggingface-pytorch-inference:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04",
@@ -170,19 +170,23 @@ def test_jumpstart_instance_variants():
170170
}
171171
},
172172
"variants": {
173-
"p2": {"properties": {"image_uri": "$gpu_image_uri"}},
174-
"p3": {"properties": {"image_uri": "$gpu_image_uri"}},
175-
"ml.p3.200xlarge": {"properties": {"image_uri": "$gpu_image_uri_2"}},
176-
"p4": {"properties": {"image_uri": "$gpu_image_uri"}},
177-
"g4dn": {"properties": {"image_uri": "$gpu_image_uri"}},
178-
"m2": {"properties": {"image_uri": "$cpu_image_uri"}},
179-
"c2": {"properties": {"image_uri": "$cpu_image_uri"}},
180-
"local": {"properties": {"image_uri": "$cpu_image_uri"}},
173+
"p2": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
174+
"p3": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
175+
"ml.p3.200xlarge": {"regional_properties": {"image_uri": "$gpu_image_uri_2"}},
176+
"p4": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
177+
"g4dn": {"regional_properties": {"image_uri": "$gpu_image_uri"}},
178+
"m2": {"regional_properties": {"image_uri": "$cpu_image_uri"}},
179+
"c2": {"regional_properties": {"image_uri": "$cpu_image_uri"}},
180+
"local": {"regional_properties": {"image_uri": "$cpu_image_uri"}},
181181
"ml.g5.48xlarge": {
182-
"properties": {"environment_variables": {"TENSOR_PARALLEL_DEGREE": "8"}}
182+
"regional_properties": {
183+
"environment_variables": {"TENSOR_PARALLEL_DEGREE": "8"}
184+
}
183185
},
184186
"ml.g5.12xlarge": {
185-
"properties": {"environment_variables": {"TENSOR_PARALLEL_DEGREE": "4"}}
187+
"regional_properties": {
188+
"environment_variables": {"TENSOR_PARALLEL_DEGREE": "4"}
189+
}
186190
},
187191
},
188192
}

0 commit comments

Comments
 (0)