feature: Add support for accelerator in Clarify (#2249)

Co-authored-by: Ajay Karpur <[email protected]>
Co-authored-by: Ahsan Khan <[email protected]>

Author: 3 people

src/sagemaker/clarify.py

@@ -123,6 +123,7 @@ def __init__(
         content_type=None,
         content_template=None,
         custom_attributes=None,
+        accelerator_type=None,
     ):
         """Initializes a configuration of a model and the endpoint to be created for it.
 
@@ -151,6 +152,9 @@ def __init__(
                 Section 3.3.6. Field Value Components (
                 https://tools.ietf.org/html/rfc7230#section-3.2.6) of the Hypertext Transfer
                 Protocol (HTTP/1.1).
+            accelerator_type (str): The Elastic Inference accelerator type to deploy to the model
+                endpoint instance for making inferences to the model, see
+                https://docs.aws.amazon.com/sagemaker/latest/dg/ei.html.
         """
         self.predictor_config = {
             "model_name": model_name,
@@ -178,9 +182,8 @@ def __init__(
                     f" Please include a placeholder $features."
                 )
             self.predictor_config["content_template"] = content_template
-
-        if custom_attributes is not None:
-            self.predictor_config["custom_attributes"] = custom_attributes
+        _set(custom_attributes, "custom_attributes", self.predictor_config)
+        _set(accelerator_type, "accelerator_type", self.predictor_config)
 
     def get_predictor_config(self):
         """Returns part of the predictor dictionary of the analysis config."""

tests/unit/test_clarify.py

@@ -92,13 +92,15 @@ def test_model_config():
     accept_type = "text/csv"
     content_type = "application/jsonlines"
     custom_attributes = "c000b4f9-df62-4c85-a0bf-7c525f9104a4"
+    accelerator_type = "ml.eia1.medium"
     model_config = ModelConfig(
         model_name=model_name,
         instance_type=instance_type,
         instance_count=instance_count,
         accept_type=accept_type,
         content_type=content_type,
         custom_attributes=custom_attributes,
+        accelerator_type=accelerator_type,
     )
     expected_config = {
         "model_name": model_name,
@@ -107,6 +109,7 @@ def test_model_config():
         "accept_type": accept_type,
         "content_type": content_type,
         "custom_attributes": custom_attributes,
+        "accelerator_type": accelerator_type,
     }
     assert expected_config == model_config.get_predictor_config()