change: AutoML improvements (#281)

* change: AutoML improvements

* AutoML improvements

* address comments

Author: chuyang-deng

buildspec.yml

@@ -19,7 +19,7 @@ phases:
       - start_time=`date +%s`
       - |
         if has-matching-changes "tests/" "src/*.py" "setup.py" "setup.cfg" "buildspec.yml"; then
-          tox -e py36 -- tests/integ -m "not local_mode" -n 48  --reruns 3 --reruns-delay 5 --durations 50 --boto-config '{"region_name": "us-west-2"}'
+          tox -e py36 -- tests/integ -m "not local_mode" -n 48  --reruns 3 --reruns-delay 5 --durations 50 --boto-config '{"region_name": "us-east-2"}'
         fi
       - ./ci-scripts/displaytime.sh 'py36 tests/integ' $start_time
 

src/sagemaker/automl/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2017-2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright 2019 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

src/sagemaker/automl/automl.py

@@ -74,8 +74,9 @@ def fit(self, inputs=None, wait=True, logs=True, job_name=None):
         """Create an AutoML Job with the input dataset.
 
         Args:
-            inputs (str or list[str]): Local path or S3 Uri where the training data is stored. If a
-                local path is provided, the dataset will be uploaded to an S3 location.
+            inputs (str or list[str] or AutoMLInput): Local path or S3 Uri where the training data
+                is stored. Or an AutoMLInput object. If a local path is provided, the dataset will
+                be uploaded to an S3 location.
             wait (bool): Whether the call should wait until the job completes (default: True).
             logs (bool): Whether to show the logs produced by the job.
                 Only meaningful when wait is True (default: True).
@@ -95,7 +96,7 @@ def fit(self, inputs=None, wait=True, logs=True, job_name=None):
                 inputs = self.sagemaker_session.upload_data(inputs, key_prefix="auto-ml-input-data")
         self._prepare_for_auto_ml_job(job_name=job_name)
 
-        self.latest_auto_ml_job = _AutoMLJob.start_new(self, inputs)  # pylint: disable=W0201
+        self.latest_auto_ml_job = AutoMLJob.start_new(self, inputs)  # pylint: disable=W0201
         if wait:
             self.latest_auto_ml_job.wait(logs=logs)
 
@@ -385,9 +386,48 @@ def _prepare_for_auto_ml_job(self, job_name=None):
             self.output_path = "s3://{}/".format(self.sagemaker_session.default_bucket())
 
 
-class _AutoMLJob(_Job):
+class AutoMLInput(object):
+    """Accepts parameters that specify an S3 input for an auto ml job and provides
+    a method to turn those parameters into a dictionary."""
+
+    def __init__(self, inputs, target_attribute_name, compression=None):
+        """Convert an S3 Uri or a list of S3 Uri to an AutoMLInput object.
+
+        :param inputs (str, list[str]): a string or a list of string that points to (a)
+            S3 location(s) where input data is stored.
+        :param target_attribute_name (str): the target attribute name for regression
+            or classification.
+        :param compression (str): if training data is compressed, the compression type.
+            The default value is None.
+        """
+        self.inputs = inputs
+        self.target_attribute_name = target_attribute_name
+        self.compression = compression
+
+    def to_request_dict(self):
+        """Generates a request dictionary using the parameters provided to the class."""
+        # Create the request dictionary.
+        auto_ml_input = []
+        if isinstance(self.inputs, string_types):
+            self.inputs = [self.inputs]
+        for entry in self.inputs:
+            input_entry = {
+                "DataSource": {"S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": entry}},
+                "TargetAttributeName": self.target_attribute_name,
+            }
+            if self.compression is not None:
+                input_entry["CompressionType"] = self.compression
+            auto_ml_input.append(input_entry)
+        return auto_ml_input
+
+
+class AutoMLJob(_Job):
     """A class for interacting with CreateAutoMLJob API."""
 
+    def __init__(self, sagemaker_session, job_name, inputs):
+        self.inputs = inputs
+        super(AutoMLJob, self).__init__(sagemaker_session=sagemaker_session, job_name=job_name)
+
     @classmethod
     def start_new(cls, auto_ml, inputs):
         """Create a new Amazon SageMaker AutoML job from auto_ml.
@@ -399,7 +439,7 @@ def start_new(cls, auto_ml, inputs):
                 :meth:`~sagemaker.automl.AutoML.fit`.
 
         Returns:
-            sagemaker.automl._AutoMLJob: Constructed object that captures
+            sagemaker.automl.AutoMLJob: Constructed object that captures
             all information about the started AutoML job.
         """
         config = cls._load_config(inputs, auto_ml)
@@ -410,7 +450,7 @@ def start_new(cls, auto_ml, inputs):
         auto_ml_args["tags"] = auto_ml.tags
 
         auto_ml.sagemaker_session.auto_ml(**auto_ml_args)
-        return cls(auto_ml.sagemaker_session, auto_ml._current_job_name)
+        return cls(auto_ml.sagemaker_session, auto_ml._current_job_name, inputs)
 
     @classmethod
     def _load_config(cls, inputs, auto_ml, expand_role=True, validate_uri=True):
@@ -432,9 +472,12 @@ def _load_config(cls, inputs, auto_ml, expand_role=True, validate_uri=True):
         # InputDataConfig
         # OutputConfig
 
-        input_config = cls._format_inputs_to_input_config(
-            inputs, validate_uri, auto_ml.compression_type, auto_ml.target_attribute_name
-        )
+        if isinstance(inputs, AutoMLInput):
+            input_config = inputs.to_request_dict()
+        else:
+            input_config = cls._format_inputs_to_input_config(
+                inputs, validate_uri, auto_ml.compression_type, auto_ml.target_attribute_name
+            )
         output_config = _Job._prepare_output_config(auto_ml.output_path, auto_ml.output_kms_key)
 
         role = auto_ml.sagemaker_session.expand_role(auto_ml.role) if expand_role else auto_ml.role
@@ -486,7 +529,9 @@ def _format_inputs_to_input_config(
             return None
 
         channels = []
-        if isinstance(inputs, string_types):
+        if isinstance(inputs, AutoMLInput):
+            channels.append(inputs.to_request_dict())
+        elif isinstance(inputs, string_types):
             channel = _Job._format_string_uri_input(
                 inputs,
                 validate_uri,
@@ -540,6 +585,10 @@ def _prepare_auto_ml_stop_condition(
 
         return stopping_condition
 
+    def describe(self):
+        """Prints out a response from the DescribeAutoMLJob API call."""
+        return self.sagemaker_session.describe_auto_ml_job(self.job_name)
+
     def wait(self, logs=True):
         """Wait for the AutoML job to finish.
         Args:

src/sagemaker/automl/candidate_estimator.py

@@ -56,24 +56,14 @@ def get_steps(self):
                 )
 
                 inputs = training_job["InputDataConfig"]
