Skip to content

Commit 71c20dc

Browse files
JingyaHuangtrajanikant
JingyaHuang
authored and
trajanikant
committed
add neuronx uri
1 parent abd94ab commit 71c20dc

File tree

4 files changed

+166
-1
lines changed

4 files changed

+166
-1
lines changed

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

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
1+
{
2+
"training": {
3+
"processors": ["trn"],
4+
"version_aliases": {"4.28": "4.28.1"},
5+
"versions": {
6+
"4.28.1": {
7+
"version_aliases": {"pytorch1.13": "pytorch1.13.0"},
8+
"pytorch1.13.0": {
9+
"py_versions": ["py38"],
10+
"repository": "huggingface-pytorch-training-neuronx",
11+
"registries": {
12+
"af-south-1": "626614931356",
13+
"ap-east-1": "871362719292",
14+
"ap-northeast-1": "763104351884",
15+
"ap-northeast-2": "763104351884",
16+
"ap-northeast-3": "364406365360",
17+
"ap-south-1": "763104351884",
18+
"ap-south-2": "772153158452",
19+
"ap-southeast-1": "763104351884",
20+
"ap-southeast-2": "763104351884",
21+
"ap-southeast-4": "457447274322",
22+
"ca-central-1": "763104351884",
23+
"cn-north-1": "727897471807",
24+
"cn-northwest-1": "727897471807",
25+
"eu-central-1": "763104351884",
26+
"eu-central-2": "380420809688",
27+
"eu-north-1": "763104351884",
28+
"eu-west-1": "763104351884",
29+
"eu-west-2": "763104351884",
30+
"eu-west-3": "763104351884",
31+
"eu-south-1": "692866216735",
32+
"eu-south-2": "503227376785",
33+
"me-south-1": "217643126080",
34+
"sa-east-1": "763104351884",
35+
"us-east-1": "763104351884",
36+
"us-east-2": "763104351884",
37+
"us-gov-east-1": "446045086412",
38+
"us-gov-west-1": "442386744353",
39+
"us-iso-east-1": "886529160074",
40+
"us-isob-east-1": "094389454867",
41+
"us-west-1": "763104351884",
42+
"us-west-2": "763104351884"
43+
},
44+
"container_version": {"trn": "ubuntu20.04"},
45+
"sdk_versions": ["sdk2.9.1"]
46+
}
47+
}
48+
}
49+
},
50+
"inference": {
51+
"processors": ["inf"],
52+
"version_aliases": {"4.28": "4.28.1"},
53+
"versions": {
54+
"4.28.1": {
55+
"version_aliases": {"pytorch1.13": "pytorch1.13.0"},
56+
"pytorch1.13.0": {
57+
"py_versions": ["py38"],
58+
"repository": "huggingface-pytorch-inference-neuronx",
59+
"registries": {
60+
"af-south-1": "626614931356",
61+
"ap-east-1": "871362719292",
62+
"ap-northeast-1": "763104351884",
63+
"ap-northeast-2": "763104351884",
64+
"ap-northeast-3": "364406365360",
65+
"ap-south-1": "763104351884",
66+
"ap-south-2": "772153158452",
67+
"ap-southeast-1": "763104351884",
68+
"ap-southeast-2": "763104351884",
69+
"ap-southeast-4": "457447274322",
70+
"ca-central-1": "763104351884",
71+
"cn-north-1": "727897471807",
72+
"cn-northwest-1": "727897471807",
73+
"eu-central-1": "763104351884",
74+
"eu-central-2": "380420809688",
75+
"eu-north-1": "763104351884",
76+
"eu-west-1": "763104351884",
77+
"eu-west-2": "763104351884",
78+
"eu-west-3": "763104351884",
79+
"eu-south-1": "692866216735",
80+
"eu-south-2": "503227376785",
81+
"me-south-1": "217643126080",
82+
"sa-east-1": "763104351884",
83+
"us-east-1": "763104351884",
84+
"us-east-2": "763104351884",
85+
"us-gov-east-1": "446045086412",
86+
"us-gov-west-1": "442386744353",
87+
"us-iso-east-1": "886529160074",
88+
"us-isob-east-1": "094389454867",
89+
"us-west-1": "763104351884",
90+
"us-west-2": "763104351884"
91+
},
92+
"container_version": {"inf": "ubuntu20.04"},
93+
"sdk_versions": ["sdk2.9.1"]
94+
}
95+
}
96+
}
97+
}
98+
}

