File tree Expand file tree Collapse file tree 2 files changed +11
-4
lines changed Expand file tree Collapse file tree 2 files changed +11
-4
lines changed Original file line number Diff line number Diff line change @@ -498,7 +498,7 @@ def transformer(
498
498
volume_kms_key (str): Optional. KMS key ID for encrypting the volume
499
499
attached to the ML compute instance (default: None).
500
500
"""
501
- self ._create_sagemaker_model (instance_type )
501
+ self ._create_sagemaker_model (instance_type , tags = tags )
502
502
if self .enable_network_isolation ():
503
503
env = None
504
504
@@ -895,11 +895,14 @@ def _is_marketplace(self):
895
895
return True
896
896
return False
897
897
898
- def _create_sagemaker_model (self , * args ): # pylint: disable=unused-argument
898
+ def _create_sagemaker_model (self , * args , ** kwargs ): # pylint: disable=unused-argument
899
899
"""Create a SageMaker Model Entity
900
900
901
901
Args:
902
- *args: Arguments coming from the caller. This class does not require
902
+ args: Positional arguments coming from the caller. This class does not require
903
+ any so they are ignored.
904
+
905
+ kwargs: Keyword arguments coming from the caller. This class does not require
903
906
any so they are ignored.
904
907
"""
905
908
if self .algorithm_arn :
Original file line number Diff line number Diff line change @@ -361,16 +361,20 @@ def test_model_create_transformer(sagemaker_session):
361
361
return_value = DESCRIBE_MODEL_PACKAGE_RESPONSE
362
362
)
363
363
364
+ tags = [{"Key" : "k" , "Value" : "v" }]
364
365
model = DummyFrameworkModel (sagemaker_session = sagemaker_session )
366
+ instance_type = "ml.m4.xlarge"
365
367
model .name = "auto-generated-model"
366
368
transformer = model .transformer (
367
- instance_count = 1 , instance_type = "ml.m4.xlarge" , env = {"test" : True }
369
+ instance_count = 1 , instance_type = instance_type , env = {"test" : True }, tags = tags
368
370
)
369
371
assert isinstance (transformer , sagemaker .transformer .Transformer )
370
372
assert transformer .model_name == "auto-generated-model"
371
373
assert transformer .instance_type == "ml.m4.xlarge"
372
374
assert transformer .env == {"test" : True }
373
375
376
+ sagemaker .model .Model ._create_sagemaker_model .assert_called_with (instance_type , tags = tags )
377
+
374
378
375
379
def test_model_package_enable_network_isolation_with_no_product_id (sagemaker_session ):
376
380
sagemaker_session .sagemaker_client .describe_model_package = Mock (
You can’t perform that action at this time.
0 commit comments