feat: combined inference + training script artifact

Author: evakravi

src/sagemaker/jumpstart/artifacts.py

@@ -235,10 +235,11 @@ def _retrieve_model_uri(
 def _retrieve_script_uri(
     model_id: str,
     model_version: str,
-    script_scope: Optional[str],
-    region: Optional[str],
-    tolerate_vulnerable_model: bool,
-    tolerate_deprecated_model: bool,
+    script_scope: Optional[str] = None,
+    region: Optional[str] = False,
+    tolerate_vulnerable_model: bool = False,
+    tolerate_deprecated_model: bool = False,
+    include_training_script: bool = False,
 ):
     """Retrieves the script S3 URI associated with the model matching the given arguments.
 
@@ -259,6 +260,8 @@ def _retrieve_script_uri(
         tolerate_deprecated_model (bool): True if deprecated versions of model
             specifications should be tolerated (exception not raised). If False, raises
             an exception if the version of the model is deprecated.
+        include_training_script (bool): True if training script should be packaged along with
+            inference script. (Default: False.)
     Returns:
         str: the model script URI for the corresponding model.
 
@@ -281,8 +284,17 @@ def _retrieve_script_uri(
     )
 
     if script_scope == JumpStartScriptScope.INFERENCE:
-        model_script_key = model_specs.hosting_script_key
+        if not include_training_script:
+            model_script_key = model_specs.hosting_script_key
+        else:
+            model_script_key = getattr(model_specs, "training_prepacked_script_key", None)
+            if model_script_key is None:
+                raise ValueError(
+                    f"Cannot include training script for {model_id} with version {model_version}."
+                )
     elif script_scope == JumpStartScriptScope.TRAINING:
+        if include_training_script:
+            raise ValueError("Can only include training script for inference jobs.")
         model_script_key = model_specs.training_script_key
 
     bucket = os.environ.get(

src/sagemaker/jumpstart/types.py

@@ -293,6 +293,7 @@ class JumpStartModelSpecs(JumpStartDataHolderType):
         "training_vulnerabilities",
         "deprecated",
         "metrics",
+        "training_prepacked_script_key",
     ]
 
     def __init__(self, spec: Dict[str, Any]):
@@ -330,6 +331,9 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.training_vulnerabilities: List[str] = json_obj["training_vulnerabilities"]
         self.deprecated: bool = bool(json_obj["deprecated"])
         self.metrics: Optional[List[Dict[str, str]]] = json_obj.get("metrics", None)
+        self.training_prepacked_script_key: Optional[str] = json_obj.get(
+            "training_prepacked_script_key", None
+        )
 
         if self.training_supported:
             self.training_ecr_specs: JumpStartECRSpecs = JumpStartECRSpecs(

src/sagemaker/script_uris.py

@@ -29,6 +29,7 @@ def retrieve(
     script_scope=None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
+    include_training_script: bool = False,
 ) -> str:
     """Retrieves the script S3 URI associated with the model matching the given arguments.
 
@@ -47,6 +48,8 @@ def retrieve(
         tolerate_deprecated_model (bool): ``True`` if deprecated models should be tolerated
             without raising an exception. ``False`` if these models should raise an exception.
             (Default: False).
+        include_training_script (bool): True if training script should be packaged along with
+            inference script. (Default: False.)
     Returns:
         str: The model script URI for the corresponding model.
 
@@ -67,4 +70,5 @@ def retrieve(
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
+        include_training_script,
     )

tests/unit/sagemaker/jumpstart/constants.py

@@ -1070,6 +1070,127 @@
             },
         ],
     },
