feature: allow custom model name during deploy

Author: imujjwal96

src/sagemaker/estimator.py

@@ -331,7 +331,7 @@ def attach(cls, training_job_name, sagemaker_session=None, model_channel_name='m
         return estimator
 
     def deploy(self, initial_instance_count, instance_type, accelerator_type=None, endpoint_name=None,
-               use_compiled_model=False, update_endpoint=False, wait=True, **kwargs):
+               use_compiled_model=False, update_endpoint=False, wait=True, model_name=None, **kwargs):
         """Deploy the trained model to an Amazon SageMaker endpoint and return a ``sagemaker.RealTimePredictor`` object.
 
         More information:
@@ -351,11 +351,13 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             update_endpoint (bool): Flag to update the model in an existing Amazon SageMaker endpoint.
                 If True, this will deploy a new EndpointConfig to an already existing endpoint and delete resources
                 corresponding to the previous EndpointConfig. Default: False
+            wait (bool): Whether the call should wait until the deployment of model completes (default: True).
+            model_name (str): Name to use for creating an Amazon SageMaker model. If not specified, the name of
+                the training job is used.
             tags(List[dict[str, str]]): Optional. The list of tags to attach to this specific endpoint. 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
-            wait (bool): Whether the call should wait until the deployment of model completes (default: True).
 
             **kwargs: Passed to invocation of ``create_model()``. Implementations may customize
                 ``create_model()`` to accept ``**kwargs`` to customize model creation during deploy.
@@ -367,6 +369,7 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
         """
         self._ensure_latest_training_job()
         endpoint_name = endpoint_name or self.latest_training_job.name
+        model_name = model_name or self.latest_training_job.name
         self.deploy_instance_type = instance_type
         if use_compiled_model:
             family = '_'.join(instance_type.split('.')[:-1])
@@ -376,6 +379,7 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             model = self._compiled_models[family]
         else:
             model = self.create_model(**kwargs)
+        model.name = model_name
         return model.deploy(
             instance_type=instance_type,
             initial_instance_count=initial_instance_count,

tests/integ/test_tf_script_mode.py

@@ -136,19 +136,21 @@ def test_mnist_async(sagemaker_session):
     training_job_name = estimator.latest_training_job.name
     time.sleep(20)
     endpoint_name = training_job_name
+    model_name = 'model-name-1'
     _assert_training_job_tags_match(sagemaker_session.sagemaker_client,
                                     estimator.latest_training_job.name, TAGS)
     with tests.integ.timeout.timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
         estimator = TensorFlow.attach(training_job_name=training_job_name,
                                       sagemaker_session=sagemaker_session)
         predictor = estimator.deploy(initial_instance_count=1, instance_type='ml.c4.xlarge',
-                                     endpoint_name=endpoint_name)
+                                     endpoint_name=endpoint_name, model_name=model_name)
 
         result = predictor.predict(np.zeros(784))
         print('predict result: {}'.format(result))
         _assert_endpoint_tags_match(sagemaker_session.sagemaker_client, predictor.endpoint, TAGS)
         _assert_model_tags_match(sagemaker_session.sagemaker_client,
                                  estimator.latest_training_job.name, TAGS)
+        _assert_model_name_match(sagemaker_session.sagemaker_client, endpoint_name, model_name)
 
 
 def test_deploy_with_input_handlers(sagemaker_session, instance_type):
@@ -208,3 +210,8 @@ def _assert_training_job_tags_match(sagemaker_client, training_job_name, tags):
     training_job_description = sagemaker_client.describe_training_job(
         TrainingJobName=training_job_name)
     _assert_tags_match(sagemaker_client, training_job_description['TrainingJobArn'], tags)
+
+
+def _assert_model_name_match(sagemaker_client, endpoint_config_name, model_name):
+    endpoint_config_description = sagemaker_client.describe_endpoint(EndpointConfigName=endpoint_config_name)
+    assert model_name == endpoint_config_description['ProductionVariants'][0]['ModelName']

tests/unit/test_estimator.py

@@ -1140,6 +1140,31 @@ def test_deploy_with_update_endpoint(sagemaker_session):
     sagemaker_session.create_endpoint.assert_not_called()
 
 
+def test_deploy_with_model_name(sagemaker_session):
+    estimator = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH,
+                          sagemaker_session=sagemaker_session)
+    estimator.set_hyperparameters(**HYPERPARAMS)
+    estimator.fit({'train': 's3://bucket/training-prefix'})
+    model_name = 'model-name'
+    estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE, model_name=model_name)
+
+    sagemaker_session.create_model.assert_called_once()
+    args, kwargs = sagemaker_session.create_model.call_args
+    assert args[0] == model_name
+
+
+def test_deploy_with_no_model_name(sagemaker_session):
+    estimator = Estimator(IMAGE_NAME, ROLE, INSTANCE_COUNT, INSTANCE_TYPE, output_path=OUTPUT_PATH,
+                          sagemaker_session=sagemaker_session)
+    estimator.set_hyperparameters(**HYPERPARAMS)
+    estimator.fit({'train': 's3://bucket/training-prefix'})
+    estimator.deploy(INSTANCE_COUNT, INSTANCE_TYPE)
+
+    sagemaker_session.create_model.assert_called_once()
+    args, kwargs = sagemaker_session.create_model.call_args
+    assert args[0].startswith(IMAGE_NAME)
+
+
 @patch('sagemaker.estimator.LocalSession')
 @patch('sagemaker.estimator.Session')
 def test_local_mode(session_class, local_session_class):

tests/unit/test_tuner.py

@@ -566,7 +566,8 @@ def test_deploy_default(tuner):
 
     tuner.estimator.sagemaker_session.create_model.assert_called_once()
     args = tuner.estimator.sagemaker_session.create_model.call_args[0]
-    assert args[0].startswith(IMAGE_NAME)
+
+    assert args[0] == 'neo'
     assert args[1] == ROLE
     assert args[2]['Image'] == IMAGE_NAME
     assert args[2]['ModelDataUrl'] == MODEL_DATA