chore: use ValueError

Author: evakravi

src/sagemaker/jumpstart/artifacts.py

@@ -208,8 +208,6 @@ def _retrieve_model_uri(
         VulnerableJumpStartModelError: If any of the dependencies required by the script have
             known security vulnerabilities.
         DeprecatedJumpStartModelError: If the version of the model is deprecated.
-        NotImplementedError: If the combination of arguments doesn't support combined model
-            and script artifact.
     """
     if region is None:
         region = JUMPSTART_DEFAULT_REGION_NAME
@@ -234,13 +232,13 @@ def _retrieve_model_uri(
         else:
             model_artifact_key = getattr(model_specs, "hosting_prepacked_artifact_key", None)
             if model_artifact_key is None:
-                raise NotImplementedError(error_msg_no_combined_artifact)
+                raise ValueError(error_msg_no_combined_artifact)
 
     elif model_scope == JumpStartScriptScope.TRAINING:
         if not include_script:
             model_artifact_key = model_specs.training_artifact_key
         else:
-            raise NotImplementedError(error_msg_no_combined_artifact)
+            raise ValueError(error_msg_no_combined_artifact)
 
     bucket = os.environ.get(
         ENV_VARIABLE_JUMPSTART_MODEL_ARTIFACT_BUCKET_OVERRIDE

src/sagemaker/model_uris.py

@@ -60,8 +60,6 @@ def retrieve(
         VulnerableJumpStartModelError: If any of the dependencies required by the script have
             known security vulnerabilities.
         DeprecatedJumpStartModelError: If the version of the model is deprecated.
-        NotImplementedError: If the combination of arguments doesn't support combined model
-            and script artifact.
     """
     if not jumpstart_utils.is_jumpstart_model_input(model_id, model_version):
         raise ValueError("Must specify `model_id` and `model_version` when retrieving model URIs.")

tests/unit/sagemaker/model_uris/jumpstart/test_combined_artifact.py

@@ -39,18 +39,18 @@ def test_jumpstart_combined_artifacts(patched_get_model_specs):
         "prepack/v1.0.0/infer-prepack-huggingface-text2text-flan-t5-xxl-fp16.tar.gz"
     )
 
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(ValueError):
         model_uris.retrieve(
             region="us-west-2",
-            model_scope="transfer_learning",
+            model_scope="training",
             model_id=model_id_combined_model_artifact,
             model_version="*",
             include_script=True,
         )
 
     model_id_combined_model_artifact_unsupported = "xgboost-classification-model"
 
-    with pytest.raises(NotImplementedError):
+    with pytest.raises(ValueError):
         model_uris.retrieve(
             region="us-west-2",
             model_scope="inference",