+    "mock-model-training-prepacked-script-key": {
+        "model_id": "sklearn-classification-linear",
+        "url": "https://scikit-learn.org/stable/",
+        "version": "1.0.0",
+        "min_sdk_version": "2.68.1",
+        "training_supported": True,
+        "incremental_training_supported": False,
+        "hosting_ecr_specs": {
+            "framework": "sklearn",
+            "framework_version": "0.23-1",
+            "py_version": "py3",
+        },
+        "hosting_artifact_key": "sklearn-infer/infer-sklearn-classification-linear.tar.gz",
+        "hosting_script_key": "source-directory-tarballs/sklearn/inference/classification/v1.0.0/sourcedir.tar.gz",
+        "inference_vulnerable": False,
+        "inference_dependencies": [],
+        "inference_vulnerabilities": [],
+        "training_vulnerable": False,
+        "training_dependencies": [],
+        "training_vulnerabilities": [],
+        "deprecated": False,
+        "hyperparameters": [
+            {
+                "name": "tol",
+                "type": "float",
+                "default": 0.0001,
+                "min": 1e-20,
+                "max": 50,
+                "scope": "algorithm",
+            },
+            {
+                "name": "penalty",
+                "type": "text",
+                "default": "l2",
+                "options": ["l1", "l2", "elasticnet", "none"],
+                "scope": "algorithm",
+            },
+            {
+                "name": "alpha",
+                "type": "float",
+                "default": 0.0001,
+                "min": 1e-20,
+                "max": 999,
+                "scope": "algorithm",
+            },
+            {
+                "name": "l1_ratio",
+                "type": "float",
+                "default": 0.15,
+                "min": 0,
+                "max": 1,
+                "scope": "algorithm",
+            },
+            {
+                "name": "sagemaker_submit_directory",
+                "type": "text",
+                "default": "/opt/ml/input/data/code/sourcedir.tar.gz",
+                "scope": "container",
+            },
+            {
+                "name": "sagemaker_program",
+                "type": "text",
+                "default": "transfer_learning.py",
+                "scope": "container",
+            },
+            {
+                "name": "sagemaker_container_log_level",
+                "type": "text",
+                "default": "20",
+                "scope": "container",
+            },
+        ],
+        "training_script_key": "source-directory-tarballs/sklearn/transfer_learning/classification/"
+        "v1.0.0/sourcedir.tar.gz",
+        "training_prepacked_script_key": "some/key/to/training_prepacked_script_key.tar.gz",
+        "training_ecr_specs": {
+            "framework_version": "0.23-1",
+            "framework": "sklearn",
+            "py_version": "py3",
+        },
+        "training_artifact_key": "sklearn-training/train-sklearn-classification-linear.tar.gz",
+        "inference_environment_variables": [
+            {
+                "name": "SAGEMAKER_PROGRAM",
+                "type": "text",
+                "default": "inference.py",
+                "scope": "container",
+            },
+            {
+                "name": "SAGEMAKER_SUBMIT_DIRECTORY",
+                "type": "text",
+                "default": "/opt/ml/model/code",
+                "scope": "container",
+            },
+            {
+                "name": "SAGEMAKER_CONTAINER_LOG_LEVEL",
+                "type": "text",
+                "default": "20",
+                "scope": "container",
+            },
+            {
+                "name": "MODEL_CACHE_ROOT",
+                "type": "text",
+                "default": "/opt/ml/model",
+                "scope": "container",
+            },
+            {"name": "SAGEMAKER_ENV", "type": "text", "default": "1", "scope": "container"},
+            {
+                "name": "SAGEMAKER_MODEL_SERVER_WORKERS",
+                "type": "text",
+                "default": "1",
+                "scope": "container",
+            },
+            {
+                "name": "SAGEMAKER_MODEL_SERVER_TIMEOUT",
+                "type": "text",
+                "default": "3600",
+                "scope": "container",
+            },
+        ],
+    },
 }
 
 BASE_SPEC = {
@@ -1093,6 +1214,7 @@
     "training_artifact_key": "pytorch-training/train-pytorch-ic-mobilenet-v2.tar.gz",
     "hosting_script_key": "source-directory-tarballs/pytorch/inference/ic/v1.0.0/sourcedir.tar.gz",
     "training_script_key": "source-directory-tarballs/pytorch/transfer_learning/ic/v1.0.0/sourcedir.tar.gz",
+    "training_prepacked_script_key": None,
     "hyperparameters": [
         {
             "name": "epochs",

tests/unit/sagemaker/script_uris/jumpstart/test_combined_script_artifact.py

@@ -0,0 +1,60 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+from __future__ import absolute_import
+
+from mock.mock import patch
+
+from sagemaker import script_uris
+import pytest
+
+from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
+
+
+@patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
+def test_jumpstart_combined_artifacts(patched_get_model_specs):
+
+    patched_get_model_specs.side_effect = get_prototype_model_spec
+
+    model_id_combined_script_artifact = "mock-model-training-prepacked-script-key"
+
+    uri = script_uris.retrieve(
+        region="us-west-2",
+        script_scope="inference",
+        model_id=model_id_combined_script_artifact,
+        model_version="*",
+        include_training_script=True,
+    )
+    assert (
+        uri == "s3://jumpstart-cache-prod-us-west-2/some/key/to/"
+        "training_prepacked_script_key.tar.gz"
+    )
+
+    with pytest.raises(ValueError):
+        script_uris.retrieve(
+            region="us-west-2",
+            script_scope="training",
+            model_id=model_id_combined_script_artifact,
+            model_version="*",
+            include_training_script=True,
+        )
+
+    model_id_combined_script_artifact_unsupported = "xgboost-classification-model"
+
+    with pytest.raises(ValueError):
+        script_uris.retrieve(
+            region="us-west-2",
+            script_scope="inference",
+            model_id=model_id_combined_script_artifact_unsupported,
+            model_version="*",
+            include_training_script=True,
+        )