Fix: change the default naming for base_name_from_image

Author: jerrypeng7773

src/sagemaker/estimator.py

@@ -102,6 +102,7 @@ class EstimatorBase(with_metaclass(ABCMeta, object)):  # pylint: disable=too-man
     MPI_CUSTOM_MPI_OPTIONS = "sagemaker_mpi_custom_mpi_options"
     SM_DDP_CUSTOM_MPI_OPTIONS = "sagemaker_distributed_dataparallel_custom_mpi_options"
     CONTAINER_CODE_CHANNEL_SOURCEDIR_PATH = "/opt/ml/input/data/code/sourcedir.tar.gz"
+    JOB_CLASS_NAME = "training-job"
 
     def __init__(
         self,
@@ -576,7 +577,9 @@ def _ensure_base_job_name(self):
         self.base_job_name = (
             self.base_job_name
             or get_jumpstart_base_name_if_jumpstart_model(self.source_dir, self.model_uri)
-            or base_name_from_image(self.training_image_uri())
+            or base_name_from_image(
+                self.training_image_uri(), default_base_name=EstimatorBase.JOB_CLASS_NAME
+            )
         )
 
     def _get_or_create_name(self, name=None):
@@ -982,7 +985,9 @@ def fit(self, inputs=None, wait=True, logs="All", job_name=None, experiment_conf
 
     def _compilation_job_name(self):
         """Placeholder docstring"""
-        base_name = self.base_job_name or base_name_from_image(self.training_image_uri())
+        base_name = self.base_job_name or base_name_from_image(
+            self.training_image_uri(), default_base_name=EstimatorBase.JOB_CLASS_NAME
+        )
         return name_from_base("compilation-" + base_name)
 
     def compile_model(

src/sagemaker/model.py

@@ -635,7 +635,7 @@ def _ensure_base_name_if_needed(self, image_uri, script_uri, model_uri):
             self._base_name = (
                 self._base_name
                 or get_jumpstart_base_name_if_jumpstart_model(script_uri, model_uri)
-                or utils.base_name_from_image(image_uri)
+                or utils.base_name_from_image(image_uri, default_base_name=Model.__name__)
             )
 
     def _set_model_name_if_needed(self):

src/sagemaker/processing.py

@@ -45,6 +45,8 @@
 class Processor(object):
     """Handles Amazon SageMaker Processing tasks."""
 
+    JOB_CLASS_NAME = "processing-job"
+
     def __init__(
         self,
         role,
@@ -280,7 +282,9 @@ def _generate_current_job_name(self, job_name=None):
         if self.base_job_name:
             base_name = self.base_job_name
         else:
-            base_name = base_name_from_image(self.image_uri)
+            base_name = base_name_from_image(
+                self.image_uri, default_base_name=Processor.JOB_CLASS_NAME
+            )
 
         return name_from_base(base_name)
 

src/sagemaker/transformer.py

@@ -25,6 +25,8 @@
 class Transformer(object):
     """A class for handling creating and interacting with Amazon SageMaker transform jobs."""
 
+    JOB_CLASS_NAME = "transform-job"
+
     def __init__(
         self,
         model_name,
@@ -240,7 +242,7 @@ def _retrieve_base_name(self):
         image_uri = self._retrieve_image_uri()
 
         if image_uri:
-            return base_name_from_image(image_uri)
+            return base_name_from_image(image_uri, default_base_name=Transformer.JOB_CLASS_NAME)
 
         return self.model_name
 

src/sagemaker/tuner.py

@@ -29,7 +29,7 @@
 from sagemaker.amazon.hyperparameter import Hyperparameter as hp  # noqa
 from sagemaker.analytics import HyperparameterTuningJobAnalytics
 from sagemaker.deprecations import removed_function
-from sagemaker.estimator import Framework
+from sagemaker.estimator import Framework, EstimatorBase
 from sagemaker.inputs import TrainingInput
 from sagemaker.job import _Job
 from sagemaker.jumpstart.utils import add_jumpstart_tags, get_jumpstart_base_name_if_jumpstart_model
@@ -367,7 +367,9 @@ def _prepare_job_name_for_tuning(self, job_name=None):
                 estimator = (
                     self.estimator or self.estimator_dict[sorted(self.estimator_dict.keys())[0]]
                 )
-                base_name = base_name_from_image(estimator.training_image_uri())
+                base_name = base_name_from_image(
+                    estimator.training_image_uri(), default_base_name=EstimatorBase.JOB_CLASS_NAME
+                )
 
                 jumpstart_base_name = get_jumpstart_base_name_if_jumpstart_model(
                     getattr(estimator, "source_dir", None),

src/sagemaker/utils.py

@@ -91,11 +91,12 @@ def unique_name_from_base(base, max_length=63):
     return "{}-{}-{}".format(trimmed, ts, unique)
 
 
-def base_name_from_image(image):
+def base_name_from_image(image, default_base_name=None):
     """Extract the base name of the image to use as the 'algorithm name' for the job.
 
     Args:
         image (str): Image name.
+        default_base_name (str): The default base name
 
     Returns:
         str: Algorithm name, as extracted from the image name.
@@ -104,7 +105,7 @@ def base_name_from_image(image):
         if is_pipeline_parameter_string(image) and image.default_value:
             image_str = image.default_value
         else:
-            return "base_name"
+            return default_base_name if default_base_name else "base_name"
     else:
         image_str = image
 

src/sagemaker/workflow/airflow.py

@@ -20,6 +20,8 @@
 from sagemaker import fw_utils, job, utils, s3, session, vpc_utils
 from sagemaker.amazon import amazon_estimator
 from sagemaker.tensorflow import TensorFlow
+from sagemaker.estimator import EstimatorBase
+from sagemaker.processing import Processor
 
 
 def prepare_framework(estimator, s3_operations):
@@ -151,7 +153,8 @@ def training_base_config(estimator, inputs=None, job_name=None, mini_batch_size=
         estimator._current_job_name = job_name
     else:
         base_name = estimator.base_job_name or utils.base_name_from_image(
-            estimator.training_image_uri()
+            estimator.training_image_uri(),
+            default_base_name=EstimatorBase.JOB_CLASS_NAME,
         )
         estimator._current_job_name = utils.name_from_base(base_name)
 
@@ -1138,7 +1141,7 @@ def processing_config(
         processor._current_job_name = (
             utils.name_from_base(base_name)
             if base_name is not None
-            else utils.base_name_from_image(processor.image_uri)
+            else utils.base_name_from_image(processor.image_uri, Processor.JOB_CLASS_NAME)
         )
 
     config = {

tests/unit/sagemaker/model/test_model.py

@@ -287,7 +287,7 @@ def test_create_sagemaker_model_generates_model_name(
     )
     model._create_sagemaker_model(INSTANCE_TYPE)
 
-    base_name_from_image.assert_called_with(MODEL_IMAGE)
+    base_name_from_image.assert_called_with(MODEL_IMAGE, default_base_name="Model")
     name_from_base.assert_called_with(base_name_from_image.return_value)
 
     sagemaker_session.create_model.assert_called_with(
@@ -317,7 +317,7 @@ def test_create_sagemaker_model_generates_model_name_each_time(
     model._create_sagemaker_model(INSTANCE_TYPE)
     model._create_sagemaker_model(INSTANCE_TYPE)
 
-    base_name_from_image.assert_called_once_with(MODEL_IMAGE)
+    base_name_from_image.assert_called_once_with(MODEL_IMAGE, default_base_name="Model")
     name_from_base.assert_called_with(base_name_from_image.return_value)
     assert 2 == name_from_base.call_count
 

tests/unit/test_utils.py

@@ -29,6 +29,8 @@
 
 import sagemaker
 from sagemaker.session_settings import SessionSettings
+from tests.unit.sagemaker.workflow.helpers import CustomStep
+from sagemaker.workflow.parameters import ParameterString
 
 BUCKET_WITHOUT_WRITING_PERMISSION = "s3://bucket-without-writing-permission"
 
@@ -82,6 +84,46 @@ def test_name_from_image(base_name_from_image, name_from_base):
     name_from_base.assert_called_with(base_name_from_image.return_value, max_length=max_length)
 
 
+@pytest.mark.parametrize(
+    "inputs",
+    [
+        (
+            CustomStep(name="test-custom-step").properties.OutputDataConfig.S3OutputPath,
+            None,
+            "base_name",
+        ),
+        (
+            CustomStep(name="test-custom-step").properties.OutputDataConfig.S3OutputPath,
+            "whatever",
+            "whatever",
+        ),
+        (ParameterString(name="image_uri"), None, "base_name"),
+        (ParameterString(name="image_uri"), "whatever", "whatever"),
+        (
+            ParameterString(
+                name="image_uri",
+                default_value="922956235488.dkr.ecr.us-west-2.amazonaws.com/analyzer",
+            ),
+            None,
+            "analyzer",
+        ),
+        (
+            ParameterString(
+                name="image_uri",
+                default_value="922956235488.dkr.ecr.us-west-2.amazonaws.com/analyzer",
+            ),
+            "whatever",
+            "analyzer",
+        ),
+    ],
+)
+def test_base_name_from_image_with_pipeline_param(inputs):
+    image, default_base_name, expected = inputs
+    assert expected == sagemaker.utils.base_name_from_image(
+        image=image, default_base_name=default_base_name
+    )
+
+
 @patch("sagemaker.utils.sagemaker_timestamp")
 def test_name_from_base(sagemaker_timestamp):
     sagemaker.utils.name_from_base(NAME, short=False)