Skip to content

Commit 2a6f1b1

Browse files
JingyaHuangJingyaHuang
committed
add neuronx uri
1 parent 3a8a2e7 commit 2a6f1b1

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
@@ -405,16 +405,46 @@ def huggingface_neuron_latest_inference_pytorch_version():
405405
return "1.9"
406406

407407

408+
@pytest.fixture(scope="module")
409+
def huggingface_neuronx_latest_inference_pytorch_version():
410+
return "1.13"
411+
412+
413+
@pytest.fixture(scope="module")
414+
def huggingface_neuronx_latest_training_pytorch_version():
415+
return "1.13"
416+
417+
408418
@pytest.fixture(scope="module")
409419
def huggingface_neuron_latest_inference_transformer_version():
410420
return "4.12"
411421

412422

423+
@pytest.fixture(scope="module")
424+
def huggingface_neuronx_latest_inference_transformer_version():
425+
return "4.28"
426+
427+
428+
@pytest.fixture(scope="module")
429+
def huggingface_neuronx_latest_training_transformer_version():
430+
return "4.28"
431+
432+
413433
@pytest.fixture(scope="module")
414434
def huggingface_neuron_latest_inference_py_version():
415435
return "py37"
416436

417437

438+
@pytest.fixture(scope="module")
439+
def huggingface_neuronx_latest_inference_py_version():
440+
return "py38"
441+
442+
443+
@pytest.fixture(scope="module")
444+
def huggingface_neuronx_latest_training_py_version():
445+
return "py38"
446+
447+
418448
@pytest.fixture(scope="module")
419449
def pytorch_neuron_version():
420450
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
@@ -268,7 +268,27 @@ def test_huggingface_neuron(
268268
pytorch_version=huggingface_neuron_latest_inference_pytorch_version,
269269
py_version=huggingface_neuron_latest_inference_py_version,
270270
)
271-
container = huggingface_model.prepare_container_def("ml.inf.xlarge")
271+
container = huggingface_model.prepare_container_def("ml.inf1.xlarge")
272+
assert container["Image"]
273+
274+
275+
def test_huggingface_neuronx(
276+
sagemaker_session,
277+
huggingface_neuronx_latest_inference_pytorch_version,
278+
huggingface_neuronx_latest_inference_transformer_version,
279+
huggingface_neuronx_latest_inference_py_version,
280+
):
281+
282+
inputs = "s3://mybucket/train"
283+
huggingface_model = HuggingFaceModel(
284+
model_data=inputs,
285+
transformers_version=huggingface_neuronx_latest_inference_transformer_version,
286+
role=ROLE,
287+
sagemaker_session=sagemaker_session,
288+
pytorch_version=huggingface_neuronx_latest_inference_pytorch_version,
289+
py_version=huggingface_neuronx_latest_inference_py_version,
290+
)
291+
container = huggingface_model.prepare_container_def("ml.inf2.xlarge")
272292
assert container["Image"]
273293

274294

0 commit comments

Comments
 (0)