fix: tags for jumpstart model package models (#4061)

Author: evakravi

src/sagemaker/jumpstart/model.py

@@ -310,13 +310,34 @@ def _is_valid_model_id_hook():
 
         super(JumpStartModel, self).__init__(**model_init_kwargs.to_kwargs_dict())
 
-    def _create_sagemaker_model(self, *args, **kwargs):  # pylint: disable=unused-argument
+    def _create_sagemaker_model(
+        self,
+        instance_type=None,
+        accelerator_type=None,
+        tags=None,
+        serverless_inference_config=None,
+        **kwargs,
+    ):
         """Create a SageMaker Model Entity
 
         Args:
-            args: Positional arguments coming from the caller. This class does not require
-                any so they are ignored.
-
+            instance_type (str): Optional. The EC2 instance type that this Model will be
+                used for, this is only used to determine if the image needs GPU
+                support or not. (Default: None).
+            accelerator_type (str): Optional. Type of Elastic Inference accelerator to
+                attach to an endpoint for model loading and inference, for
+                example, 'ml.eia1.medium'. If not specified, no Elastic
+                Inference accelerator will be attached to the endpoint. (Default: None).
+            tags (List[dict[str, str]]): Optional. The list of tags to add to
+                the model. Example: >>> tags = [{'Key': 'tagname', 'Value':
+                'tagvalue'}] For more information about tags, see
+                https://boto3.amazonaws.com/v1/documentation
+                /api/latest/reference/services/sagemaker.html#SageMaker.Client.add_tags
+                (Default: None).
+            serverless_inference_config (sagemaker.serverless.ServerlessInferenceConfig):
+                Optional. Specifies configuration related to serverless endpoint. Instance type is
+                not provided in serverless inference. So this is used to find image URIs.
+                (Default: None).
             kwargs: Keyword arguments coming from the caller. This class does not require
                 any so they are ignored.
         """
@@ -347,10 +368,16 @@ def _create_sagemaker_model(self, *args, **kwargs):  # pylint: disable=unused-ar
                 container_def,
                 vpc_config=self.vpc_config,
                 enable_network_isolation=self.enable_network_isolation(),
-                tags=kwargs.get("tags"),
+                tags=tags,
             )
         else:
-            super(JumpStartModel, self)._create_sagemaker_model(*args, **kwargs)
+            super(JumpStartModel, self)._create_sagemaker_model(
+                instance_type=instance_type,
+                accelerator_type=accelerator_type,
+                tags=tags,
+                serverless_inference_config=serverless_inference_config,
+                **kwargs,
+            )
 
     def deploy(
         self,

src/sagemaker/model.py

@@ -1375,7 +1375,10 @@ def deploy(
                 self._base_name = "-".join((self._base_name, compiled_model_suffix))
 
         self._create_sagemaker_model(
-            instance_type, accelerator_type, tags, serverless_inference_config
+            instance_type=instance_type,
+            accelerator_type=accelerator_type,
+            tags=tags,
+            serverless_inference_config=serverless_inference_config,
         )
 
         serverless_inference_config_dict = (

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

@@ -589,7 +589,10 @@ def test_jumpstart_model_package_arn(
 
         model = JumpStartModel(model_id=model_id)
 
-        model.deploy()
+        tag = {"Key": "foo", "Value": "bar"}
+        tags = [tag]
+
+        model.deploy(tags=tags)
 
         self.assertEqual(
             mock_session.return_value.create_model.call_args[0][2],
@@ -599,6 +602,8 @@ def test_jumpstart_model_package_arn(
             },
         )
 
+        self.assertIn(tag, mock_session.return_value.create_model.call_args[1]["tags"])
+
     @mock.patch("sagemaker.jumpstart.model.is_valid_model_id")
     @mock.patch("sagemaker.jumpstart.factory.model.Session")
     @mock.patch("sagemaker.jumpstart.accessors.JumpStartModelsAccessor.get_model_specs")

tests/unit/sagemaker/model/test_deploy.py

@@ -159,7 +159,12 @@ def test_deploy_accelerator_type(
         accelerator_type=ACCELERATOR_TYPE,
     )
 
-    create_sagemaker_model.assert_called_with(INSTANCE_TYPE, ACCELERATOR_TYPE, None, None)
+    create_sagemaker_model.assert_called_with(
+        instance_type=INSTANCE_TYPE,
+        accelerator_type=ACCELERATOR_TYPE,
+        tags=None,
+        serverless_inference_config=None,
+    )
     production_variant.assert_called_with(
         MODEL_NAME,
         INSTANCE_TYPE,
@@ -271,7 +276,12 @@ def test_deploy_tags(create_sagemaker_model, production_variant, name_from_base,
     tags = [{"Key": "ModelName", "Value": "TestModel"}]
     model.deploy(instance_type=INSTANCE_TYPE, initial_instance_count=INSTANCE_COUNT, tags=tags)
 
-    create_sagemaker_model.assert_called_with(INSTANCE_TYPE, None, tags, None)
+    create_sagemaker_model.assert_called_with(
+        instance_type=INSTANCE_TYPE,
+        accelerator_type=None,
+        tags=tags,
+        serverless_inference_config=None,
+    )
     sagemaker_session.endpoint_from_production_variants.assert_called_with(
         name=ENDPOINT_NAME,
         production_variants=[BASE_PRODUCTION_VARIANT],
@@ -463,7 +473,12 @@ def test_deploy_serverless_inference(production_variant, create_sagemaker_model,
         serverless_inference_config=serverless_inference_config,
     )
 
-    create_sagemaker_model.assert_called_with(None, None, None, serverless_inference_config)
+    create_sagemaker_model.assert_called_with(
+        instance_type=None,
+        accelerator_type=None,
+        tags=None,
+        serverless_inference_config=serverless_inference_config,
+    )
     production_variant.assert_called_with(
         MODEL_NAME,
         None,

tests/unit/sagemaker/model/test_model_package.py

@@ -197,7 +197,7 @@ def test_create_sagemaker_model_include_tags(sagemaker_session):
         sagemaker_session=sagemaker_session,
     )
 
-    model_package._create_sagemaker_model(tags=tags)
+    model_package.deploy(tags=tags, instance_type="ml.p2.xlarge", initial_instance_count=1)
 
     sagemaker_session.create_model.assert_called_with(
         model_name,