update tests

Author: bencrabtree

src/sagemaker/jumpstart/constants.py

@@ -170,7 +170,7 @@
 
 JUMPSTART_DEFAULT_MANIFEST_FILE_S3_KEY = "models_manifest.json"
 
-HUB_CONTENT_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub_content/(.*?)/(.*?)/(.*?)/(.*?)$"
+HUB_CONTENT_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub-content/(.*?)/(.*?)/(.*?)/(.*?)$"
 HUB_ARN_REGEX = r"arn:(.*?):sagemaker:(.*?):(.*?):hub/(.*?)$"
 
 INFERENCE_ENTRY_POINT_SCRIPT_NAME = "inference.py"

src/sagemaker/jumpstart/types.py

@@ -867,7 +867,7 @@ def from_json(self, json_obj: Dict[str, Any]) -> None:
         self.inference_enable_network_isolation: bool = json_obj.get(
             "inference_enable_network_isolation", False
         )
-        self.resource_name_base: bool = json_obj.get("resource_name_base")
+        self.resource_name_base: Optional[str] = json_obj.get("resource_name_base")
 
         self.hosting_eula_key: Optional[str] = json_obj.get("hosting_eula_key")
 

src/sagemaker/jumpstart/utils.py

@@ -870,7 +870,7 @@ def get_info_from_hub_resource_arn(
             account_id=account_id,
             hub_name=hub_name,
         )
-
+    
     return None
 
 

tests/unit/sagemaker/accept_types/jumpstart/test_accept_types.py

@@ -45,7 +45,7 @@ def test_jumpstart_default_accept_types(
     assert default_accept_type == "application/json"
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version=model_version, s3_client=mock_client
+        region=region, model_id=model_id, version=model_version, s3_client=mock_client, hub_arn=None
     )
 
 
@@ -73,5 +73,5 @@ def test_jumpstart_supported_accept_types(
     ]
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version=model_version, s3_client=mock_client
+        region=region, model_id=model_id, version=model_version, s3_client=mock_client, hub_arn=None
     )

tests/unit/sagemaker/content_types/jumpstart/test_content_types.py

@@ -45,7 +45,7 @@ def test_jumpstart_default_content_types(
     assert default_content_type == "application/x-text"
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version=model_version, s3_client=mock_client
+        region=region, model_id=model_id, version=model_version, s3_client=mock_client, hub_arn=None
     )
 
 
@@ -72,5 +72,5 @@ def test_jumpstart_supported_content_types(
     ]
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version=model_version, s3_client=mock_client
+        region=region, model_id=model_id, version=model_version, s3_client=mock_client, hub_arn=None
     )

tests/unit/sagemaker/deserializers/jumpstart/test_deserializers.py

@@ -47,7 +47,7 @@ def test_jumpstart_default_deserializers(
     assert isinstance(default_deserializer, base_deserializers.JSONDeserializer)
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version=model_version, s3_client=mock_client
+        region=region, model_id=model_id, version=model_version, s3_client=mock_client, hub_arn=None
     )
 
 
@@ -79,5 +79,5 @@ def test_jumpstart_deserializer_options(
     )
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version=model_version, s3_client=mock_client
+        region=region, model_id=model_id, version=model_version, s3_client=mock_client, hub_arn=None
     )

tests/unit/sagemaker/environment_variables/jumpstart/test_default.py

@@ -48,7 +48,7 @@ def test_jumpstart_default_environment_variables(patched_get_model_specs):
     }
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version="*", s3_client=mock_client
+        region=region, model_id=model_id, version="*", s3_client=mock_client, hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()
@@ -68,7 +68,7 @@ def test_jumpstart_default_environment_variables(patched_get_model_specs):
     }
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version="1.*", s3_client=mock_client
+        region=region, model_id=model_id, version="1.*", s3_client=mock_client, hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()
@@ -122,7 +122,7 @@ def test_jumpstart_sdk_environment_variables(patched_get_model_specs):
     }
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version="*", s3_client=mock_client
+        region=region, model_id=model_id, version="*", s3_client=mock_client, hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()
@@ -143,7 +143,7 @@ def test_jumpstart_sdk_environment_variables(patched_get_model_specs):
     }
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version="1.*", s3_client=mock_client
+        region=region, model_id=model_id, version="1.*", s3_client=mock_client, hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()

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

@@ -47,6 +47,7 @@ def test_jumpstart_default_hyperparameters(patched_get_model_specs):
         model_id=model_id,
         version="*",
         s3_client=mock_client,
+        hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()
@@ -63,7 +64,7 @@ def test_jumpstart_default_hyperparameters(patched_get_model_specs):
         region=region,
         model_id=model_id,
         version="1.*",
-        s3_client=mock_client,
+        s3_client=mock_client, hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()
@@ -88,7 +89,7 @@ def test_jumpstart_default_hyperparameters(patched_get_model_specs):
         region=region,
         model_id=model_id,
         version="1.*",
-        s3_client=mock_client,
+        s3_client=mock_client, hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()

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

@@ -139,7 +139,7 @@ def add_options_to_hyperparameter(*largs, **kwargs):
         region=region,
         model_id=model_id,
         version=model_version,
-        s3_client=mock_client,
+        s3_client=mock_client, hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()
@@ -437,7 +437,7 @@ def add_options_to_hyperparameter(*largs, **kwargs):
     )
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version=model_version, s3_client=mock_client
+        region=region, model_id=model_id, version=model_version, s3_client=mock_client, hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()
@@ -491,7 +491,7 @@ def test_jumpstart_validate_all_hyperparameters(patched_get_model_specs):
     )
 
     patched_get_model_specs.assert_called_once_with(
-        region=region, model_id=model_id, version=model_version, s3_client=mock_client
+        region=region, model_id=model_id, version=model_version, s3_client=mock_client, hub_arn=None
     )
 
     patched_get_model_specs.reset_mock()

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