-                candidate_step = {
-                    "name": step_name,
-                    "inputs": inputs,
-                    "type": step_type,
-                    "desc": training_job,
-                }
+                candidate_step = CandidateStep(step_name, inputs, step_type, training_job)
                 candidate_steps.append(candidate_step)
             elif step_type == "TransformJob":
                 transform_job = self.sagemaker_session.sagemaker_client.describe_transform_job(
                     TransformJobName=step_name
                 )
                 inputs = transform_job["TransformInput"]
-                candidate_step = {
-                    "name": step_name,
-                    "inputs": inputs,
-                    "type": step_type,
-                    "desc": transform_job,
-                }
+                candidate_step = CandidateStep(step_name, inputs, step_type, transform_job)
                 candidate_steps.append(candidate_step)
         return candidate_steps
 
@@ -313,3 +303,33 @@ def _process_steps(self, steps):
             step_type = step["CandidateStepType"].split("::")[2]
             processed_steps.append({"name": step_name, "type": step_type})
         return processed_steps
+
+
+class CandidateStep(object):
+    """A class that maintains an AutoML Candidate step's name, inputs, type, and description."""
+
+    def __init__(self, name, inputs, step_type, description):
+        self._name = name
+        self._inputs = inputs
+        self._type = step_type
+        self._description = description
+
+    @property
+    def name(self):
+        """Name of the candidate step -> (str)"""
+        return self._name
+
+    @property
+    def inputs(self):
+        """Inputs of the candidate step -> (dict)"""
+        return self._inputs
+
+    @property
+    def type(self):
+        """Type of the candidate step, Training or Transform -> (str)"""
+        return self._type
+
+    @property
+    def description(self):
+        """Description of candidate step job -> (dict)"""
+        return self._description

src/sagemaker/job.py

