feature: Support Amazon SageMaker AutoMLStep

Co-authored-by: Xinlu Tu <[email protected]>
Co-authored-by: Xinlu Tu <[email protected]>

@xinlutu2
feat: Close feature gaps between Python SageMaker SDK and CreateAutoMLJob API includes ENSEMBLING mode

@xinlutu2
feature: add AutoMLStep for SageMaker Pipelines Workflows

@xinlutu2
feature: add AutoMLStep integration test

Author: 3 people

doc/workflows/pipelines/sagemaker.workflow.pipelines.rst

@@ -168,3 +168,5 @@ Steps
 .. autoclass:: sagemaker.workflow.clarify_check_step.ClarifyCheckStep
 
 .. autoclass:: sagemaker.workflow.fail_step.FailStep
+
+.. autoclass:: sagemaker.workflow.automl_step.AutoMLStep

setup.py

@@ -48,7 +48,7 @@ def read_requirements(filename):
 # Declare minimal set for installation
 required_packages = [
     "attrs>=20.3.0,<23",
-    "boto3>=1.20.21,<2.0",
+    "boto3>=1.24.69,<2.0",
     "google-pasta",
     "numpy>=1.9.0,<2.0",
     "protobuf>=3.1,<4.0",

src/sagemaker/automl/automl.py

@@ -63,3 +63,17 @@ class UnexpectedClientError(AsyncInferenceError):
 
     def __init__(self, message):
         super().__init__(message=message)
+
+
+class AutoMLStepInvalidModeError(Exception):
+    """Raised when the automl mode passed into AutoMLStep in invalid"""
+
+    fmt = (
+        "Mode in AutoMLJobConfig must be defined for AutoMLStep. "
+        "AutoMLStep currently only supports ENSEMBLING mode"
+    )
+
+    def __init__(self, **kwargs):
+        msg = self.fmt.format(**kwargs)
+        Exception.__init__(self, msg)
+        self.kwargs = kwargs

src/sagemaker/exceptions.py

@@ -1646,6 +1646,7 @@ def auto_ml(
         job_objective=None,
         generate_candidate_definitions_only=False,
         tags=None,
+        model_deploy_config=None,
     ):
         """Create an Amazon SageMaker AutoML job.
 
@@ -1669,6 +1670,71 @@ def auto_ml(
                 definitions. If True, AutoML.list_candidates() cannot be called. Default: False.
             tags ([dict[str,str]]): A list of dictionaries containing key-value
                 pairs.
+            model_deploy_config (dict): Specifies how to generate the endpoint name
+                for an automatic one-click Autopilot model deployment.
+                Contains "AutoGenerateEndpointName" and "EndpointName"
+        """
+        auto_ml_job_request = self._get_auto_ml_request(
+            input_config=input_config,
+            output_config=output_config,
+            auto_ml_job_config=auto_ml_job_config,
+            role=role,
+            job_name=job_name,
+            problem_type=problem_type,
+            job_objective=job_objective,
+            generate_candidate_definitions_only=generate_candidate_definitions_only,
+            tags=tags,
+            model_deploy_config=model_deploy_config,
+        )
+
+        def submit(request):
+            LOGGER.info("Creating auto-ml-job with name: %s", job_name)
+            LOGGER.debug("auto ml request: %s", json.dumps(request), indent=4)
+            self.sagemaker_client.create_auto_ml_job(**request)
+
+        self._intercept_create_request(auto_ml_job_request, submit, self.auto_ml.__name__)
+
+    def _get_auto_ml_request(
+        self,
+        input_config,
+        output_config,
+        auto_ml_job_config,
+        role,
+        job_name,
+        problem_type=None,
+        job_objective=None,
+        generate_candidate_definitions_only=False,
+        tags=None,
+        model_deploy_config=None,
+    ):
+        """Constructs a request compatible for creating an Amazon SageMaker AutoML job.
+
+        Args:
+            input_config (list[dict]): A list of Channel objects. Each channel contains "DataSource"
+                and "TargetAttributeName", "CompressionType" is an optional field.
+            output_config (dict): The S3 URI where you want to store the training results and
+                optional KMS key ID.
+            auto_ml_job_config (dict): A dict of AutoMLJob config, containing "StoppingCondition",
+                "SecurityConfig", optionally contains "VolumeKmsKeyId".
+            role (str): The Amazon Resource Name (ARN) of an IAM role that
+                Amazon SageMaker can assume to perform tasks on your behalf.
+            job_name (str): A string that can be used to identify an AutoMLJob. Each AutoMLJob
+                should have a unique job name.
+            problem_type (str): The type of problem of this AutoMLJob. Valid values are
+                "Regression", "BinaryClassification", "MultiClassClassification". If None,
+                SageMaker AutoMLJob will infer the problem type automatically.
+            job_objective (dict): AutoMLJob objective, contains "AutoMLJobObjectiveType" (optional),
+                "MetricName" and "Value".
+            generate_candidate_definitions_only (bool): Indicates whether to only generate candidate
+                definitions. If True, AutoML.list_candidates() cannot be called. Default: False.
+            tags ([dict[str,str]]): A list of dictionaries containing key-value
+                pairs.
+            model_deploy_config (dict): Specifies how to generate the endpoint name
+                for an automatic one-click Autopilot model deployment.
+                Contains "AutoGenerateEndpointName" and "EndpointName"
+
+        Returns:
+            Dict: a automl request dict
         """
         auto_ml_job_request = {
             "AutoMLJobName": job_name,
@@ -1678,6 +1744,8 @@ def auto_ml(
             "RoleArn": role,
             "GenerateCandidateDefinitionsOnly": generate_candidate_definitions_only,
         }
+        if model_deploy_config is not None:
+            auto_ml_job_request["ModelDeployConfig"] = model_deploy_config
 
         if job_objective is not None:
             auto_ml_job_request["AutoMLJobObjective"] = job_objective
@@ -1688,9 +1756,7 @@ def auto_ml(
         if tags is not None:
             auto_ml_job_request["Tags"] = tags
 
-        LOGGER.info("Creating auto-ml-job with name: %s", job_name)
-        LOGGER.debug("auto ml request: %s", json.dumps(auto_ml_job_request, indent=4))
-        self.sagemaker_client.create_auto_ml_job(**auto_ml_job_request)
+        return auto_ml_job_request
 
     def describe_auto_ml_job(self, job_name):
         """Calls the DescribeAutoMLJob API for the given job name and returns the response.

src/sagemaker/session.py

@@ -0,0 +1,166 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""The `AutoMLStep` definition for SageMaker Pipelines Workflows"""
+from __future__ import absolute_import
+
+from typing import Union, Optional, List
+
+from sagemaker import Session, Model
+from sagemaker.exceptions import AutoMLStepInvalidModeError
+from sagemaker.workflow.entities import RequestType
+
+from sagemaker.workflow.pipeline_context import _JobStepArguments
+from sagemaker.workflow.properties import Properties
+from sagemaker.workflow.retry import RetryPolicy
+from sagemaker.workflow.steps import ConfigurableRetryStep, CacheConfig, Step, StepTypeEnum
+from sagemaker.workflow.utilities import validate_step_args_input
+from sagemaker.workflow.step_collections import StepCollection
+
+
+class AutoMLStep(ConfigurableRetryStep):
+    """`AutoMLStep` for SageMaker Pipelines Workflows."""
+
+    def __init__(
+        self,
+        name: str,
+        step_args: _JobStepArguments,
+        display_name: str = None,
+        description: str = None,
+        cache_config: CacheConfig = None,
+        depends_on: Optional[List[Union[str, Step, StepCollection]]] = None,
+        retry_policies: List[RetryPolicy] = None,
+    ):
+        """Construct a `AutoMLStep`, given a `AutoML` instance.
+
+        In addition to the `AutoML` instance, the other arguments are those
+        that are supplied to the `fit` method of the `sagemaker.automl.automl.AutoML`.
+
+        Args:
+            name (str): The name of the `AutoMLStep`.
+            step_args (_JobStepArguments): The arguments for the `AutoMLStep` definition.
+            display_name (str): The display name of the `AutoMLStep`.
+            description (str): The description of the `AutoMLStep`.
+            cache_config (CacheConfig): A `sagemaker.workflow.steps.CacheConfig` instance.
+            depends_on (List[Union[str, Step, StepCollection]]): A list of `Step`/`StepCollection`
+                names or `Step` instances or `StepCollection` instances that this `AutoMLStep`
+                depends on.
+            retry_policies (List[RetryPolicy]): A list of retry policies.
+        """
+        super(AutoMLStep, self).__init__(
+            name, StepTypeEnum.AUTOML, display_name, description, depends_on, retry_policies
+        )
+
+        validate_step_args_input(
+            step_args=step_args,
+            expected_caller={Session.auto_ml.__name__},
+            error_message="The step_args of AutoMLStep must be obtained " "from automl.fit().",
+        )
+
+        self.step_args = step_args.args
+        self.cache_config = cache_config
+
+        root_property = Properties(step_name=name, shape_name="DescribeAutoMLJobResponse")
+
+        best_candidate_properties = Properties(step_name=name, path="bestCandidateProperties")
+        best_candidate_properties.__dict__["modelInsightsJsonReportPath"] = Properties(
+            step_name=name, path="bestCandidateProperties.modelInsightsJsonReportPath"
+        )
+        best_candidate_properties.__dict__["explainabilityJsonReportPath"] = Properties(
+            step_name=name, path="bestCandidateProperties.explainabilityJsonReportPath"
+        )
+
+        root_property.__dict__["bestCandidateProperties"] = best_candidate_properties
+        self._properties = root_property
+
+    @property
+    def arguments(self) -> RequestType:
+        """The arguments dictionary that is used to call `create_auto_ml_job`.
+
+        NOTE: The `CreateAutoMLJob` request is not quite the
+            args list that workflow needs.
+
+        The `AutoMLJobName`, `ModelDeployConfig` and `GenerateCandidateDefinitionsOnly`
+            attribute cannot be included.
+        """
+        request_dict = self.step_args
+        if "AutoMLJobConfig" not in request_dict:
+            raise AutoMLStepInvalidModeError()
+        if (
+            "Mode" not in request_dict["AutoMLJobConfig"]
+            or request_dict["AutoMLJobConfig"]["Mode"] != "ENSEMBLING"
+        ):
+            raise AutoMLStepInvalidModeError()
+
+        if "ModelDeployConfig" in request_dict:
+            request_dict.pop("ModelDeployConfig", None)
+        if "GenerateCandidateDefinitionsOnly" in request_dict:
+            request_dict.pop("GenerateCandidateDefinitionsOnly", None)
+        request_dict.pop("AutoMLJobName", None)
+        return request_dict
+
+    @property
+    def properties(self):
+        """A `Properties` object representing the `DescribeAutoMLJobResponse` data model."""
+        return self._properties
+
+    def to_request(self) -> RequestType:
+        """Updates the dictionary with cache configuration."""
+        request_dict = super().to_request()
+        if self.cache_config:
+            request_dict.update(self.cache_config.config)
+
+        return request_dict
+
+    def get_best_auto_ml_model(self, role, sagemaker_session=None):
+        """Get the best candidate model artifacts, image uri and env variables for the best model.
+
+        Args:
+            role (str): An AWS IAM role (either name or full ARN). The Amazon
+                SageMaker AutoML jobs and APIs that create Amazon SageMaker
+                endpoints use this role to access training data and model
+                artifacts.
+            sagemaker_session (sagemaker.session.Session): A SageMaker Session
+                object, used for SageMaker interactions.
+                If the best model will be used as part of ModelStep, then sagemaker_session
+                should be class:`~sagemaker.workflow.pipeline_context.PipelineSession`. Example::
+                     model = Model(sagemaker_session=PipelineSession())
+                     model_step = ModelStep(step_args=model.register())
+        """
+        inference_container = self.properties.BestCandidate.InferenceContainers[0]
+        inference_container_environment = inference_container.Environment
+        image = inference_container.Image
+        model_data = inference_container.ModelDataUrl
+        model = Model(
+            image_uri=image,
+            model_data=model_data,
+            env={
+                "MODEL_NAME": inference_container_environment["MODEL_NAME"],
+                "SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT": inference_container_environment[
+                    "SAGEMAKER_DEFAULT_INVOCATIONS_ACCEPT"
+                ],
+                "SAGEMAKER_SUBMIT_DIRECTORY": inference_container_environment[
+                    "SAGEMAKER_SUBMIT_DIRECTORY"
+                ],
+                "SAGEMAKER_INFERENCE_SUPPORTED": inference_container_environment[
+                    "SAGEMAKER_INFERENCE_SUPPORTED"
+                ],
+                "SAGEMAKER_INFERENCE_OUTPUT": inference_container_environment[
+                    "SAGEMAKER_INFERENCE_OUTPUT"
+                ],
+                "SAGEMAKER_PROGRAM": inference_container_environment["SAGEMAKER_PROGRAM"],
+            },
+            sagemaker_session=sagemaker_session,
+            role=role,
+        )
+
+        return model

src/sagemaker/workflow/automl_step.py

@@ -71,6 +71,7 @@ class StepTypeEnum(Enum, metaclass=DefaultEnumMeta):
     CLARIFY_CHECK = "ClarifyCheck"
     EMR = "EMR"
     FAIL = "Fail"
+    AUTOML = "AutoML"
 
 
 @attr.s

src/sagemaker/workflow/steps.py

@@ -0,0 +1,31 @@
+1080,4,setosa,versicolor,virginica
+5.9,3.0,4.2,1.5,1
+6.9,3.1,5.4,2.1,2
+5.1,3.3,1.7,0.5,0
+6.0,3.4,4.5,1.6,1
+5.5,2.5,4.0,1.3,1
+6.2,2.9,4.3,1.3,1
+5.5,4.2,1.4,0.2,0
+6.3,2.8,5.1,1.5,2
+5.6,3.0,4.1,1.3,1
+6.7,2.5,5.8,1.8,2
+7.1,3.0,5.9,2.1,2
+4.3,3.0,1.1,0.1,0
+5.6,2.8,4.9,2.0,2
+5.5,2.3,4.0,1.3,1
+6.0,2.2,4.0,1.0,1
+5.1,3.5,1.4,0.2,0
+5.7,2.6,3.5,1.0,1
+4.8,3.4,1.9,0.2,0
+5.1,3.4,1.5,0.2,0
+5.7,2.5,5.0,2.0,2
+5.4,3.4,1.7,0.2,0
+5.6,3.0,4.5,1.5,1
+6.3,2.9,5.6,1.8,2
+6.3,2.5,4.9,1.5,1
+5.8,2.7,3.9,1.2,1
+6.1,3.0,4.6,1.4,1
+5.2,4.1,1.5,0.1,0
+6.7,3.1,4.7,1.5,1
+6.7,3.3,5.7,2.5,2
+6.4,2.9,4.3,1.3,1