Skip to content

Commit baccba5

Browse files
jeniyatjeniyat
committed
update w/ seperate json files
1 parent b92e9c5 commit baccba5

File tree

3 files changed

+72
-8
lines changed

3 files changed

+72
-8
lines changed

‎src/sagemaker/image_uri_config/huggingface-neuron.json

+49
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,49 @@
1+
{
2+
"inference": {
3+
"processors": ["neuron"],
4+
"version_aliases": {
5+
"4.12": "4.12.3"
6+
},
7+
"versions": {
8+
"4.12.3": {
9+
"version_aliases": {
10+
"pytorch1.9": "pytorch1.9.1"
11+
},
12+
"pytorch1.9.1": {
13+
"py_versions": ["py37"],
14+
"repository": "huggingface-pytorch-inference",
15+
"registries": {
16+
"af-south-1": "626614931356",
17+
"ap-east-1": "871362719292",
18+
"ap-northeast-1": "763104351884",
19+
"ap-northeast-2": "763104351884",
20+
"ap-northeast-3": "364406365360",
21+
"ap-south-1": "763104351884",
22+
"ap-southeast-1": "763104351884",
23+
"ap-southeast-2": "763104351884",
24+
"ca-central-1": "763104351884",
25+
"cn-north-1": "727897471807",
26+
"cn-northwest-1": "727897471807",
27+
"eu-central-1": "763104351884",
28+
"eu-north-1": "763104351884",
29+
"eu-west-1": "763104351884",
30+
"eu-west-2": "763104351884",
31+
"eu-west-3": "763104351884",
32+
"eu-south-1": "692866216735",
33+
"me-south-1": "217643126080",
34+
"sa-east-1": "763104351884",
35+
"us-east-1": "763104351884",
36+
"us-east-2": "763104351884",
37+
"us-gov-west-1": "442386744353",
38+
"us-iso-east-1": "886529160074",
39+
"us-west-1": "763104351884",
40+
"us-west-2": "763104351884"
41+
},
42+
"container_version": {"neuron": "ubuntu18.04"},
43+
"repo_versions": ["v1.0"],
44+
"sdk_versions": ["sdk1.17.1"]
45+
}
46+
}
47+
}
48+
}
49+
}

‎src/sagemaker/image_uri_config/huggingface.json

+2-2
Original file line numberDiff line numberDiff line change
@@ -562,7 +562,7 @@
562562
},
563563

564564
"inference": {
565-
"processors": ["gpu", "cpu", "neuron"],
565+
"processors": ["gpu", "cpu"],
566566
"version_aliases": {
567567
"4.6": "4.6.1",
568568
"4.10": "4.10.2",
@@ -913,7 +913,7 @@
913913
"us-west-2": "763104351884"
914914
},
915915
"repository": "huggingface-pytorch-inference",
916-
"container_version": {"gpu": "cu111-ubuntu20.04", "cpu": "ubuntu20.04", "neuron": "sdk1.17.1-ubuntu18.04-v1.0" }
916+
"container_version": {"gpu": "cu111-ubuntu20.04", "cpu": "ubuntu20.04"}
917917
},
918918
"tensorflow2.5.1": {
919919
"py_versions": ["py37"],

‎src/sagemaker/image_uris.py

+21-6
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
from sagemaker.spark import defaults
2525
from sagemaker.jumpstart import artifacts
2626

27-
2827
logger = logging.getLogger(__name__)
2928

3029
ECR_URI_TEMPLATE = "{registry}.dkr.{hostname}/{repository}"
@@ -100,7 +99,6 @@ def retrieve(
10099
DeprecatedJumpStartModelError: If the version of the model is deprecated.
101100
"""
102101
if is_jumpstart_model_input(model_id, model_version):
103-
104102
return artifacts._retrieve_image_uri(
105103
model_id,
106104
model_version,
@@ -118,17 +116,22 @@ def retrieve(
118116
tolerate_vulnerable_model,
119117
tolerate_deprecated_model,
120118
)
121-
122119
if training_compiler_config is None:
123-
config = _config_for_framework_and_scope(framework, image_scope, accelerator_type)
120+
if framework == HUGGING_FACE_FRAMEWORK and instance_type == "neuron":
121+
config = _config_for_framework_and_scope(
122+
framework + "-neuron", image_scope, accelerator_type
123+
)
124+
else:
125+
config = _config_for_framework_and_scope(framework, image_scope, accelerator_type)
124126
elif framework == HUGGING_FACE_FRAMEWORK:
125127
config = _config_for_framework_and_scope(
126-
framework + "-training-compiler", image_scope, accelerator_type
128+
framework + "-training-compiler", image_scope, accelerator_type, instance_type
127129
)
128130
else:
129131
raise ValueError(
130132
"Unsupported Configuration: Training Compiler is only supported with HuggingFace"
131133
)
134+
132135
original_version = version
133136
version = _validate_version_and_set_if_needed(version, config, framework)
134137
version_config = config["versions"][_version_for_config(version, config)]
@@ -169,6 +172,9 @@ def retrieve(
169172
]:
170173
_version = version
171174
if processor == "neuron":
175+
sdk_version = _get_latest_versions(version_config["sdk_versions"])
176+
repo_versions = _get_latest_versions(version_config["repo_versions"])
177+
container_version = sdk_version + "-" + container_version + "-" + repo_versions
172178
repo += "-{0}".format(processor)
173179

174180
tag_prefix = f"{pt_or_tf_version}-transformers{_version}"
@@ -208,8 +214,12 @@ def retrieve(
208214
return ECR_URI_TEMPLATE.format(registry=registry, hostname=hostname, repository=repo)
209215

210216

211-
def _config_for_framework_and_scope(framework, image_scope, accelerator_type=None):
217+
def _config_for_framework_and_scope(
218+
framework, image_scope, accelerator_type=None, instance_type=None
219+
):
212220
"""Loads the JSON config for the given framework and image scope."""
221+
if framework == HUGGING_FACE_FRAMEWORK and instance_type == "neuron":
222+
framework = framework + "_" + instance_type
213223
config = config_for_framework(framework)
214224

215225
if accelerator_type:
@@ -250,6 +260,11 @@ def config_for_framework(framework):
250260
return json.load(f)
251261

252262

263+
def _get_latest_versions(list_of_versions):
264+
"""Raises a ``ValueError`` if ``accelerator_type`` is invalid."""
265+
return sorted(list_of_versions, reverse=True)[0]
266+
267+
253268
def _validate_accelerator_type(accelerator_type):
254269
"""Raises a ``ValueError`` if ``accelerator_type`` is invalid."""
255270
if not accelerator_type.startswith("ml.eia") and accelerator_type != "local_sagemaker_notebook":

0 commit comments

Comments
 (0)