Merge branch 'master' into lmi-bumps

tosterberg · web-flow · commit 42f80887932f · 2024-10-18T11:26:22.000-07:00
diff --git a/requirements/extras/feature-processor_requirements.txt b/requirements/extras/feature-processor_requirements.txt
@@ -1,2 +1,2 @@
-pyspark==3.3.1
+pyspark==3.3.2
 sagemaker-feature-store-pyspark-3.3
diff --git a/requirements/extras/test_requirements.txt b/requirements/extras/test_requirements.txt
@@ -24,7 +24,7 @@ Jinja2==3.1.4
 pyvis==0.2.1
 pandas>=1.3.5,<1.5
 scikit-learn==1.3.0
-cloudpickle==2.2.1
+cloudpickle>=2.2.1
 PyYAML==6.0
 # TODO find workaround
 xgboost>=1.6.2,<=1.7.6
diff --git a/src/sagemaker/fw_utils.py b/src/sagemaker/fw_utils.py
@@ -153,6 +153,7 @@
     "2.1.2",
     "2.2.0",
     "2.3.1",
+    "2.4.1",
 ]
 
 TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]
diff --git a/src/sagemaker/image_uri_config/huggingface-llm-neuronx.json b/src/sagemaker/image_uri_config/huggingface-llm-neuronx.json
@@ -4,7 +4,7 @@
             "inf2"
         ],
         "version_aliases": {
-            "0.0": "0.0.24"
+            "0.0": "0.0.25"
         },
         "versions": {
             "0.0.16": {
@@ -301,6 +301,39 @@
                 "container_version": {
                     "inf2": "ubuntu22.04"
                 }
+            },
+            "0.0.25": {
+                "py_versions": [
+                    "py310"
+                ],
+                "registries": {
+                    "ap-northeast-1": "763104351884",
+                    "ap-south-1": "763104351884",
+                    "ap-south-2": "772153158452",
+                    "ap-southeast-1": "763104351884",
+                    "ap-southeast-2": "763104351884",
+                    "ap-southeast-4": "457447274322",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-south-2": "503227376785",
+                    "eu-west-1": "763104351884",
+                    "eu-west-3": "763104351884",
+                    "il-central-1": "780543022126",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
+                    "us-west-2": "763104351884",
+                    "ca-west-1": "204538143572"
+                },
+                "tag_prefix": "2.1.2-optimum0.0.25",
+                "repository": "huggingface-pytorch-tgi-inference",
+                "container_version": {
+                    "inf2": "ubuntu22.04"
+                }
             }
         }
     }
diff --git a/src/sagemaker/image_uri_config/huggingface-llm.json b/src/sagemaker/image_uri_config/huggingface-llm.json
@@ -12,7 +12,7 @@
             "1.2": "1.2.0",
             "1.3": "1.3.3",
             "1.4": "1.4.5",
-            "2.0": "2.2.0"
+            "2.0": "2.3.1"
         },
         "versions": {
             "0.6.0": {
@@ -719,6 +719,53 @@
                 "container_version": {
                     "gpu": "cu121-ubuntu22.04-v2.0"
                 }
+            },
+            "2.3.1": {
+                "py_versions": [
+                    "py311"
+                ],
+                "registries": {
+                    "af-south-1": "626614931356",
+                    "il-central-1": "780543022126",
+                    "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",
+                    "ca-central-1": "763104351884",
+                    "cn-north-1": "727897471807",
+                    "cn-northwest-1": "727897471807",
+                    "eu-central-1": "763104351884",
+                    "eu-central-2": "380420809688",
+                    "eu-north-1": "763104351884",
+                    "eu-west-1": "763104351884",
+                    "eu-west-2": "763104351884",
+                    "eu-west-3": "763104351884",      
+                    "eu-south-1": "692866216735",
+                    "eu-south-2": "503227376785",
+                    "me-south-1": "217643126080",
+                    "me-central-1": "914824155844",
+                    "sa-east-1": "763104351884",
+                    "us-east-1": "763104351884",
+                    "us-east-2": "763104351884",
+                    "us-gov-east-1": "446045086412",
+                    "us-gov-west-1": "442386744353",
+                    "us-iso-east-1": "886529160074",
+                    "us-isob-east-1": "094389454867",
+                    "us-west-1": "763104351884",
+                    "us-west-2": "763104351884",
+                    "ca-west-1": "204538143572"
+                },
+                "tag_prefix": "2.4.0-tgi2.3.1",
+                "repository": "huggingface-pytorch-tgi-inference",
+                "container_version": {
+                    "gpu": "cu124-ubuntu22.04"
+                }
             }
         }
     }