@@ -48,7 +48,7 @@ def test_jumpstart_common_image_uri(
         region="us-west-2",
         model_id="pytorch-ic-mobilenet-v2",
         version="*",
-        s3_client=mock_client,
+        s3_client=mock_client, hub_arn=None
     )
     patched_verify_model_region_and_return_specs.assert_called_once()
 
@@ -68,7 +68,7 @@ def test_jumpstart_common_image_uri(
         region="us-west-2",
         model_id="pytorch-ic-mobilenet-v2",
         version="1.*",
-        s3_client=mock_client,
+        s3_client=mock_client, hub_arn=None
     )
     patched_verify_model_region_and_return_specs.assert_called_once()
 
@@ -88,7 +88,7 @@ def test_jumpstart_common_image_uri(
         region=sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME,
         model_id="pytorch-ic-mobilenet-v2",
         version="*",
-        s3_client=mock_client,
+        s3_client=mock_client, hub_arn=None
     )
     patched_verify_model_region_and_return_specs.assert_called_once()
 
@@ -108,7 +108,7 @@ def test_jumpstart_common_image_uri(
         region=sagemaker_constants.JUMPSTART_DEFAULT_REGION_NAME,
         model_id="pytorch-ic-mobilenet-v2",
         version="1.*",
-        s3_client=mock_client,
+        s3_client=mock_client, hub_arn=None
     )
     patched_verify_model_region_and_return_specs.assert_called_once()
 

tests/unit/sagemaker/jumpstart/constants.py

@@ -6239,6 +6239,7 @@
     "training_volume_size": 456,
     "inference_enable_network_isolation": True,
     "training_enable_network_isolation": False,
+    "resource_name_base": "dfsdfsds",
     "hosting_resource_requirements": {"num_accelerators": 1, "min_memory_mb": 34360},
     "dynamic_container_deployment_supported": True,
 }

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

@@ -282,11 +282,12 @@ def test_prepacked(
             ],
         )
 
+    @mock.patch("sagemaker.jumpstart.artifacts.resource_names._retrieve_resource_name_base")
     @mock.patch("sagemaker.session.Session.account_id")
     @mock.patch("sagemaker.jumpstart.factory.model._retrieve_model_init_kwargs")
     @mock.patch("sagemaker.jumpstart.factory.model._retrieve_model_deploy_kwargs")
     @mock.patch("sagemaker.jumpstart.factory.estimator._retrieve_estimator_fit_kwargs")
-    @mock.patch("sagemaker.jumpstart.estimator.construct_hub_arn_from_name")
+    @mock.patch("sagemaker.jumpstart.utils.construct_hub_arn_from_name")
     @mock.patch("sagemaker.jumpstart.estimator.is_valid_model_id")
     @mock.patch("sagemaker.jumpstart.factory.model.Session")
     @mock.patch("sagemaker.jumpstart.factory.estimator.Session")
@@ -310,7 +311,9 @@ def test_hub_model(
         mock_retrieve_model_deploy_kwargs: mock.Mock,
         mock_retrieve_model_init_kwargs: mock.Mock,
         mock_get_caller_identity: mock.Mock,
+        mock_retrieve_resource_name_base: mock.Mock,
     ):
+        mock_retrieve_resource_name_base.return_value = "go-blue"
         mock_get_caller_identity.return_value = "123456789123"
         mock_estimator_deploy.return_value = default_predictor
 
@@ -372,11 +375,11 @@ def test_hub_model(
             f"some-training-dataset-doesn't-matter",
         }
 
-        estimator.fit(channels)
+        estimator.fit(channels, job_name="go-blue")
 
-        mock_estimator_fit.assert_called_once_with(inputs=channels, wait=True)
+        mock_estimator_fit.assert_called_once_with(inputs=channels, wait=True, job_name="go-blue")
 
-        estimator.deploy()
+        estimator.deploy(endpoint_name="go-blue", model_name="go-blue")
 
         mock_estimator_deploy.assert_called_once_with(
             instance_type="ml.p2.xlarge",
@@ -386,6 +389,8 @@ def test_hub_model(
             source_dir="s3://jumpstart-cache-prod-us-west-2/source-directory-tarballs/"
             "pytorch/inference/ic/v1.0.0/sourcedir.tar.gz",
             entry_point="inference.py",
+            endpoint_name="go-blue",
+            model_name="go-blue",
             env={
                 "SAGEMAKER_PROGRAM": "inference.py",
                 "ENDPOINT_SERVER_TIMEOUT": "3600",
@@ -414,7 +419,7 @@ def test_hub_model(
         )
 
         mock_construct_hub_arn_from_name.assert_called_once_with(
-            hub_name="my-mock-hub", region=None, sagemaker_session=None
+            hub_name="my-mock-hub", region=None, session=None
         )
 
     @mock.patch("sagemaker.utils.sagemaker_timestamp")
@@ -1475,6 +1480,7 @@ def test_incremental_training_with_unsupported_model_logs_warning(
             model_id=model_id,
             model_version="*",
             region=region,
+            hub_arn=None,
             tolerate_deprecated_model=False,
             tolerate_vulnerable_model=False,
             sagemaker_session=sagemaker_session,
@@ -1526,6 +1532,7 @@ def test_incremental_training_with_supported_model_doesnt_log_warning(
             model_id=model_id,
             model_version="*",
             region=region,
+            hub_arn=None,
             tolerate_deprecated_model=False,
             tolerate_vulnerable_model=False,
             sagemaker_session=sagemaker_session,