‎tests/conftest.py

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -407,16 +407,46 @@ def huggingface_neuron_latest_inference_pytorch_version():
407407
return "1.9"
408408

409409

410+
@pytest.fixture(scope="module")
411+
def huggingface_neuronx_latest_inference_pytorch_version():
412+
return "1.13"
413+
414+
415+
@pytest.fixture(scope="module")
416+
def huggingface_neuronx_latest_training_pytorch_version():
417+
return "1.13"
418+
419+
410420
@pytest.fixture(scope="module")
411421
def huggingface_neuron_latest_inference_transformer_version():
412422
return "4.12"
413423

414424

425+
@pytest.fixture(scope="module")
426+
def huggingface_neuronx_latest_inference_transformer_version():
427+
return "4.28"
428+
429+
430+
@pytest.fixture(scope="module")
431+
def huggingface_neuronx_latest_training_transformer_version():
432+
return "4.28"
433+
434+
415435
@pytest.fixture(scope="module")
416436
def huggingface_neuron_latest_inference_py_version():
417437
return "py37"
418438

419439

440+
@pytest.fixture(scope="module")
441+
def huggingface_neuronx_latest_inference_py_version():
442+
return "py38"
443+
444+
445+
@pytest.fixture(scope="module")
446+
def huggingface_neuronx_latest_training_py_version():
447+
return "py38"
448+
449+
420450
@pytest.fixture(scope="module")
421451
def pytorch_neuron_version():
422452
return "1.11"

‎tests/unit/sagemaker/huggingface/huggingface_utils.py

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
REGION = "us-east-1"
1818
GPU_INSTANCE_TYPE = "ml.p2.xlarge"
19+
NEURONX_INSTANCE_TYPE = "ml.trn1.2xlarge"
1920

2021

2122
def get_full_gpu_image_uri(
@@ -33,3 +34,19 @@ def get_full_gpu_image_uri(
3334
base_framework_version=base_framework_version,
3435
container_version="cu110-ubuntu18.04",
3536
)
37+
38+
def get_full_neuronx_image_uri(
39+
version,
40+
base_framework_version,
41+
region=REGION,
42+
instance_type=NEURONX_INSTANCE_TYPE,
43+
):
44+
return image_uris.retrieve(
45+
"huggingface",
46+
region,
47+
version=version,
48+
instance_type=instance_type,
49+
image_scope="training",
50+
base_framework_version=base_framework_version,
51+
container_version="cu110-ubuntu18.04",
52+
)

‎tests/unit/sagemaker/huggingface/test_estimator.py

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -269,7 +269,27 @@ def test_huggingface_neuron(
269269
pytorch_version=huggingface_neuron_latest_inference_pytorch_version,
270270
py_version=huggingface_neuron_latest_inference_py_version,
271271
)
272-
container = huggingface_model.prepare_container_def("ml.inf.xlarge")
272+
container = huggingface_model.prepare_container_def("ml.inf1.xlarge")
273+
assert container["Image"]
274+
275+
276+
def test_huggingface_neuronx(
277+
sagemaker_session,
278+
huggingface_neuronx_latest_inference_pytorch_version,
279+
huggingface_neuronx_latest_inference_transformer_version,
280+
huggingface_neuronx_latest_inference_py_version,
281+
):
282+
283+
inputs = "s3://mybucket/train"
284+
huggingface_model = HuggingFaceModel(
285+
model_data=inputs,
286+
transformers_version=huggingface_neuronx_latest_inference_transformer_version,
287+
role=ROLE,
288+
sagemaker_session=sagemaker_session,
289+
pytorch_version=huggingface_neuronx_latest_inference_pytorch_version,
290+
py_version=huggingface_neuronx_latest_inference_py_version,
291+
)
292+
container = huggingface_model.prepare_container_def("ml.inf2.xlarge")
273293
assert container["Image"]
274294

275295

0 commit comments

Comments
 (0)