feature: allow custom model name during deploy

Author: imujjwal96

src/sagemaker/estimator.py

@@ -328,7 +328,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, **kwargs):
+               model_name=None, use_compiled_model=False, update_endpoint=False, **kwargs):
         """Deploy the trained model to an Amazon SageMaker endpoint and return a ``sagemaker.RealTimePredictor`` object.
 
         More information:
@@ -344,6 +344,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
                 For more information: https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html
             endpoint_name (str): Name to use for creating an Amazon SageMaker endpoint. If not specified, the name of
                 the training job is used.
+            model_name (str): Name to use for creating an Amazon SageMaker model. If not specified, the name of
+                the training job is used.
             use_compiled_model (bool): Flag to select whether to use compiled (optimized) model. Default: False.
             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
@@ -363,6 +365,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])
@@ -372,6 +375,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/unit/test_estimator.py

@@ -1131,6 +1131,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):