Merge branch 'master' into MLX-1269

CHANGELOG.md

@@ -1,5 +1,26 @@
 # Changelog
 
+## v2.214.2 (2024-04-01)
+
+### Bug Fixes and Other Changes
+
+ * Skip JS Tune integration test
+ * bump apache-airflow version to 2.8.3
+ * bump onnx version to >=1.15.0
+ * Updates for DJL 0.27.0 release
+ * Tune (local mode) support for Jumpstart Models
+ * attach jumpstart estimator for gated model
+
+## v2.214.1 (2024-03-27)
+
+### Bug Fixes and Other Changes
+
+ * Update schema dependency version
+ * remove failing deprecated tests from suite
+ * update readme, trigger p311 tests
+ * JumpStart list models flaky tests
+ * fix badge in README
+
 ## v2.214.0 (2024-03-22)
 
 ### Features

VERSION

@@ -1 +1 @@
-2.214.1.dev0
+2.214.3.dev0

requirements/extras/test_requirements.txt

@@ -12,7 +12,7 @@ awslogs==0.14.0
 black==22.3.0
 stopit==1.1.2
 # Update tox.ini to have correct version of airflow constraints file
-apache-airflow==2.8.2
+apache-airflow==2.8.3
 apache-airflow-providers-amazon==7.2.1
 attrs>=23.1.0,<24
 fabric==2.6.0
@@ -31,6 +31,8 @@ transformers>=4.36.0
 sentencepiece==0.1.99
 # https://github.com/triton-inference-server/server/issues/6246
 tritonclient[http]<2.37.0
-onnx==1.14.1
+onnx>=1.15.0
 # tf2onnx==1.15.1
 nbformat>=5.9,<6
+accelerate>=0.24.1,<=0.27.0
+schema==0.7.5

src/sagemaker/image_uri_config/djl-deepspeed.json

@@ -3,6 +3,38 @@
         "inference"
     ],
     "versions": {
+        "0.27.0": {
+            "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-southeast-1": "763104351884",
+                "ap-southeast-2": "763104351884",
+                "ap-southeast-3": "907027046896",
+                "ca-central-1": "763104351884",
+                "cn-north-1": "727897471807",
+                "cn-northwest-1": "727897471807",
+                "eu-central-1": "763104351884",
+                "eu-north-1": "763104351884",
+                "eu-west-1": "763104351884",
+                "eu-west-2": "763104351884",
+                "eu-west-3": "763104351884",
+                "eu-south-1": "692866216735",
+                "me-south-1": "217643126080",
+                "sa-east-1": "763104351884",
+                "us-east-1": "763104351884",
+                "us-east-2": "763104351884",
+                "us-west-1": "763104351884",
+                "us-west-2": "763104351884",
+                "ca-west-1": "204538143572"
+            },
+            "repository": "djl-inference",
+            "tag_prefix": "0.27.0-deepspeed0.12.6-cu121"
+        },
         "0.26.0": {
             "registries": {
                 "af-south-1": "626614931356",

src/sagemaker/image_uri_config/djl-neuronx.json

@@ -3,6 +3,24 @@
         "inference"
     ],
     "versions": {
+         "0.27.0": {
+            "registries": {
+                "ap-northeast-1": "763104351884",
+                "ap-south-1": "763104351884",
+                "ap-southeast-1": "763104351884",
+                "ap-southeast-2": "763104351884",
+                "eu-central-1": "763104351884",
+                "eu-west-1": "763104351884",
+                "eu-west-3": "763104351884",
+                "sa-east-1": "763104351884",
+                "us-east-1": "763104351884",
+                "us-east-2": "763104351884",
+                "us-west-2": "763104351884",
+                "ca-west-1": "204538143572"
+            },
+            "repository": "djl-inference",
+            "tag_prefix": "0.27.0-neuronx-sdk2.18.0"
+        },
          "0.26.0": {
             "registries": {
                 "ap-northeast-1": "763104351884",

src/sagemaker/image_uri_config/djl-tensorrtllm.json

@@ -3,6 +3,38 @@
         "inference"
     ],
     "versions": {
+        "0.27.0": {
+            "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-southeast-1": "763104351884",
+                "ap-southeast-2": "763104351884",
+                "ap-southeast-3": "907027046896",
+                "ca-central-1": "763104351884",
+                "cn-north-1": "727897471807",
+                "cn-northwest-1": "727897471807",
+                "eu-central-1": "763104351884",
+                "eu-north-1": "763104351884",
+                "eu-west-1": "763104351884",
+                "eu-west-2": "763104351884",
+                "eu-west-3": "763104351884",
+                "eu-south-1": "692866216735",
+                "me-south-1": "217643126080",
+                "sa-east-1": "763104351884",
+                "us-east-1": "763104351884",
+                "us-east-2": "763104351884",
+                "us-west-1": "763104351884",
+                "us-west-2": "763104351884",
+                "ca-west-1": "204538143572"
+            },
+            "repository": "djl-inference",
+            "tag_prefix": "0.27.0-tensorrtllm0.8.0-cu122"
+        },
         "0.26.0": {
             "registries": {
                 "af-south-1": "626614931356",

src/sagemaker/jumpstart/estimator.py

@@ -37,6 +37,7 @@
 from sagemaker.jumpstart.utils import (
     validate_model_id_and_get_type,
     resolve_model_sagemaker_config_field,
+    verify_model_region_and_return_specs,
 )
 from sagemaker.utils import stringify_object, format_tags, Tags
 from sagemaker.model_monitor.data_capture_config import DataCaptureConfig
@@ -729,11 +730,27 @@ def attach(
 
         model_version = model_version or "*"
 
+        additional_kwargs = {"model_id": model_id, "model_version": model_version}
+
+        model_specs = verify_model_region_and_return_specs(
+            model_id=model_id,
+            version=model_version,
+            region=sagemaker_session.boto_region_name,
+            scope=JumpStartScriptScope.TRAINING,
+            tolerate_deprecated_model=True,  # model is already trained, so tolerate if deprecated
+            tolerate_vulnerable_model=True,  # model is already trained, so tolerate if vulnerable
+            sagemaker_session=sagemaker_session,
+        )
+
+        # eula was already accepted if the model was successfully trained
+        if model_specs.is_gated_model():
+            additional_kwargs.update({"environment": {"accept_eula": "true"}})
+
         return cls._attach(
             training_job_name=training_job_name,
             sagemaker_session=sagemaker_session,
             model_channel_name=model_channel_name,
-            additional_kwargs={"model_id": model_id, "model_version": model_version},
+            additional_kwargs=additional_kwargs,
         )
 
     def deploy(