@@ -195,13 +195,14 @@ def _format_string_uri_input(
             target_attribute_name:
         """
         if isinstance(uri_input, str) and validate_uri and uri_input.startswith("s3://"):
-            return s3_input(
+            s3_input_result = s3_input(
                 uri_input,
                 content_type=content_type,
                 input_mode=input_mode,
                 compression=compression,
                 target_attribute_name=target_attribute_name,
             )
+            return s3_input_result
         if isinstance(uri_input, str) and validate_uri and uri_input.startswith("file://"):
             return file_input(uri_input)
         if isinstance(uri_input, str) and validate_uri:
@@ -210,13 +211,14 @@ def _format_string_uri_input(
                 '"file://"'.format(uri_input)
             )
         if isinstance(uri_input, str):
-            return s3_input(
+            s3_input_result = s3_input(
                 uri_input,
                 content_type=content_type,
                 input_mode=input_mode,
                 compression=compression,
                 target_attribute_name=target_attribute_name,
             )
+            return s3_input_result
         if isinstance(uri_input, (s3_input, file_input, FileSystemInput)):
             return uri_input
 

tests/integ/test_auto_ml.py

@@ -17,7 +17,7 @@
 
 import pytest
 
-from sagemaker.automl.automl import AutoML
+from sagemaker.automl.automl import AutoML, AutoMLInput
 from sagemaker.automl.candidate_estimator import CandidateEstimator
 from sagemaker.exceptions import UnexpectedStatusException
 from sagemaker.utils import unique_name_from_base
@@ -47,7 +47,7 @@
         "DataSource": {
             "S3DataSource": {
                 "S3DataType": "S3Prefix",
-                "S3Uri": "s3://sagemaker-us-west-2-{}/{}/input/iris_training.csv".format(
+                "S3Uri": "s3://sagemaker-us-east-2-{}/{}/input/iris_training.csv".format(
                     DEV_ACCOUNT, PREFIX
                 ),
             }
@@ -60,7 +60,7 @@
     "SecurityConfig": {"EnableInterContainerTrafficEncryption": False},
 }
 EXPECTED_DEFAULT_OUTPUT_CONFIG = {
-    "S3OutputPath": "s3://sagemaker-us-west-2-{}/".format(DEV_ACCOUNT)
+    "S3OutputPath": "s3://sagemaker-us-east-2-{}/".format(DEV_ACCOUNT)
 }
 
 
@@ -90,8 +90,21 @@ def test_auto_ml_fit_local_input(sagemaker_session):
         auto_ml.fit(inputs)
 
 
+def test_auto_ml_input_object_fit(sagemaker_session):
+    auto_ml = AutoML(
+        role=ROLE,
+        target_attribute_name=TARGET_ATTRIBUTE_NAME,
+        sagemaker_session=sagemaker_session,
+        max_candidates=1,
+    )
+    s3_input = sagemaker_session.upload_data(path=TRAINING_DATA, key_prefix=PREFIX + "/input")
+    inputs = AutoMLInput(inputs=s3_input, target_attribute_name=TARGET_ATTRIBUTE_NAME)
+    with timeout(minutes=AUTO_ML_DEFAULT_TIMEMOUT_MINUTES):
+        auto_ml.fit(inputs)
+
+
 def test_auto_ml_fit_optional_args(sagemaker_session):
-    output_path = "s3://sagemaker-us-west-2-{}/{}".format(DEV_ACCOUNT, "specified_ouput_path")
+    output_path = "s3://sagemaker-us-east-2-{}/{}".format(DEV_ACCOUNT, "specified_ouput_path")
     problem_type = "MulticlassClassification"
     job_objective = {"MetricName": "Accuracy"}
     auto_ml = AutoML(

tests/unit/sagemaker/automl/test_auto_ml.py

@@ -15,7 +15,7 @@
 import pytest
 from mock import Mock, patch
 
-from sagemaker.automl.automl import AutoML, _AutoMLJob
+from sagemaker.automl.automl import AutoML, AutoMLJob, AutoMLInput
 from sagemaker.automl.candidate_estimator import CandidateEstimator
 
 MODEL_DATA = "s3://bucket/model.tar.gz"
@@ -189,7 +189,7 @@ def test_auto_ml_default_channel_name(sagemaker_session):
         role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
     )
     inputs = DEFAULT_S3_INPUT_DATA
-    _AutoMLJob.start_new(auto_ml, inputs)
+    AutoMLJob.start_new(auto_ml, inputs)
     sagemaker_session.auto_ml.assert_called_once()
     _, args = sagemaker_session.auto_ml.call_args
     assert args["input_config"] == [
@@ -210,7 +210,7 @@ def test_auto_ml_invalid_input_data_format(sagemaker_session):
 
     expected_error_msg = "Cannot format input {}. Expecting one of str or list of strings."
     with pytest.raises(ValueError, message=expected_error_msg.format(inputs)):
-        _AutoMLJob.start_new(auto_ml, inputs)
+        AutoMLJob.start_new(auto_ml, inputs)
     sagemaker_session.auto_ml.assert_not_called()
 
 
@@ -330,6 +330,26 @@ def test_auto_ml_local_input(sagemaker_session):
     assert args["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] == DEFAULT_S3_INPUT_DATA
 
 
+def test_auto_ml_input(sagemaker_session):
+    inputs = AutoMLInput(
+        inputs=DEFAULT_S3_INPUT_DATA, target_attribute_name="target", compression="Gzip"
+    )
+    auto_ml = AutoML(
+        role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
+    )
+    auto_ml.fit(inputs)
+    _, args = sagemaker_session.auto_ml.call_args
+    assert args["input_config"] == [
+        {
+            "CompressionType": "Gzip",
+            "DataSource": {
+                "S3DataSource": {"S3DataType": "S3Prefix", "S3Uri": DEFAULT_S3_INPUT_DATA}
+            },
+            "TargetAttributeName": TARGET_ATTRIBUTE_NAME,
+        }
+    ]
+
+
 def test_describe_auto_ml_job(sagemaker_session):
     auto_ml = AutoML(
         role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session