fix: formatting issues

Author: evakravi

src/sagemaker/image_uris.py

@@ -155,22 +155,14 @@ def retrieve(
 
     if is_jumpstart_model_input(model_id, model_version):
         return artifacts._retrieve_image_uri(
-            model_id,
-            model_version,
-            image_scope,
-            hub_arn,
-            framework,
-            region,
-            version,
-            py_version,
-            instance_type,
-            accelerator_type,
-            container_version,
-            distribution,
-            base_framework_version,
-            training_compiler_config,
-            tolerate_vulnerable_model,
-            tolerate_deprecated_model,
+            model_id=model_id,
+            model_version=model_version,
+            image_scope=image_scope,
+            hub_arn=hub_arn,
+            region=region,
+            instance_type=instance_type,
+            tolerate_vulnerable_model=tolerate_vulnerable_model,
+            tolerate_deprecated_model=tolerate_deprecated_model,
             sagemaker_session=sagemaker_session,
             config_name=config_name,
             model_type=model_type,

src/sagemaker/jumpstart/artifacts/image_uris.py

@@ -33,16 +33,8 @@ def _retrieve_image_uri(
     model_version: str,
     image_scope: str,
     hub_arn: Optional[str] = None,
-    framework: Optional[str] = None,
     region: Optional[str] = None,
-    version: Optional[str] = None,
-    py_version: Optional[str] = None,
     instance_type: Optional[str] = None,
-    accelerator_type: Optional[str] = None,
-    container_version: Optional[str] = None,
-    distribution: Optional[str] = None,
-    base_framework_version: Optional[str] = None,
-    training_compiler_config: Optional[str] = None,
     tolerate_vulnerable_model: bool = False,
     tolerate_deprecated_model: bool = False,
     sagemaker_session: Session = DEFAULT_JUMPSTART_SAGEMAKER_SESSION,
@@ -64,30 +56,11 @@ def _retrieve_image_uri(
         image_scope (str): The image type, i.e. what it is used for.
             Valid values: "training", "inference", "eia". If ``accelerator_type`` is set,
             ``image_scope`` is ignored.
-        framework (str): The name of the framework or algorithm.
         region (str): The AWS region. (Default: None).
-        version (str): The framework or algorithm version. This is required if there is
-            more than one supported version for the given framework or algorithm.
-            (Default: None).
-        py_version (str): The Python version. This is required if there is
-            more than one supported Python version for the given framework version.
         instance_type (str): The SageMaker instance type. For supported types, see
             https://aws.amazon.com/sagemaker/pricing/instance-types. This is required if
             there are different images for different processor types.
             (Default: None).
-        accelerator_type (str): Elastic Inference accelerator type. For more, see
-            https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.
-            (Default: None).
-        container_version (str): the version of docker image.
-            Ideally the value of parameter should be created inside the framework.
-            For custom use, see the list of supported container versions:
-            https://github.com/aws/deep-learning-containers/blob/master/available_images.md.
-            (Default: None).
-        distribution (dict): A dictionary with information on how to run distributed training.
-            (Default: None).
-        training_compiler_config (:class:`~sagemaker.training_compiler.TrainingCompilerConfig`):
-            A configuration class for the SageMaker Training Compiler.
-            (Default: None).
         tolerate_vulnerable_model (bool): True if vulnerable versions of model
             specifications should be tolerated (exception not raised). If False, raises an
             exception if the script used by this version of the model has dependencies with known
@@ -145,7 +118,7 @@ def _retrieve_image_uri(
             f"with {instance_type} instance type in {region}. "
             "Please try another instance type or region."
         )
-    elif image_scope == JumpStartScriptScope.TRAINING:
+    if image_scope == JumpStartScriptScope.TRAINING:
         training_instance_type_variants = model_specs.training_instance_type_variants
         if training_instance_type_variants:
             image_uri = training_instance_type_variants.get_image_uri(
@@ -162,3 +135,5 @@ def _retrieve_image_uri(
             f"with {instance_type} instance type in {region}. "
             "Please try another instance type or region."
         )
+
+    raise ValueError(f"Invalid scope: {image_scope}")

tests/unit/sagemaker/hyperparameters/jumpstart/test_validate.py

@@ -21,7 +21,7 @@
 from sagemaker.jumpstart.exceptions import JumpStartHyperparametersError
 from sagemaker.jumpstart.types import JumpStartHyperparameter
 
-from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec, get_spec_from_base_spec
+from tests.unit.sagemaker.jumpstart.utils import get_prototype_model_spec
 
 region = "us-west-2"
 mock_client = boto3.client("s3")

tests/unit/sagemaker/image_uris/jumpstart/test_sklearn.py

@@ -53,14 +53,15 @@ def test_jumpstart_sklearn_image_uri(patched_get_model_specs, session):
         sagemaker_session=session,
     ).serving_image_uri(region, instance_type)
 
-    ## framework classes dont use digest.
+    # framework classes dont use digest.
     assert (
         framework_class_uri
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:1.2-1-cpu-py3"
+        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:1.2-1"
+        "-cpu-py3"
     )
     assert (
-        uri
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn@sha256:e09bbb7686077a1db23d316b699020a786a6e1636b2b89384be9651368c40f95"
+        uri == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn@"
+        "sha256:e09bbb7686077a1db23d316b699020a786a6e1636b2b89384be9651368c40f95"
     )
 
     # training
@@ -84,14 +85,14 @@ def test_jumpstart_sklearn_image_uri(patched_get_model_specs, session):
         sagemaker_session=session,
     ).training_image_uri(region=region)
 
-    ## framework classes dont use digest.
+    # framework classes dont use digest.
     assert (
         framework_class_uri
         == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn:1.2-1-cpu-py3"
     )
     assert (
-        uri
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn@sha256:e09bbb7686077a1db23d316b699020a786a6e1636b2b89384be9651368c40f95"
+        uri == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-scikit-learn"
+        "@sha256:e09bbb7686077a1db23d316b699020a786a6e1636b2b89384be9651368c40f95"
     )
 
     with pytest.raises(ValueError):

tests/unit/sagemaker/image_uris/jumpstart/test_xgboost.py

@@ -58,8 +58,8 @@ def test_jumpstart_xgboost_image_uri(patched_get_model_specs, session):
         == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.7-1"
     )
     assert (
-        uri
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost@sha256:ba417ec6d8d3e0c6b5f463bc9202e3b498b42260a29b61875f34beb6d99d8444"
+        uri == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost@sha256:"
+        "ba417ec6d8d3e0c6b5f463bc9202e3b498b42260a29b61875f34beb6d99d8444"
     )
 
     # training
@@ -89,6 +89,6 @@ def test_jumpstart_xgboost_image_uri(patched_get_model_specs, session):
         == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost:1.7-1"
     )
     assert (
-        uri
-        == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost@sha256:ba417ec6d8d3e0c6b5f463bc9202e3b498b42260a29b61875f34beb6d99d8444"
+        uri == "246618743249.dkr.ecr.us-west-2.amazonaws.com/sagemaker-xgboost"
+        "@sha256:ba417ec6d8d3e0c6b5f463bc9202e3b498b42260a29b61875f34beb6d99d8444"
     )

tests/unit/sagemaker/jumpstart/constants.py

@@ -12,6 +12,7 @@
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
 
+# flake8: noqa: E501
 
 SPECIAL_MODEL_SPECS_DICT = {
     "js-model-class-model-prepacked": {

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -243,10 +243,13 @@ def test_prepacked(
         mock_estimator_init.assert_called_once_with(
             instance_type="ml.p3.16xlarge",
             instance_count=1,
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04",
-            model_uri="s3://jumpstart-cache-prod-us-west-2/huggingface-training/train-huggingface"
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface"
+            "-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04",
+            model_uri="s3://jumpstart-cache-prod-us-west-2/huggingface-training"
+            "/train-huggingface"
             "-text2text-flan-t5-base.tar.gz",
-            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/huggingface/transfer_learning/text2text/prepack/v2.0.0/sourcedir.tar.gz",
+            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs"
+            "/huggingface/transfer_learning/text2text/prepack/v2.0.0/sourcedir.tar.gz",
             entry_point="transfer_learning.py",
             hyperparameters={
                 "epochs": "1",
@@ -323,7 +326,8 @@ def test_prepacked(
         mock_estimator_deploy.assert_called_once_with(
             instance_type="ml.g5.2xlarge",
             initial_instance_count=1,
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-tgi-inference:2.1.1-tgi1.4.0-gpu-py310-cu121-ubuntu20.04",
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+            "huggingface-pytorch-tgi-inference:2.1.1-tgi1.4.0-gpu-py310-cu121-ubuntu20.04",
             env={
                 "SAGEMAKER_PROGRAM": "inference.py",
                 "ENDPOINT_SERVER_TIMEOUT": "3600",
@@ -412,8 +416,10 @@ def test_gated_model_s3_uri(
         mock_estimator_init.assert_called_once_with(
             instance_type="ml.g5.12xlarge",
             instance_count=1,
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04",
-            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/meta/transfer_learning/textgeneration/v1.0.6/sourcedir.tar.gz",
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-"
+            "pytorch-training:2.0.0-transformers4.28.1-gpu-py310-cu118-ubuntu20.04",
+            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
+            "meta/transfer_learning/textgeneration/v1.0.6/sourcedir.tar.gz",
             entry_point="transfer_learning.py",
             hyperparameters={
                 "int8_quantization": "False",
@@ -457,7 +463,8 @@ def test_gated_model_s3_uri(
             encrypt_inter_container_traffic=True,
             environment={
                 "accept_eula": "true",
-                "SageMakerGatedModelS3Uri": "s3://sagemaker-repository-pdx/model-data-model-package_llama2-7b-f-v4-71eeccf76ddf33f2a18d2e16b9c7f302",
+                "SageMakerGatedModelS3Uri": "s3://sagemaker-repository-pdx/"
+                "model-data-model-package_llama2-7b-f-v4-71eeccf76ddf33f2a18d2e16b9c7f302",
             },
             tags=[
                 {
@@ -1919,7 +1926,8 @@ def test_model_artifact_variant_estimator(
             instance_count=1,
             image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.10.0-cpu-py38",
             model_uri="s3://jumpstart-cache-prod-us-west-2/hello-world-1",
-            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/pytorch/transfer_learning/ic/prepack/v1.1.0/sourcedir.tar.gz",
+            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
+            "pytorch/transfer_learning/ic/prepack/v1.1.0/sourcedir.tar.gz",
             entry_point="transfer_learning.py",
             hyperparameters={
                 "train_only_top_layer": "True",
@@ -1950,8 +1958,10 @@ def test_model_artifact_variant_estimator(
             instance_type="ml.p3.2xlarge",
             instance_count=1,
             image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-training:1.10.0-gpu-py38",
-            model_uri="s3://jumpstart-cache-prod-us-west-2/pytorch-training/v2.0.0/train-pytorch-ic-mobilenet-v2.tar.gz",
-            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/pytorch/transfer_learning/ic/prepack/v1.1.0/sourcedir.tar.gz",
+            model_uri="s3://jumpstart-cache-prod-us-west-2/pytorch-training/"
+            "v2.0.0/train-pytorch-ic-mobilenet-v2.tar.gz",
+            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
+            "pytorch/transfer_learning/ic/prepack/v1.1.0/sourcedir.tar.gz",
             entry_point="transfer_learning.py",
             hyperparameters={
                 "train_only_top_layer": "True",
@@ -2061,9 +2071,12 @@ def test_estimator_initialization_with_config_name(
         mock_estimator_init.assert_called_with(
             instance_type="ml.p2.xlarge",
             instance_count=1,
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-training:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04",
-            model_uri="s3://jumpstart-cache-prod-us-west-2/artifacts/meta-textgeneration-llama-2-7b/gpu-training/model/",
-            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/pytorch/transfer_learning/ic/prepack/v1.1.0/sourcedir.tar.gz",
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+            "huggingface-pytorch-training:1.13.1-neuronx-py310-sdk2.14.1-ubuntu20.04",
+            model_uri="s3://jumpstart-cache-prod-us-west-2/artifacts/"
+            "meta-textgeneration-llama-2-7b/gpu-training/model/",
+            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
+            "pytorch/transfer_learning/ic/prepack/v1.1.0/sourcedir.tar.gz",
             entry_point="transfer_learning.py",
             hyperparameters={
                 "train_only_top_layer": "True",
@@ -2131,7 +2144,8 @@ def test_estimator_set_config_name(
             "pytorch-training:1.13.1-py310-sdk2.14.1-ubuntu20.04",
             model_uri="s3://jumpstart-cache-prod-us-west-2/artifacts/meta-textgeneration-llama-2-7b/"
             "gpu-training-budget/model/",
-            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/pytorch/transfer_learning/ic/prepack/v1.1.0/sourcedir.tar.gz",
+            source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
+            "pytorch/transfer_learning/ic/prepack/v1.1.0/sourcedir.tar.gz",
             entry_point="transfer_learning.py",
             hyperparameters={
                 "train_only_top_layer": "True",

tests/unit/sagemaker/jumpstart/model/test_model.py

@@ -282,7 +282,9 @@ def test_non_prepacked_inference_component_based_endpoint_no_default_pass_custom
             "1.10.2-transformers4.17.0-gpu-py38-cu113-ubuntu20.04",
             model_data={
                 "S3DataSource": {
-                    "S3Uri": "s3://jumpstart-cache-prod-us-west-2/huggingface-txt2img/huggingface-txt2img-conflictx-complex-lineart/artifacts/inference-prepack/v1.0.0/",
+                    "S3Uri": "s3://jumpstart-cache-prod-us-west-2/"
+                    "huggingface-txt2img/huggingface-txt2img-conflictx"
+                    "-complex-lineart/artifacts/inference-prepack/v1.0.0/",
                     "S3DataType": "S3Prefix",
                     "CompressionType": "None",
                 }
@@ -370,7 +372,9 @@ def test_prepacked(
             "1.10.2-transformers4.17.0-gpu-py38-cu113-ubuntu20.04",
             model_data={
                 "S3DataSource": {
-                    "S3Uri": "s3://jumpstart-cache-prod-us-west-2/huggingface-txt2img/huggingface-txt2img-conflictx-complex-lineart/artifacts/inference-prepack/v1.0.0/",
+                    "S3Uri": "s3://jumpstart-cache-prod-us-west-2/"
+                    "huggingface-txt2img/huggingface-txt2img-conflictx"
+                    "-complex-lineart/artifacts/inference-prepack/v1.0.0/",
                     "S3DataType": "S3Prefix",
                     "CompressionType": "None",
                 }
@@ -1188,7 +1192,8 @@ def test_jumpstart_model_package_arn(
         self.assertEqual(
             mock_session.create_model.call_args[0][2],
             {
-                "ModelPackageName": "arn:aws:sagemaker:us-west-2:594846645681:model-package/llama2-7b-f-v4-71eeccf76ddf33f2a18d2e16b9c7f302"
+                "ModelPackageName": "arn:aws:sagemaker:us-west-2:594846645681:"
+                "model-package/llama2-7b-f-v4-71eeccf76ddf33f2a18d2e16b9c7f302"
             },
         )
 
@@ -1269,7 +1274,8 @@ def test_jumpstart_model_package_arn_unsupported_region(
             JumpStartModel(model_id=model_id, region="us-west-1")
         assert (
             str(e.value) == "Model package arn for 'js-model-package-arn' not supported in "
-            "us-west-1. Please try one of the following regions: us-west-2, us-east-2, us-east-1, eu-west-1, ap-southeast-1, ap-southeast-2."
+            "us-west-1. Please try one of the following regions: "
+            "us-west-2, us-east-2, us-east-1, eu-west-1, ap-southeast-1, ap-southeast-2."
         )
 
     @mock.patch(
@@ -1383,8 +1389,10 @@ def test_model_data_s3_prefix_model(
         JumpStartModel(model_id=model_id, instance_type="ml.p2.xlarge")
 
         mock_model_init.assert_called_once_with(
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-pytorch-inference:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04",
-            model_data="s3://jumpstart-cache-prod-us-west-2/huggingface-infer/prepack/v1.1.2/infer-prepack-huggingface-text2text-flan-t5-xxl-fp16.tar.gz",
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/huggingface-"
+            "pytorch-inference:1.13.1-transformers4.26.0-gpu-py39-cu117-ubuntu20.04",
+            model_data="s3://jumpstart-cache-prod-us-west-2/huggingface-infer/"
+            "prepack/v1.1.2/infer-prepack-huggingface-text2text-flan-t5-xxl-fp16.tar.gz",
             env={
                 "SAGEMAKER_PROGRAM": "inference.py",
                 "ENDPOINT_SERVER_TIMEOUT": "3600",
@@ -1463,10 +1471,12 @@ def test_model_artifact_variant_model(
         JumpStartModel(model_id=model_id, instance_type="ml.p3.2xlarge")
 
         mock_model_init.assert_called_once_with(
-            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.10.0-gpu-py38",
+            image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/"
+            "pytorch-inference:1.10.0-gpu-py38",
             model_data={
                 "S3DataSource": {
-                    "S3Uri": "s3://jumpstart-cache-prod-us-west-2/pytorch-ic/pytorch-ic-mobilenet-v2/artifacts/inference-prepack/v1.0.0/",
+                    "S3Uri": "s3://jumpstart-cache-prod-us-west-2/"
+                    "pytorch-ic/pytorch-ic-mobilenet-v2/artifacts/inference-prepack/v1.0.0/",
                     "S3DataType": "S3Prefix",
                     "CompressionType": "None",
                 }
@@ -1949,7 +1959,8 @@ def test_model_deployment_config_additional_model_data_source(
             image_uri="763104351884.dkr.ecr.us-west-2.amazonaws.com/pytorch-inference:1.10.0-cpu-py38",
             model_data={
                 "S3DataSource": {
-                    "S3Uri": "s3://jumpstart-cache-prod-us-west-2/pytorch-ic/pytorch-ic-mobilenet-v2/artifacts/inference-prepack/v1.0.0/",
+                    "S3Uri": "s3://jumpstart-cache-prod-us-west-2/pytorch-ic"
+                    "/pytorch-ic-mobilenet-v2/artifacts/inference-prepack/v1.0.0/",
                     "S3DataType": "S3Prefix",
                     "CompressionType": "None",
                 }

tests/unit/sagemaker/jumpstart/test_artifacts.py

@@ -34,7 +34,7 @@
 from sagemaker.jumpstart.artifacts.model_uris import _retrieve_model_uri
 from sagemaker.jumpstart.enums import JumpStartScriptScope, JumpStartModelType
 
-from tests.unit.sagemaker.jumpstart.utils import get_spec_from_base_spec, get_special_model_spec
+from tests.unit.sagemaker.jumpstart.utils import get_special_model_spec
 from tests.unit.sagemaker.workflow.conftest import mock_client