fix: correct AutoML imports and expose current_job_name (#300)

Author: chuyang-deng

src/sagemaker/__init__.py

@@ -57,4 +57,7 @@
 from sagemaker.session import s3_input  # noqa: F401
 from sagemaker.session import get_execution_role  # noqa: F401
 
+from sagemaker.automl.automl import AutoML, AutoMLJob, AutoMLInput  # noqa: F401
+from sagemaker.automl.candidate_estimator import CandidateEstimator, CandidateStep  # noqa: F401
+
 __version__ = pkg_resources.require("sagemaker")[0].version

src/sagemaker/automl/automl.py

@@ -63,7 +63,7 @@ def __init__(
         self.generate_candidate_definitions_only = generate_candidate_definitions_only
         self.tags = tags
 
-        self._current_job_name = None
+        self.current_job_name = None
         self._auto_ml_job_desc = None
         self._best_candidate = None
         self.sagemaker_session = sagemaker_session or Session()
@@ -111,7 +111,7 @@ def describe_auto_ml_job(self, job_name=None):
             dict: A dictionary response with the AutoML Job description.
         """
         if job_name is None:
-            job_name = self._current_job_name
+            job_name = self.current_job_name
         self._auto_ml_job_desc = self.sagemaker_session.describe_auto_ml_job(job_name)
         return self._auto_ml_job_desc
 
@@ -128,7 +128,7 @@ def best_candidate(self, job_name=None):
             return self._best_candidate
 
         if job_name is None:
-            job_name = self._current_job_name
+            job_name = self.current_job_name
         if self._auto_ml_job_desc is None:
             self._auto_ml_job_desc = self.sagemaker_session.describe_auto_ml_job(job_name)
         elif self._auto_ml_job_desc["AutoMLJobName"] != job_name:
@@ -168,7 +168,7 @@ def list_candidates(
             list: A list of dictionaries with candidates information
         """
         if job_name is None:
-            job_name = self._current_job_name
+            job_name = self.current_job_name
 
         list_candidates_args = {"job_name": job_name}
 
@@ -249,6 +249,7 @@ def deploy(
             candidate = CandidateEstimator(candidate, sagemaker_session=sagemaker_session)
 
         inference_containers = candidate.containers
+        endpoint_name = endpoint_name or self.current_job_name
 
         return self._deploy_inference_pipeline(
             inference_containers,
@@ -373,14 +374,14 @@ def _prepare_for_auto_ml_job(self, job_name=None):
                 created from base_job_name or "sagemaker-auto-ml".
         """
         if job_name is not None:
-            self._current_job_name = job_name
+            self.current_job_name = job_name
         else:
             if self.base_job_name:
                 base_name = self.base_job_name
             else:
                 base_name = "sagemaker-auto-ml"
             # CreateAutoMLJob API validates that member length less than or equal to 32
-            self._current_job_name = name_from_base(base_name, max_length=32)
+            self.current_job_name = name_from_base(base_name, max_length=32)
 
         if self.output_path is None:
             self.output_path = "s3://{}/".format(self.sagemaker_session.default_bucket())
@@ -426,6 +427,7 @@ class AutoMLJob(_Job):
 
     def __init__(self, sagemaker_session, job_name, inputs):
         self.inputs = inputs
+        self.job_name = job_name
         super(AutoMLJob, self).__init__(sagemaker_session=sagemaker_session, job_name=job_name)
 
     @classmethod
@@ -444,13 +446,13 @@ def start_new(cls, auto_ml, inputs):
         """
         config = cls._load_config(inputs, auto_ml)
         auto_ml_args = config.copy()
-        auto_ml_args["job_name"] = auto_ml._current_job_name
+        auto_ml_args["job_name"] = auto_ml.current_job_name
         auto_ml_args["problem_type"] = auto_ml.problem_type
         auto_ml_args["job_objective"] = auto_ml.job_objective
         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, inputs)
+        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):

src/sagemaker/session.py

@@ -1326,12 +1326,12 @@ def wait_for_auto_ml_job(self, job, poll=5):
         """Wait for an Amazon SageMaker AutoML job to complete.
 
         Args:
-            job (str): Name of the transform job to wait for.
+            job (str): Name of the auto ml job to wait for.
             poll (int): Polling interval in seconds (default: 5).
         Returns:
-            (dict): Return value from the ``DescribeTransformJob`` API.
+            (dict): Return value from the ``DescribeAutoMLJob`` API.
         Raises:
-            exceptions.UnexpectedStatusException: If the transform job fails.
+            exceptions.UnexpectedStatusException: If the auto ml job fails.
         """
         desc = _wait_until(lambda: _auto_ml_job_status(self.sagemaker_client, job), poll)
         self._check_job_status(job, desc, "AutoMLJobStatus")

tests/integ/test_auto_ml.py

@@ -16,16 +16,14 @@
 import time
 
 import pytest
+from sagemaker import AutoML, CandidateEstimator, AutoMLInput
 
-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
 from tests.integ import DATA_DIR, AUTO_ML_DEFAULT_TIMEMOUT_MINUTES
 from tests.integ.timeout import timeout
 
 DEV_ACCOUNT = 142577830533
-# ROLE = "arn:aws:iam::142577830533:role/SageMakerRole"
 ROLE = "SageMakerRole"
 PREFIX = "sagemaker/beta-automl-xgboost"
 HOSTING_INSTANCE_TYPE = "ml.c4.xlarge"

tests/unit/sagemaker/automl/test_auto_ml.py

@@ -14,9 +14,7 @@
 
 import pytest
 from mock import Mock, patch
-
-from sagemaker.automl.automl import AutoML, AutoMLJob, AutoMLInput
-from sagemaker.automl.candidate_estimator import CandidateEstimator
+from sagemaker import AutoML, AutoMLJob, AutoMLInput, CandidateEstimator
 
 MODEL_DATA = "s3://bucket/model.tar.gz"
 MODEL_IMAGE = "mi"
@@ -363,10 +361,10 @@ def test_list_candidates_default(sagemaker_session):
     auto_ml = AutoML(
         role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
     )
-    auto_ml._current_job_name = "current_job_name"
+    auto_ml.current_job_name = "current_job_name"
     auto_ml.list_candidates()
     sagemaker_session.list_candidates.assert_called_once()
-    sagemaker_session.list_candidates.assert_called_with(job_name=auto_ml._current_job_name)
+    sagemaker_session.list_candidates.assert_called_with(job_name=auto_ml.current_job_name)
 
 
 def test_list_candidates_with_optional_args(sagemaker_session):
@@ -409,7 +407,7 @@ def test_best_candidate_default_job_name(sagemaker_session):
     auto_ml = AutoML(
         role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
     )
-    auto_ml._current_job_name = JOB_NAME
+    auto_ml.current_job_name = JOB_NAME
     auto_ml._auto_ml_job_desc = AUTO_ML_DESC
     best_candidate = auto_ml.best_candidate()
     sagemaker_session.describe_auto_ml_job.assert_not_called()
@@ -420,7 +418,7 @@ def test_best_candidate_job_no_desc(sagemaker_session):
     auto_ml = AutoML(
         role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
     )
-    auto_ml._current_job_name = JOB_NAME
+    auto_ml.current_job_name = JOB_NAME
     best_candidate = auto_ml.best_candidate()
     sagemaker_session.describe_auto_ml_job.assert_called_once()
     sagemaker_session.describe_auto_ml_job.assert_called_with(JOB_NAME)
@@ -441,7 +439,7 @@ def test_best_candidate_job_name_not_match(sagemaker_session):
     auto_ml = AutoML(
         role=ROLE, target_attribute_name=TARGET_ATTRIBUTE_NAME, sagemaker_session=sagemaker_session
     )
-    auto_ml._current_job_name = JOB_NAME
+    auto_ml.current_job_name = JOB_NAME
     auto_ml._auto_ml_job_desc = AUTO_ML_DESC
     best_candidate = auto_ml.best_candidate(job_name=JOB_NAME_2)
     sagemaker_session.describe_auto_ml_job.assert_called_once()