address PR comments

Author: laurenyu

src/sagemaker/algorithm.py

@@ -391,6 +391,11 @@ def _is_marketplace(self):
         """Placeholder docstring"""
         return "ProductId" in self.algorithm_spec
 
+    def _ensure_base_job_name(self):
+        """Set ``self.base_job_name`` if it is not set already."""
+        if self.base_job_name is None:
+            self.base_job_name = self.algorithm_arn.split("/")[-1]
+
     def _prepare_for_training(self, job_name=None):
         # Validate hyperparameters
         # an explicit call to set_hyperparameters() will also validate the hyperparameters

src/sagemaker/estimator.py

@@ -331,14 +331,25 @@ def prepare_workflow_for_training(self, job_name=None):
     def _ensure_base_job_name(self):
         """Set ``self.base_job_name`` if it is not set already."""
         # honor supplied base_job_name or generate it
-        if self.base_job_name:
-            return
-
-        if isinstance(self, sagemaker.algorithm.AlgorithmEstimator):
-            self.base_job_name = self.algorithm_arn.split("/")[-1]  # pylint: disable=no-member
-        else:
+        if self.base_job_name is None:
             self.base_job_name = base_name_from_image(self.train_image())
 
+    def _get_or_create_name(self, name=None):
+        """Generate a name based on the base job name or training image if needed.
+
+        Args:
+            name (str): User-supplied name. If not specified, a name is generated from
+                the base job name or training image.
+
+        Returns:
+            str: Either the user-supplied name or a generated name.
+        """
+        if name:
+            return name
+
+        self._ensure_base_job_name()
+        return name_from_base(self.base_job_name)
+
     def _prepare_for_training(self, job_name=None):
         """Set any values in the estimator that need to be set before training.
 
@@ -347,11 +358,7 @@ def _prepare_for_training(self, job_name=None):
                 specified, one is generated, using the base name given to the
                 constructor if applicable.
         """
-        if job_name is not None:
-            self._current_job_name = job_name
-        else:
-            self._ensure_base_job_name()
-            self._current_job_name = name_from_base(self.base_job_name)
+        self._current_job_name = self._get_or_create_name(job_name)
 
         # if output_path was specified we use it otherwise initialize here.
         # For Local Mode with local_code=True we don't need an explicit output_path
