chore: always include training script if available

Author: evakravi

src/sagemaker/jumpstart/artifacts.py

@@ -239,7 +239,6 @@ def _retrieve_script_uri(
     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.
 
@@ -260,8 +259,6 @@ 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.
 
@@ -284,17 +281,11 @@ def _retrieve_script_uri(
     )
 
     if script_scope == JumpStartScriptScope.INFERENCE:
-        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}."
-                )
+        model_script_key = (
+            getattr(model_specs, "training_prepacked_script_key") or model_specs.hosting_script_key
+        )
+
     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/script_uris.py

@@ -29,7 +29,6 @@ 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.
 
@@ -48,8 +47,6 @@ 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.
 
@@ -70,5 +67,4 @@ def retrieve(
         region,
         tolerate_vulnerable_model,
         tolerate_deprecated_model,
-        include_training_script,
     )

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

@@ -15,9 +15,8 @@
 from mock.mock import patch
 
 from sagemaker import script_uris
-import pytest
 
-from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec, get_special_model_spec
+from tests.unit.sagemaker.jumpstart.utils import get_special_model_spec
 
 
 @patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")
@@ -32,31 +31,8 @@ def test_jumpstart_combined_artifacts(patched_get_model_specs):
         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,
-        )
-
-    patched_get_model_specs.side_effect = get_prototype_model_spec
-
-    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,
-        )