Neuron URIs update (#1626)

Resolve recipes correctly before launching (#1529)
fixes. (#1532)
fix recipe path. (#1566)

Author: schinmayee

src/sagemaker/image_uri_config/hyperpod-recipes-neuron.json

@@ -0,0 +1,52 @@
+{
+    "training": {
+        "processors": [
+            "neuronx"
+        ],
+        "version_aliases": {
+            "2.1.2": "2.1.2"
+        },
+        "versions": {
+            "2.1.2": {
+                "py_versions": [
+                    "py310"
+                ],
+                "repository": "pytorch-training-neuronx",
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "ap-east-1": "871362719292",
+                    "ap-northeast-1": "763104351884",
+                    "ap-northeast-2": "763104351884",
+                    "ap-northeast-3": "364406365360",
+                    "ap-south-1": "763104351884",
+                    "ap-south-2": "772153158452",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-3": "907027046896",
+                    "ap-southeast-4": "457447274322",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-north-1": "763104351884",
+                    "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884",
+                    "ca-west-1": "204538143572",
+                    "ca-central-1": "763104351884"
+                },
+                "container_version": {
+                    "neuronx": "ubuntu20.04"
+                },
+                "sdk_versions": [
+                    "sdk2.20.2"
+                ]
+            }
+        }
+    }
+}

src/sagemaker/image_uris.py

@@ -224,7 +224,7 @@ def retrieve(
         container_version = version_config["container_version"][processor]
 
     # Append sdk version in case of trainium instances
-    if repo in ["pytorch-training-neuron"]:
+    if repo in ["pytorch-training-neuron", "pytorch-training-neuronx"]:
         if not sdk_version:
             sdk_version = _get_latest_versions(version_config["sdk_versions"])
         container_version = sdk_version + "-" + container_version

src/sagemaker/pytorch/estimator.py

@@ -177,7 +177,7 @@ def __init__(
                 a directory with any other training source code dependencies aside from the entry
                 point file (default: None). If ``source_dir`` is an S3 URI, it must
                 point to a tar.gz file. Structure within this directory are preserved
-                when training on Amazon SageMaker.
+                when training on Amazon SageMaker. Must be a local path when using training_recipe.
             hyperparameters (dict[str, str] or dict[str, PipelineVariable]): Hyperparameters
                 that will be used for training (default: None). The hyperparameters are made
                 accessible as a dict[str, str] to the training code on
@@ -635,6 +635,10 @@ def _setup_for_training_recipe(cls, training_recipe, recipe_overrides, source_di
         if source_dir is None:
             args["source_dir"] = "."
         else:
+            if not os.path.exists(source_dir):
+                raise ValueError(
+                    "When using training_recipe, source_dir must be a local directory."
+                )
             args["source_dir"] = source_dir
 
         recipe_name = os.path.splitext(os.path.basename(training_recipe))[0]

src/sagemaker/pytorch/training_recipes.json

@@ -7,5 +7,9 @@
         "version":  "2.4.1",
         "additional_args": {}
     },
-    "neuron_image": "855988369404.dkr.ecr.us-west-2.amazonaws.com/chinmayee-dev:neuron_sept26_v1"
+    "neuron_image" : {
+        "framework": "hyperpod-recipes-neuron",
+        "version":  "2.1.2",
+        "additional_args": {}
+    }
 }

tests/unit/test_pytorch.py

@@ -23,6 +23,7 @@
 from sagemaker import image_uris
 from sagemaker.pytorch import defaults
 from sagemaker.pytorch import PyTorch, PyTorchPredictor, PyTorchModel
+from sagemaker.pytorch.estimator import _get_training_recipe_image_uri
 from sagemaker.instance_group import InstanceGroup
 from sagemaker.session_settings import SessionSettings
 
@@ -67,6 +68,13 @@
     "neuronx-distributed-training/refs/heads/main/examples/"
     "conf/hf_llama3_8B_config.yaml"
 )
+RECIPE_GPU_IMAGE = (
+    "658645717510.dkr.ecr.us-west-2.amazonaws.com/smdistributed-modelparallel:2.4.1-gpu-py311"
+)
+RECIPE_NEURON_IMAGE = (
+    "763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+    "pytorch-training-neuronx:2.1.2-neuronx-py310-sdk2.20.2-ubuntu20.04"
+)
 
 
 @pytest.fixture(name="sagemaker_session")
@@ -1085,3 +1093,16 @@ def test_training_recipe_for_trainium_custom_source_dir(sagemaker_session):
         },
     }
     assert pytorch.distribution == expected_distribution
+
+
+def test_training_recipe_images_uri():
+    gpu_image_cfg = {"framework": "pytorch-smp", "version": "2.4.1", "additional_args": {}}
+    gpu_image_uri = _get_training_recipe_image_uri(gpu_image_cfg, "us-west-2")
+    assert gpu_image_uri == RECIPE_GPU_IMAGE
+    neuron_image_cfg = {
+        "framework": "hyperpod-recipes-neuron",
+        "version": "2.1.2",
+        "additional_args": {},
+    }
+    neuron_image_uri = _get_training_recipe_image_uri(neuron_image_cfg, "us-west-2")
+    assert neuron_image_uri == RECIPE_NEURON_IMAGE