diff --git a/src/sagemaker/image_uri_config/pytorch-smp.json b/src/sagemaker/image_uri_config/pytorch-smp.json
@@ -8,7 +8,8 @@
             "2.1": "2.1.2",
             "2.2": "2.3.1",
             "2.2.0": "2.3.1",
-            "2.3.1": "2.5.0"
+            "2.3.1": "2.5.0",
+            "2.4.1": "2.6.0"
         },
         "versions": {
             "2.0.1": {
@@ -160,6 +161,31 @@
                     "us-west-2": "658645717510"
                 },
                 "repository": "smdistributed-modelparallel"
+            },
+            "2.6.0": {
+                "py_versions": [
+                    "py311"
+                ],
+                "registries": {
+                    "ap-northeast-1": "658645717510",
+                    "ap-northeast-2": "658645717510",
+                    "ap-northeast-3": "658645717510",
+                    "ap-south-1": "658645717510",
+                    "ap-southeast-1": "658645717510",
+                    "ap-southeast-2": "658645717510",
+                    "ca-central-1": "658645717510",
+                    "eu-central-1": "658645717510",
+                    "eu-north-1": "658645717510",
+                    "eu-west-1": "658645717510",
+                    "eu-west-2": "658645717510",
+                    "eu-west-3": "658645717510",
+                    "sa-east-1": "658645717510",
+                    "us-east-1": "658645717510",
+                    "us-east-2": "658645717510",
+                    "us-west-1": "658645717510",
+                    "us-west-2": "658645717510"
+                },
+                "repository": "smdistributed-modelparallel"
             }
         }
     }
diff --git a/src/sagemaker/image_uris.py b/src/sagemaker/image_uris.py
@@ -699,6 +699,7 @@ def get_training_image_uri(
                         or "2.1" in framework_version
                         or "2.2" in framework_version
                         or "2.3" in framework_version
+                        or "2.4" in framework_version
                     ):
                         container_version = "cu121"
                     else:
diff --git a/tests/unit/sagemaker/image_uris/test_huggingface_llm.py b/tests/unit/sagemaker/image_uris/test_huggingface_llm.py
@@ -45,6 +45,7 @@
         "2.0.1": "2.1.1-tgi2.0.1-gpu-py310-cu121-ubuntu22.04",
         "2.0.2": "2.3.0-tgi2.0.2-gpu-py310-cu121-ubuntu22.04",
         "2.2.0": "2.3.0-tgi2.2.0-gpu-py310-cu121-ubuntu22.04-v2.0",
+        "2.3.1": "2.4.0-tgi2.3.1-gpu-py311-cu124-ubuntu22.04",
     },
     "inf2": {
         "0.0.16": "1.13.1-optimum0.0.16-neuronx-py310-ubuntu22.04",
@@ -56,6 +57,7 @@
         "0.0.22": "2.1.2-optimum0.0.22-neuronx-py310-ubuntu22.04",
         "0.0.23": "2.1.2-optimum0.0.23-neuronx-py310-ubuntu22.04",
         "0.0.24": "2.1.2-optimum0.0.24-neuronx-py310-ubuntu22.04",
+        "0.0.25": "2.1.2-optimum0.0.25-neuronx-py310-ubuntu22.04",
     },
 }
 
diff --git a/tests/unit/sagemaker/image_uris/test_smp_v2.py b/tests/unit/sagemaker/image_uris/test_smp_v2.py
@@ -36,14 +36,17 @@ def test_smp_v2(load_config):
                 for region in ACCOUNTS.keys():
                     for instance_type in CONTAINER_VERSIONS.keys():
                         cuda_vers = CONTAINER_VERSIONS[instance_type]
-                        if "2.1" in version or "2.2" in version or "2.3" in version:
+                        if (
+                            "2.1" in version
+                            or "2.2" in version
+                            or "2.3" in version
+                            or "2.4" in version
+                        ):
                             cuda_vers = "cu121"
 
-                        if "2.3.1" == version:
+                        if "2.3.1" == version or "2.4.1" == version:
                             py_version = "py311"
 
-                        print(version, py_version)
-
                         uri = image_uris.get_training_image_uri(
                             region,
                             framework="pytorch",

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`		`-pyspark==3.3.1`
	`1`	`+pyspark==3.3.2`
`2`	`2`	`sagemaker-feature-store-pyspark-3.3`
Original file line number	Diff line number	Diff line change
`@@ -153,6 +153,7 @@`
`153`	`153`	`"2.1.2",`
`154`	`154`	`"2.2.0",`
`155`	`155`	`"2.3.1",`
	`156`	`+ "2.4.1",`
`156`	`157`	`]`
`157`	`158`
`158`	`159`	`TRAINIUM_SUPPORTED_DISTRIBUTION_STRATEGIES = ["torch_distributed"]`