Add AutoMLV2 support

Author: repushko

doc/api/training/automlv2.rst

@@ -0,0 +1,7 @@
+AutoMLV2
+--------
+
+.. automodule:: sagemaker.automl.automlv2
+    :members:
+    :undoc-members:
+    :show-inheritance:

doc/api/training/index.rst

@@ -8,6 +8,7 @@ Training APIs
    algorithm
    analytics
    automl
+   automlv2
    debugger
    estimators
    tuner

src/sagemaker/__init__.py

@@ -61,6 +61,17 @@
 
 from sagemaker.automl.automl import AutoML, AutoMLJob, AutoMLInput  # noqa: F401
 from sagemaker.automl.candidate_estimator import CandidateEstimator, CandidateStep  # noqa: F401
+from sagemaker.automl.automlv2 import (  # noqa: F401
+    AutoMLV2,
+    AutoMLJobV2,
+    LocalAutoMLDataChannel,
+    AutoMLDataChannel,
+    AutoMLTimeSeriesForecastingConfig,
+    AutoMLImageClassificationConfig,
+    AutoMLTabularConfig,
+    AutoMLTextClassificationConfig,
+    AutoMLTextGenerationConfig,
+)
 
 from sagemaker.debugger import ProfilerConfig, Profiler  # noqa: F401
 

src/sagemaker/automl/automlv2.py

@@ -50,6 +50,7 @@
     AUTO_ML_OUTPUT_CONFIG_PATH,
     AUTO_ML_JOB_CONFIG_PATH,
     AUTO_ML_JOB,
+    AUTO_ML_JOB_V2,
     COMPILATION_JOB_ROLE_ARN_PATH,
     COMPILATION_JOB_OUTPUT_CONFIG_PATH,
     COMPILATION_JOB_VPC_CONFIG_PATH,

src/sagemaker/config/__init__.py

@@ -83,6 +83,7 @@
 ENDPOINT = "Endpoint"
 INFERENCE_COMPONENT = "InferenceComponent"
 AUTO_ML_JOB = "AutoMLJob"
+AUTO_ML_JOB_V2 = "AutoMLJobV2"
 COMPILATION_JOB = "CompilationJob"
 CUSTOM_PARAMETERS = "CustomParameters"
 PIPELINE = "Pipeline"

src/sagemaker/config/config_schema.py

