feature: Add ModelClientConfig Fields for Batch Transform

Author: guoqiao1992

src/sagemaker/session.py

@@ -1980,6 +1980,7 @@ def transform(
         output_config,
         resource_config,
         experiment_config,
+        model_client_config,
         tags,
         data_processing,
     ):
@@ -2002,6 +2003,9 @@ def transform(
             experiment_config (dict): A dictionary describing the experiment configuration for the
                 job. Dictionary contains three optional keys,
                 'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
+            model_client_config (dict): A dictionary describing the model configuration for the
+                job. Dictionary contains two optional keys,
+                'InvocationsTimeoutInSeconds', and 'InvocationsMaxRetries'.
             tags (list[dict]): List of tags for labeling a transform job.
             data_processing(dict): A dictionary describing config for combining the input data and
                 transformed data. For more, see
@@ -2036,6 +2040,9 @@ def transform(
         if experiment_config and len(experiment_config) > 0:
             transform_request["ExperimentConfig"] = experiment_config
 
+        if model_client_config and len(model_client_config) > 0:
+            transform_request["ModelClientConfig"] = model_client_config
+
         LOGGER.info("Creating transform job with name: %s", job_name)
         LOGGER.debug("Transform request: %s", json.dumps(transform_request, indent=4))
         self.sagemaker_client.create_transform_job(**transform_request)

src/sagemaker/transformer.py

@@ -120,6 +120,7 @@ def transform(
         output_filter=None,
         join_source=None,
         experiment_config=None,
+        model_client_config=None,
         wait=False,
         logs=False,
     ):
@@ -171,6 +172,10 @@ def transform(
                 Dictionary contains three optional keys,
                 'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
                 (default: ``None``).
+            model_client_config (dict[str, str]): Model configuration.
+                Dictionary contains two optional keys,
+                'InvocationsTimeoutInSeconds', and 'InvocationsMaxRetries'.
+                (default: ``None``).
             wait (bool): Whether the call should wait until the job completes
                 (default: False).
             logs (bool): Whether to show the logs produced by the job.
@@ -207,6 +212,7 @@ def transform(
             output_filter,
             join_source,
             experiment_config,
+            model_client_config,
         )
 
         if wait:
@@ -341,6 +347,7 @@ def start_new(
         output_filter,
         join_source,
         experiment_config,
+        model_client_config,
     ):
         """
         Args:
@@ -354,6 +361,7 @@ def start_new(
             output_filter:
             join_source:
             experiment_config:
+            model_client_config:
         """
         config = _TransformJob._load_config(
             data, data_type, content_type, compression_type, split_type, transformer
@@ -373,6 +381,7 @@ def start_new(
             output_config=config["output_config"],
             resource_config=config["resource_config"],
             experiment_config=experiment_config,
+            model_client_config=model_client_config,
             tags=transformer.tags,
             data_processing=data_processing,
         )

tests/integ/test_transformer.py

@@ -229,6 +229,26 @@ def test_transform_mxnet_tags(
         assert tags == model_tags
 
 
+def test_transform_model_client_config(
+    mxnet_estimator, mxnet_transform_input, sagemaker_session, cpu_instance_type
+):
+    model_client_config = {"InvocationsTimeoutInSeconds": 60, "InvocationsMaxRetries": 2}
+    transformer = mxnet_estimator.transformer(1, cpu_instance_type)
+    transformer.transform(
+        mxnet_transform_input, content_type="text/csv", model_client_config=model_client_config
+    )
+
+    with timeout_and_delete_model_with_transformer(
+        transformer, sagemaker_session, minutes=TRANSFORM_DEFAULT_TIMEOUT_MINUTES
+    ):
+        transformer.wait()
+        transform_job_desc = sagemaker_session.sagemaker_client.describe_transform_job(
+            TransformJobName=transformer.latest_transform_job.name
+        )
+
+        assert model_client_config == transform_job_desc["ModelClientConfig"]
+
+
 def test_transform_byo_estimator(sagemaker_session, cpu_instance_type):
     data_path = os.path.join(DATA_DIR, "one_p_mnist")
     pickle_args = {} if sys.version_info.major == 2 else {"encoding": "latin1"}

tests/unit/test_session.py

@@ -562,6 +562,7 @@ def test_s3_input_all_arguments():
     "TrialName": "dummyT",
     "TrialComponentDisplayName": "dummyTC",
 }
+MODEL_CLIENT_CONFIG = {"InvocationsMaxRetries": 2, "InvocationsTimeoutInSeconds": 60}
 
 DEFAULT_EXPECTED_TRAIN_JOB_ARGS = {
     "OutputDataConfig": {"S3OutputPath": S3_OUTPUT},
@@ -1244,6 +1245,7 @@ def test_transform_pack_to_request(sagemaker_session):
         output_config=out_config,
         resource_config=resource_config,
         experiment_config=None,
+        model_client_config=None,
         tags=None,
         data_processing=data_processing,
     )
@@ -1269,6 +1271,7 @@ def test_transform_pack_to_request_with_optional_params(sagemaker_session):
         output_config={},
         resource_config={},
         experiment_config=EXPERIMENT_CONFIG,
+        model_client_config=MODEL_CLIENT_CONFIG,
         tags=TAGS,
         data_processing=None,
     )
@@ -1280,6 +1283,7 @@ def test_transform_pack_to_request_with_optional_params(sagemaker_session):
     assert actual_args["Environment"] == env
     assert actual_args["Tags"] == TAGS
     assert actual_args["ExperimentConfig"] == EXPERIMENT_CONFIG
+    assert actual_args["ModelClientConfig"] == MODEL_CLIENT_CONFIG
 
 
 @patch("sys.stdout", new_callable=io.BytesIO if six.PY2 else io.StringIO)

tests/unit/test_transformer.py

@@ -167,6 +167,7 @@ def test_transform_with_all_params(start_new_job, transformer):
         "TrialName": "t",
         "TrialComponentDisplayName": "tc",
     }
+    model_client_config = {"InvocationsTimeoutInSeconds": 60, "InvocationsMaxRetries": 2}
 
     transformer.transform(
         DATA,
@@ -179,6 +180,7 @@ def test_transform_with_all_params(start_new_job, transformer):
         output_filter=output_filter,
         join_source=join_source,
         experiment_config=experiment_config,
+        model_client_config=model_client_config,
     )
 
     assert transformer._current_job_name == JOB_NAME
@@ -194,6 +196,7 @@ def test_transform_with_all_params(start_new_job, transformer):
         output_filter,
         join_source,
         experiment_config,
+        model_client_config,
     )
 
 
@@ -428,6 +431,8 @@ def test_start_new(prepare_data_processing, load_config, sagemaker_session):
     split_type = "Line"
     io_filter = "$"
     join_source = "Input"
+    model_client_config = {"InvocationsTimeoutInSeconds": 60, "InvocationsMaxRetries": 2}
+
     job = _TransformJob.start_new(
         transformer=transformer,
         data=DATA,
@@ -439,6 +444,7 @@ def test_start_new(prepare_data_processing, load_config, sagemaker_session):
         output_filter=io_filter,
         join_source=join_source,
         experiment_config={"ExperimentName": "exp"},
+        model_client_config=model_client_config,
     )
 
     assert job.sagemaker_session == sagemaker_session
@@ -460,6 +466,7 @@ def test_start_new(prepare_data_processing, load_config, sagemaker_session):
         output_config=output_config,
         resource_config=resource_config,
         experiment_config={"ExperimentName": "exp"},
+        model_client_config=model_client_config,
         tags=tags,
         data_processing=prepare_data_processing.return_value,
     )