feature: allow custom model name during deploy (#792)

* feature: allow custom model name during deploy

* black check

Author: imujjwal96

src/sagemaker/estimator.py

@@ -392,6 +392,7 @@ def deploy(
         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.
@@ -413,11 +414,13 @@ def deploy(
             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.
@@ -429,6 +432,7 @@ def deploy(
         """
         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])
@@ -440,6 +444,7 @@ def deploy(
             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,

src/sagemaker/tuner.py

@@ -375,6 +375,7 @@ def deploy(
         accelerator_type=None,
         endpoint_name=None,
         wait=True,
+        model_name=None,
         **kwargs
     ):
         """Deploy the best trained or user specified model to an Amazon SageMaker endpoint and return a
@@ -393,6 +394,8 @@ def deploy(
             endpoint_name (str): Name to use for creating an Amazon SageMaker endpoint. If not specified,
                 the name of the training job is used.
             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.
             **kwargs: Other arguments needed for deployment. Please refer to the ``create_model()`` method of
                 the associated estimator to see what other arguments are needed.
 
@@ -410,6 +413,7 @@ def deploy(
             accelerator_type=accelerator_type,
             endpoint_name=endpoint_name,
             wait=wait,
+            model_name=model_name,
             **kwargs
         )
 

tests/integ/test_tf_script_mode.py

@@ -159,6 +159,7 @@ 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
     )
@@ -167,7 +168,10 @@ def test_mnist_async(sagemaker_session):
             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
+            initial_instance_count=1,
+            instance_type="ml.c4.xlarge",
+            endpoint_name=endpoint_name,
+            model_name=model_name,
         )
 
         result = predictor.predict(np.zeros(784))
@@ -176,6 +180,7 @@ def test_mnist_async(sagemaker_session):
         _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):
@@ -241,3 +246,10 @@ def _assert_training_job_tags_match(sagemaker_client, training_job_name, tags):
         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_config(
+        EndpointConfigName=endpoint_config_name
+    )
+    assert model_name == endpoint_config_description["ProductionVariants"][0]["ModelName"]

tests/integ/test_tuner.py

@@ -843,14 +843,17 @@ def test_attach_tuning_pytorch(sagemaker_session):
         time.sleep(15)
         tuner.wait()
 
+    endpoint_name = tuning_job_name
+    model_name = "model-name-1"
     attached_tuner = HyperparameterTuner.attach(
         tuning_job_name, sagemaker_session=sagemaker_session
     )
     assert attached_tuner.early_stopping_type == "Auto"
 
-    best_training_job = tuner.best_training_job()
-    with timeout_and_delete_endpoint_by_name(best_training_job, sagemaker_session):
-        predictor = attached_tuner.deploy(1, "ml.c4.xlarge")
+    with timeout_and_delete_endpoint_by_name(endpoint_name, sagemaker_session):
+        predictor = attached_tuner.deploy(
+            1, "ml.c4.xlarge", endpoint_name=endpoint_name, model_name=model_name
+        )
         data = np.zeros(shape=(1, 1, 28, 28), dtype=np.float32)
         predictor.predict(data)
 
@@ -859,6 +862,7 @@ def test_attach_tuning_pytorch(sagemaker_session):
         output = predictor.predict(data)
 
         assert output.shape == (batch_size, 10)
+        _assert_model_name_match(sagemaker_session.sagemaker_client, endpoint_name, model_name)
 
 
 @pytest.mark.canary_quick
@@ -941,3 +945,10 @@ def _fm_serializer(data):
     for row in data:
         js["instances"].append({"features": row.tolist()})
     return json.dumps(js)
+
+
+def _assert_model_name_match(sagemaker_client, endpoint_config_name, model_name):
+    endpoint_config_description = sagemaker_client.describe_endpoint_config(
+        EndpointConfigName=endpoint_config_name
+    )
+    assert model_name == endpoint_config_description["ProductionVariants"][0]["ModelName"]

tests/unit/test_estimator.py

@@ -1711,6 +1711,43 @@ 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

@@ -645,7 +645,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