change: infer base name from job name in estimator.attach()

src/sagemaker/chainer/estimator.py

@@ -274,12 +274,10 @@ class constructor
             init_params["image_name"] = image_name
             return init_params
 
-        training_job_name = init_params["base_job_name"]
-
         if framework != cls.__framework_name__:
             raise ValueError(
                 "Training job: {} didn't use image for requested framework".format(
-                    training_job_name
+                    job_details["TrainingJobName"]
                 )
             )
         return init_params

src/sagemaker/estimator.py

@@ -56,7 +56,7 @@
 from sagemaker.session import Session
 from sagemaker.session import s3_input
 from sagemaker.transformer import Transformer
-from sagemaker.utils import base_name_from_image, name_from_base, get_config_value
+from sagemaker.utils import base_from_name, base_name_from_image, name_from_base, get_config_value
 from sagemaker import vpc_utils
 
 
@@ -616,7 +616,7 @@ def attach(cls, training_job_name, sagemaker_session=None, model_channel_name="m
 
         estimator = cls(sagemaker_session=sagemaker_session, **init_params)
         estimator.latest_training_job = _TrainingJob(
-            sagemaker_session=sagemaker_session, job_name=init_params["base_job_name"]
+            sagemaker_session=sagemaker_session, job_name=training_job_name
         )
         estimator._current_job_name = estimator.latest_training_job.name
         estimator.latest_training_job.wait()
@@ -776,7 +776,7 @@ class constructor
         init_params["train_volume_size"] = job_details["ResourceConfig"]["VolumeSizeInGB"]
         init_params["train_max_run"] = job_details["StoppingCondition"]["MaxRuntimeInSeconds"]
         init_params["input_mode"] = job_details["AlgorithmSpecification"]["TrainingInputMode"]
-        init_params["base_job_name"] = job_details["TrainingJobName"]
+        init_params["base_job_name"] = base_from_name(job_details["TrainingJobName"])
         init_params["output_path"] = job_details["OutputDataConfig"]["S3OutputPath"]
         init_params["output_kms_key"] = job_details["OutputDataConfig"]["KmsKeyId"]
         if "EnableNetworkIsolation" in job_details:

src/sagemaker/mxnet/estimator.py

@@ -277,12 +277,10 @@ class constructor
             init_params["image_name"] = image_name
             return init_params
 
-        training_job_name = init_params["base_job_name"]
-
         if framework != cls.__framework_name__:
             raise ValueError(
                 "Training job: {} didn't use image for requested framework".format(
-                    training_job_name
+                    job_details["TrainingJobName"]
                 )
             )
 

src/sagemaker/pytorch/estimator.py

@@ -223,12 +223,10 @@ class constructor
             init_params["image_name"] = image_name
             return init_params
 
-        training_job_name = init_params["base_job_name"]
-
         if framework != cls.__framework_name__:
             raise ValueError(
                 "Training job: {} didn't use image for requested framework".format(
-                    training_job_name
+                    job_details["TrainingJobName"]
                 )
             )
 

src/sagemaker/rl/estimator.py

@@ -314,10 +314,9 @@ class constructor
         toolkit, toolkit_version = cls._toolkit_and_version_from_tag(tag)
 
         if not cls._is_combination_supported(toolkit, toolkit_version, framework):
-            training_job_name = init_params["base_job_name"]
             raise ValueError(
                 "Training job: {} didn't use image for requested framework".format(
-                    training_job_name
+                    job_details["TrainingJobName"]
                 )
             )
 

src/sagemaker/sklearn/estimator.py

@@ -246,12 +246,10 @@ class constructor
             init_params["image_name"] = image_name
             return init_params
 
-        training_job_name = init_params["base_job_name"]
-
         if framework and framework != cls.__framework_name__:
             raise ValueError(
                 "Training job: {} didn't use image for requested framework".format(
-                    training_job_name
+                    job_details["TrainingJobName"]
                 )
             )
 

src/sagemaker/tensorflow/estimator.py

@@ -221,11 +221,10 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
         if not script_mode:
             init_params["image_name"] = image_name
 
-        training_job_name = init_params["base_job_name"]
         if framework != cls.__framework_name__:
             raise ValueError(
                 "Training job: {} didn't use image for requested framework".format(
-                    training_job_name
+                    job_details["TrainingJobName"]
                 )
             )
 

src/sagemaker/utils.py

@@ -104,6 +104,22 @@ def base_name_from_image(image):
     return algo_name
 
 
+def base_from_name(name):
+    """Extract the base name of the resource name (for use with future resource name generation).
+
+    This function looks for timestamps that match the ones produced by
+    :func:`~sagemaker.utils.name_from_base`.
+
+    Args:
+        name (str): The resource name.
+
+    Returns:
+        str: The base name, as extracted from the resource name.
+    """
+    m = re.match(r"^(.+)-(\d{4}-\d{2}-\d{2}-\d{2}-\d{2}-\d{2}-\d{3}|\d{6}-\d{4})", name)
+    return m.group(1) if m else name
+
+
 def sagemaker_timestamp():
     """Return a timestamp with millisecond precision."""
     moment = time.time()

src/sagemaker/xgboost/estimator.py

@@ -238,7 +238,7 @@ def attach(cls, training_job_name, sagemaker_session=None, model_channel_name="m
 
         estimator = cls(sagemaker_session=sagemaker_session, **init_params)
         estimator.latest_training_job = _TrainingJob(
-            sagemaker_session=sagemaker_session, job_name=init_params["base_job_name"]
+            sagemaker_session=sagemaker_session, job_name=training_job_name
         )
         estimator._current_job_name = estimator.latest_training_job.name
         estimator.latest_training_job.wait()
@@ -268,10 +268,9 @@ def _prepare_init_params_from_job_description(cls, job_details, model_channel_na
         init_params["py_version"] = py_version
 
         if framework and framework != cls.__framework_name__:
-            training_job_name = init_params["base_job_name"]
             raise ValueError(
                 "Training job: {} didn't use image for requested framework".format(
-                    training_job_name
+                    job_details["TrainingJobName"]
                 )
             )
         init_params["framework_version"] = framework_version_from_tag(tag)

tests/unit/test_estimator.py

@@ -21,7 +21,7 @@
 import pytest
 from mock import ANY, MagicMock, Mock, patch
 
-from sagemaker import vpc_utils
+from sagemaker import utils, vpc_utils
 from sagemaker.amazon.amazon_estimator import registry
 from sagemaker.algorithm import AlgorithmEstimator
 from sagemaker.estimator import Estimator, EstimatorBase, Framework, _TrainingJob
@@ -796,6 +796,19 @@ def test_attach_framework_with_inter_container_traffic_encryption_flag(sagemaker
     assert framework_estimator.encrypt_inter_container_traffic is True
 
 
+def test_attach_framework_base_from_generated_name(sagemaker_session):
+    sagemaker_session.sagemaker_client.describe_training_job = Mock(
+        name="describe_training_job", return_value=RETURNED_JOB_DESCRIPTION
+    )
+
+    base_job_name = "neo"
+    framework_estimator = DummyFramework.attach(
+        training_job_name=utils.name_from_base("neo"), sagemaker_session=sagemaker_session
+    )
+
+    assert framework_estimator.base_job_name == base_job_name
+
+
 @patch("time.strftime", return_value=TIMESTAMP)
 def test_fit_verify_job_name(strftime, sagemaker_session):
     fw = DummyFramework(

tests/unit/test_tuner.py

@@ -787,7 +787,7 @@ def test_deploy_default(tuner):
 
     tuner.sagemaker_session.sagemaker_client.describe_hyper_parameter_tuning_job = Mock(
         name="describe_hyper_parameter_tuning_job",
-        return_value={"BestTrainingJob": {"TrainingJobName": JOB_NAME}},
+        return_value={"BestTrainingJob": {"TrainingJobName": TRAINING_JOB_NAME}},
     )
 
     tuner.sagemaker_session.sagemaker_client.list_tags = Mock(
@@ -807,7 +807,7 @@ def test_deploy_default(tuner):
     assert args[2]["ModelDataUrl"] == MODEL_DATA
 
     assert isinstance(predictor, Predictor)
-    assert predictor.endpoint_name.startswith(JOB_NAME)
+    assert predictor.endpoint_name.startswith(TRAINING_JOB_NAME)
     assert predictor.sagemaker_session == tuner.sagemaker_session
 
 
@@ -823,7 +823,7 @@ def test_deploy_estimator_dict(tuner):
         name="describe_hyper_parameter_tuning_job",
         return_value={
             "BestTrainingJob": {
-                "TrainingJobName": JOB_NAME,
+                "TrainingJobName": TRAINING_JOB_NAME,
                 "TrainingJobDefinitionName": ESTIMATOR_NAME,
             }
         },
@@ -846,7 +846,7 @@ def test_deploy_estimator_dict(tuner):
     assert args[2]["ModelDataUrl"] == MODEL_DATA
 
     assert isinstance(predictor, Predictor)
-    assert predictor.endpoint_name.startswith(JOB_NAME)
+    assert predictor.endpoint_name.startswith(TRAINING_JOB_NAME)
     assert predictor.sagemaker_session == tuner.sagemaker_session
 
 

tests/unit/test_utils.py

@@ -115,6 +115,14 @@ def test_name_from_training_arn():
     assert name == "resnet-sgd-tuningjob-11-22-38-46-002-2927640b"
 
 
+def test_base_from_name():
+    name = "mxnet-training-2020-06-29-15-19-25-475"
+    assert "mxnet-training" == sagemaker.utils.base_from_name(name)
+
+    name = "sagemaker-pytorch-200629-1611"
+    assert "sagemaker-pytorch" == sagemaker.utils.base_from_name(name)
+
+
 MESSAGE = "message"
 STATUS = "status"
 TRAINING_JOB_DESCRIPTION_1 = {