feature: add wait argument to estimator 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, **kwargs):
+               use_compiled_model=False, update_endpoint=False, wait=True, **kwargs):
         """Deploy the trained model to an Amazon SageMaker endpoint and return a ``sagemaker.RealTimePredictor`` object.
 
         More information:
@@ -355,6 +355,7 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
                     >>> 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.
@@ -381,7 +382,8 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             accelerator_type=accelerator_type,
             endpoint_name=endpoint_name,
             update_endpoint=update_endpoint,
-            tags=self.tags)
+            tags=self.tags,
+            wait=wait)
 
     @property
     def model_data(self):

src/sagemaker/model.py

@@ -228,7 +228,7 @@ def compile(self, target_instance_family, input_shape, output_path, role,
         return self
 
     def deploy(self, initial_instance_count, instance_type, accelerator_type=None, endpoint_name=None,
-               update_endpoint=False, tags=None, kms_key=None):
+               update_endpoint=False, tags=None, kms_key=None, wait=True):
         """Deploy this ``Model`` to an ``Endpoint`` and optionally return a ``Predictor``.
 
         Create a SageMaker ``Model`` and ``EndpointConfig``, and deploy an ``Endpoint`` from this ``Model``.
@@ -256,6 +256,7 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             tags(List[dict[str, str]]): The list of tags to attach to this specific endpoint.
             kms_key (str): The ARN of the KMS key that is used to encrypt the data on the
                 storage volume attached to the instance hosting the endpoint.
+            wait (bool): Whether the call should wait until the deployment of this model completes (default: True).
 
         Returns:
             callable[string, sagemaker.session.Session] or None: Invocation of ``self.predictor_cls`` on
@@ -296,7 +297,7 @@ def deploy(self, initial_instance_count, instance_type, accelerator_type=None, e
             self.sagemaker_session.update_endpoint(self.endpoint_name, endpoint_config_name)
         else:
             self.sagemaker_session.endpoint_from_production_variants(self.endpoint_name, [production_variant],
-                                                                     tags, kms_key)
+                                                                     tags, kms_key, wait)
 
         if self.predictor_cls:
             return self.predictor_cls(self.endpoint_name, self.sagemaker_session)

tests/unit/test_estimator.py

@@ -922,7 +922,8 @@ def test_fit_deploy_keep_tags(sagemaker_session):
     sagemaker_session.endpoint_from_production_variants.assert_called_with(job_name,
                                                                            variant,
                                                                            tags,
-                                                                           None)
+                                                                           None,
+                                                                           True)
 
     sagemaker_session.create_model.assert_called_with(
         ANY,

tests/unit/test_model.py

@@ -166,7 +166,8 @@ def test_deploy(sagemaker_session, tmpdir):
           'InitialInstanceCount': 1,
           'VariantName': 'AllTraffic'}],
         None,
-        None)
+        None,
+        True)
 
 
 @patch('sagemaker.fw_utils.tar_and_upload_dir', MagicMock())
@@ -182,7 +183,8 @@ def test_deploy_endpoint_name(sagemaker_session, tmpdir):
           'InitialInstanceCount': 55,
           'VariantName': 'AllTraffic'}],
         None,
-        None)
+        None,
+        True)
 
 
 @patch('sagemaker.fw_utils.tar_and_upload_dir', MagicMock())
@@ -199,7 +201,8 @@ def test_deploy_tags(sagemaker_session, tmpdir):
           'InitialInstanceCount': 1,
           'VariantName': 'AllTraffic'}],
         tags,
-        None)
+        None,
+        True)
 
 
 @patch('sagemaker.fw_utils.tar_and_upload_dir', MagicMock())
@@ -217,7 +220,8 @@ def test_deploy_accelerator_type(tfo, time, sagemaker_session):
           'VariantName': 'AllTraffic',
           'AcceleratorType': ACCELERATOR_TYPE}],
         None,
-        None)
+        None,
+        True)
 
 
 @patch('sagemaker.fw_utils.tar_and_upload_dir', MagicMock())
@@ -235,7 +239,8 @@ def test_deploy_kms_key(tfo, time, sagemaker_session):
           'InitialInstanceCount': 1,
           'VariantName': 'AllTraffic'}],
         None,
-        key)
+        key,
+        True)
 
 
 @patch('sagemaker.session.Session')