@@ -2570,12 +2570,273 @@ def logs_for_auto_ml_job(  # noqa: C901 - suppress complexity warning for this m
             exceptions.UnexpectedStatusException: If waiting and auto ml job fails.
         """
 
-        description = _wait_until(lambda: self.describe_auto_ml_job(job_name), poll)
+        description = _wait_until(lambda: self.describe_auto_ml_job_v2(job_name), poll)
 
-        instance_count, stream_names, positions, client, log_group, dot, color_wrap = _logs_init(
-            self.boto_session, description, job="AutoML"
+        (
+            instance_count,
+            stream_names,
+            positions,
+            client,
+            log_group,
+            dot,
+            color_wrap,
+        ) = _logs_init(self.boto_session, description, job="AutoML")
+
+        state = _get_initial_job_state(description, "AutoMLJobStatus", wait)
+
+        # The loop below implements a state machine that alternates between checking the job status
+        # and reading whatever is available in the logs at this point. Note, that if we were
+        # called with wait == False, we never check the job status.
+        #
+        # If wait == TRUE and job is not completed, the initial state is TAILING
+        # If wait == FALSE, the initial state is COMPLETE (doesn't matter if the job really is
+        # complete).
+        #
+        # The state table:
+        #
+        # STATE               ACTIONS                        CONDITION             NEW STATE
+        # ----------------    ----------------               -----------------     ----------------
+        # TAILING             Read logs, Pause, Get status   Job complete          JOB_COMPLETE
+        #                                                    Else                  TAILING
+        # JOB_COMPLETE        Read logs, Pause               Any                   COMPLETE
+        # COMPLETE            Read logs, Exit                                      N/A
+        #
+        # Notes:
+        # - The JOB_COMPLETE state forces us to do an extra pause and read any items that got to
+        #   Cloudwatch after the job was marked complete.
+        last_describe_job_call = time.time()
+        while True:
+            _flush_log_streams(
+                stream_names,
+                instance_count,
+                client,
+                log_group,
+                job_name,
+                positions,
+                dot,
+                color_wrap,
+            )
+            if state == LogState.COMPLETE:
+                break
+
+            time.sleep(poll)
+
+            if state == LogState.JOB_COMPLETE:
+                state = LogState.COMPLETE
+            elif time.time() - last_describe_job_call >= 30:
+                description = self.sagemaker_client.describe_auto_ml_job_v2(AutoMLJobName=job_name)
+                last_describe_job_call = time.time()
+
+                status = description["AutoMLJobStatus"]
+
+                if status in ("Completed", "Failed", "Stopped"):
+                    print()
+                    state = LogState.JOB_COMPLETE
+
+        if wait:
+            _check_job_status(job_name, description, "AutoMLJobStatus")
+            if dot:
+                print()
+
+    def create_auto_ml_v2(
+        self,
+        input_config,
+        job_name,
+        problem_config,
+        output_config,
+        job_objective=None,
+        model_deploy_config=None,
+        data_split_config=None,
+        role=None,
+        security_config=None,
+        tags=None,
+    ):
+        """Create an Amazon SageMaker AutoMLV2 job.
+
+        Args:
+            input_config (list[dict]): A list of AutoMLDataChannel objects.
+                Each channel contains "DataSource" and other optional fields.
+            job_name (str): A string that can be used to identify an AutoMLJob. Each AutoMLJob
+                should have a unique job name.
+            problem_config (object): A collection of settings specific
+                to the problem type used to configure an AutoML job V2.
+                There must be one and only one config of the following type.
+                Supported problem types are:
+
+                - Image Classification (sagemaker.automl.automlv2.ImageClassificationJobConfig),
+                - Tabular (sagemaker.automl.automlv2.TabularJobConfig),
+                - Text Classification (sagemaker.automl.automlv2.TextClassificationJobConfig),
+                - Text Generation (TextGenerationJobConfig),
+                - Time Series Forecasting (
+                    sagemaker.automl.automlv2.TimeSeriesForecastingJobConfig).
+
+            output_config (dict): The S3 URI where you want to store the training results and
+                optional KMS key ID.
+            job_objective (dict): AutoMLJob objective, contains "AutoMLJobObjectiveType" (optional),
+                "MetricName" and "Value".
+            model_deploy_config (dict): Specifies how to generate the endpoint name
+                for an automatic one-click Autopilot model deployment.
+                Contains "AutoGenerateEndpointName" and "EndpointName"
+            data_split_config (dict): This structure specifies how to split the data
+                into train and validation datasets.
+            role (str): The Amazon Resource Name (ARN) of an IAM role that
+                Amazon SageMaker can assume to perform tasks on your behalf.
+            security_config (dict): The security configuration for traffic encryption
+                or Amazon VPC settings.
+            tags (Optional[Tags]): A list of dictionaries containing key-value
+                pairs.
+        """
+
+        role = resolve_value_from_config(role, AUTO_ML_ROLE_ARN_PATH, sagemaker_session=self)
+        inferred_output_config = update_nested_dictionary_with_values_from_config(
+            output_config, AUTO_ML_OUTPUT_CONFIG_PATH, sagemaker_session=self
         )
 
+        auto_ml_job_v2_request = self._get_auto_ml_request_v2(
+            input_config=input_config,
+            job_name=job_name,
+            problem_config=problem_config,
+            output_config=inferred_output_config,
+            role=role,
+            job_objective=job_objective,
+            model_deploy_config=model_deploy_config,
+            data_split_config=data_split_config,
+            security_config=security_config,
+            tags=format_tags(tags),
+        )
+
+        def submit(request):
+            logger.info("Creating auto-ml-v2-job with name: %s", job_name)
+            logger.debug("auto ml v2 request: %s", json.dumps(request), indent=4)
+            print(json.dumps(request))
+            self.sagemaker_client.create_auto_ml_job_v2(**request)
+
+        self._intercept_create_request(
+            auto_ml_job_v2_request, submit, self.create_auto_ml_v2.__name__
+        )
+
+    def _get_auto_ml_request_v2(
+        self,
+        input_config,
+        output_config,
+        job_name,
+        problem_config,
+        role,
+        job_objective=None,
+        model_deploy_config=None,
+        data_split_config=None,
+        security_config=None,
+        tags=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" and "SampleWeightAttributeName" are
+                optional fields.
+            output_config (dict): The S3 URI where you want to store the training results and
+                optional KMS key ID.
+            job_name (str): A string that can be used to identify an AutoMLJob. Each AutoMLJob
+                should have a unique job name.
+            problem_config (object): A collection of settings specific
+                to the problem type used to configure an AutoML job V2.
+                There must be one and only one config of the following type.
+                Supported problem types are:
+
+                - Image Classification (sagemaker.automl.automlv2.ImageClassificationJobConfig),
+                - Tabular (sagemaker.automl.automlv2.TabularJobConfig),
+                - Text Classification (sagemaker.automl.automlv2.TextClassificationJobConfig),
+                - Text Generation (TextGenerationJobConfig),
+                - Time Series Forecasting (
+                    sagemaker.automl.automlv2.TimeSeriesForecastingJobConfig).
+
+            role (str): The Amazon Resource Name (ARN) of an IAM role that
+                Amazon SageMaker can assume to perform tasks on your behalf.
+            job_objective (dict): AutoMLJob objective, contains "AutoMLJobObjectiveType" (optional),
+                "MetricName" and "Value".
+            model_deploy_config (dict): Specifies how to generate the endpoint name
+                for an automatic one-click Autopilot model deployment.
+                Contains "AutoGenerateEndpointName" and "EndpointName"
+            data_split_config (dict): This structure specifies how to split the data
+                into train and validation datasets.
+            security_config (dict): The security configuration for traffic encryption
+                or Amazon VPC settings.
+            tags (Optional[Tags]): A list of dictionaries containing key-value
+                pairs.
+
+        Returns:
+            Dict: a automl v2 request dict
+        """
+        auto_ml_job_v2_request = {
+            "AutoMLJobName": job_name,
+            "AutoMLJobInputDataConfig": input_config,
+            "OutputDataConfig": output_config,
+            "AutoMLProblemTypeConfig": problem_config,
+            "RoleArn": role,
+        }
+        if job_objective is not None:
+            auto_ml_job_v2_request["AutoMLJobObjective"] = job_objective
+        if model_deploy_config is not None:
+            auto_ml_job_v2_request["ModelDeployConfig"] = model_deploy_config
+        if data_split_config is not None:
+            auto_ml_job_v2_request["DataSplitConfig"] = data_split_config
+        if security_config is not None:
+            auto_ml_job_v2_request["SecurityConfig"] = security_config
+
+        tags = _append_project_tags(format_tags(tags))
+        tags = self._append_sagemaker_config_tags(
+            tags, "{}.{}.{}".format(SAGEMAKER, AUTO_ML_JOB, TAGS)
+        )
+        if tags is not None:
+            auto_ml_job_v2_request["Tags"] = tags
+
+        return auto_ml_job_v2_request
+
+    # Done
+    def describe_auto_ml_job_v2(self, job_name):
+        """Calls the DescribeAutoMLJobV2 API for the given job name and returns the response.
+
+        Args:
+            job_name (str): The name of the AutoML job to describe.
+
+        Returns:
+            dict: A dictionary response with the AutoMLV2 Job description.
+        """
+        return self.sagemaker_client.describe_auto_ml_job_v2(AutoMLJobName=job_name)
+
+    def logs_for_auto_ml_job_v2(  # noqa: C901 - suppress complexity warning for this method
+        self, job_name, wait=False, poll=10
+    ):
+        """Display logs for a given AutoML V2 job, optionally tailing them until job is complete.
+
+        If the output is a tty or a Jupyter cell, it will be color-coded
+        based on which instance the log entry is from.
+
+        Args:
+            job_name (str): Name of the Auto ML V2 job to display the logs for.
+            wait (bool): Whether to keep looking for new log entries until the job completes
+                (default: False).
+            poll (int): The interval in seconds between polling for new log entries and job
+                completion (default: 5).
+
+        Raises:
+            exceptions.CapacityError: If waiting and auto ml job fails with CapacityError.
+            exceptions.UnexpectedStatusException: If waiting and auto ml job fails.
+        """
+
+        # Why here is the method from the session and below is the method from the sagemaker client?
+        description = _wait_until(lambda: self.describe_auto_ml_job_v2(job_name), poll)
+
+        (
+            instance_count,
+            stream_names,
+            positions,
+            client,
+            log_group,
+            dot,
+            color_wrap,
+        ) = _logs_init(self.boto_session, description, job="AutoML")
+
         state = _get_initial_job_state(description, "AutoMLJobStatus", wait)
 
         # The loop below implements a state machine that alternates between checking the job status
@@ -2618,7 +2879,7 @@ def logs_for_auto_ml_job(  # noqa: C901 - suppress complexity warning for this m
             if state == LogState.JOB_COMPLETE:
                 state = LogState.COMPLETE
             elif time.time() - last_describe_job_call >= 30:
-                description = self.sagemaker_client.describe_auto_ml_job(AutoMLJobName=job_name)
+                description = self.sagemaker_client.describe_auto_ml_job_v2(AutoMLJobName=job_name)
                 last_describe_job_call = time.time()
 
                 status = description["AutoMLJobStatus"]