fix: tags for jumpstart model package models

Author: evakravi

src/sagemaker/jumpstart/model.py

@@ -310,13 +310,32 @@ 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): 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.
+            accelerator_type (str): 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.
+            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
+            serverless_inference_config (sagemaker.serverless.ServerlessInferenceConfig):
+                Specifies configuration related to serverless endpoint. Instance type is
+                not provided in serverless inference. So this is used to find image URIs.
             kwargs: Keyword arguments coming from the caller. This class does not require
                 any so they are ignored.
         """
@@ -347,10 +366,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,

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")