@@ -910,9 +917,7 @@ def transformer(
                 based on the training image name and current timestamp.
         """
         tags = tags or self.tags
-
-        self._ensure_base_job_name()
-        model_name = model_name or name_from_base(self.base_job_name)
+        model_name = self._get_or_create_name(model_name)
 
         if self.latest_training_job is None:
             logging.warning(
@@ -2012,9 +2017,7 @@ def transformer(
         """
         role = role or self.role
         tags = tags or self.tags
-
-        self._ensure_base_job_name()
-        model_name = model_name or name_from_base(self.base_job_name)
+        model_name = self._get_or_create_name(model_name)
 
         if self.latest_training_job is not None:
             if enable_network_isolation is None:

src/sagemaker/mxnet/estimator.py

@@ -15,7 +15,6 @@
 
 import logging
 
-from sagemaker import utils
 from sagemaker.estimator import Framework
 from sagemaker.fw_utils import (
     framework_name_from_image,
@@ -219,9 +218,7 @@ def create_model(
         if "image" not in kwargs:
             kwargs["image"] = image_name or self.image_name
 
-        if "name" not in kwargs:
-            self._ensure_base_job_name()
-            kwargs["name"] = utils.name_from_base(self.base_job_name)
+        kwargs["name"] = self._get_or_create_name(kwargs.get("name"))
 
         return MXNetModel(
             self.model_data,

src/sagemaker/pytorch/estimator.py

@@ -15,7 +15,6 @@
 
 import logging
 
-from sagemaker import utils
 from sagemaker.estimator import Framework
 from sagemaker.fw_utils import (
     framework_name_from_image,
@@ -171,9 +170,7 @@ def create_model(
         if "image" not in kwargs:
             kwargs["image"] = self.image_name
 
-        if "name" not in kwargs:
-            self._ensure_base_job_name()
-            kwargs["name"] = utils.name_from_base(self.base_job_name)
+        kwargs["name"] = self._get_or_create_name(kwargs.get("name"))
 
         return PyTorchModel(
             self.model_data,

src/sagemaker/rl/estimator.py

@@ -17,7 +17,7 @@
 import logging
 import re
 
-from sagemaker import fw_utils, utils
+from sagemaker import fw_utils
 from sagemaker.estimator import Framework
 from sagemaker.model import FrameworkModel, SAGEMAKER_OUTPUT_LOCATION
 from sagemaker.mxnet.model import MXNetModel
@@ -218,18 +218,18 @@ def create_model(
         Raises:
             ValueError: If image_name is not specified and framework enum is not valid.
         """
-        self._ensure_base_job_name()
 
         base_args = dict(
             model_data=self.model_data,
             role=role or self.role,
             image=kwargs.get("image", self.image_name),
-            name=kwargs.get("name", utils.name_from_base(self.base_job_name)),
             container_log_level=self.container_log_level,
             sagemaker_session=self.sagemaker_session,
             vpc_config=self.get_vpc_config(vpc_config_override),
         )
 
+        base_args["name"] = self._get_or_create_name(kwargs.get("name"))
+
         if not entry_point and (source_dir or dependencies):
             raise AttributeError("Please provide an `entry_point`.")
 

src/sagemaker/sklearn/estimator.py

@@ -15,7 +15,6 @@
 
 import logging
 
-from sagemaker import utils
 from sagemaker.estimator import Framework
 from sagemaker.fw_registry import default_framework_uri
 from sagemaker.fw_utils import (
@@ -187,17 +186,14 @@ def create_model(
             object. See :func:`~sagemaker.sklearn.model.SKLearnModel` for full details.
         """
         role = role or self.role
+        kwargs["name"] = self._get_or_create_name(kwargs.get("name"))
 
         if "image" not in kwargs:
             kwargs["image"] = self.image_name
 
         if "enable_network_isolation" not in kwargs:
             kwargs["enable_network_isolation"] = self.enable_network_isolation()
 
-        if "name" not in kwargs:
-            self._ensure_base_job_name()
-            kwargs["name"] = utils.name_from_base(self.base_job_name)
-
         return SKLearnModel(
             self.model_data,
             role,

src/sagemaker/tensorflow/estimator.py

@@ -270,13 +270,11 @@ def create_model(
             sagemaker.tensorflow.model.TensorFlowModel: A ``TensorFlowModel`` object.
                 See :class:`~sagemaker.tensorflow.model.TensorFlowModel` for full details.
         """
+        kwargs["name"] = self._get_or_create_name(kwargs.get("name"))
+
         if "image" not in kwargs:
             kwargs["image"] = self.image_name
 
-        if "name" not in kwargs:
-            self._ensure_base_job_name()
-            kwargs["name"] = utils.name_from_base(self.base_job_name)
-
         if "enable_network_isolation" not in kwargs:
             kwargs["enable_network_isolation"] = self.enable_network_isolation()
 
@@ -446,9 +444,7 @@ def transformer(
                 based on the training image name and current timestamp.
         """
         role = role or self.role
-
-        self._ensure_base_job_name()
-        model_name = model_name or utils.name_from_base(self.base_job_name)
+        model_name = self._get_or_create_name(model_name)
 
         if self.latest_training_job is None:
             logging.warning(

src/sagemaker/xgboost/estimator.py

@@ -15,7 +15,6 @@
 
 import logging
 
-from sagemaker import utils
 from sagemaker.estimator import Framework, _TrainingJob
 from sagemaker.fw_registry import default_framework_uri
 from sagemaker.fw_utils import (
@@ -165,14 +164,11 @@ def create_model(
                 See :func:`~sagemaker.xgboost.model.XGBoostModel` for full details.
         """
         role = role or self.role
+        kwargs["name"] = self._get_or_create_name(kwargs.get("name"))
 
         if "image" not in kwargs:
             kwargs["image"] = self.image_name
 
-        if "name" not in kwargs:
-            self._ensure_base_job_name()
-            kwargs["name"] = utils.name_from_base(self.base_job_name)
-
         return XGBoostModel(
             self.model_data,
             role,