feat: Add PipelineDefinitionConfig to pipelines to toggle custom job … (#3946)

Co-authored-by: svia3 <[email protected]>

Author: svia3

doc/workflows/pipelines/sagemaker.workflow.pipelines.rst

@@ -103,6 +103,11 @@ Parallelism Configuration
 .. autoclass:: sagemaker.workflow.parallelism_config.ParallelismConfiguration
     :members:
 
+Pipeline Definition Config
+--------------------------
+
+.. autoclass:: sagemaker.workflow.pipeline_definition_config.PipelineDefinitionConfig
+
 Pipeline Experiment Config
 --------------------------
 

src/sagemaker/model_monitor/clarify_model_monitoring.py

@@ -175,6 +175,7 @@ def _create_baselining_processor(self):
             network_config=self.network_config,
         )
         baselining_processor.image_uri = self.image_uri
+        baselining_processor.base_job_name = self.base_job_name
         return baselining_processor
 
     def _upload_analysis_config(self, analysis_config, output_s3_uri, job_definition_name):

src/sagemaker/workflow/_utils.py

@@ -34,6 +34,7 @@
 )
 from sagemaker.utils import _save_model, download_file_from_url
 from sagemaker.workflow.retry import RetryPolicy
+from sagemaker.workflow.utilities import trim_request_dict
 
 if TYPE_CHECKING:
     from sagemaker.workflow.step_collections import StepCollection
@@ -412,6 +413,8 @@ def __init__(
     @property
     def arguments(self) -> RequestType:
         """The arguments dict that are used to call `create_model_package`."""
+        from sagemaker.workflow.utilities import _pipeline_config
+
         model_name = self.name
 
         if self.step_args:
@@ -492,9 +495,9 @@ def arguments(self) -> RequestType:
         if "Description" in request_dict:
             request_dict.pop("Description")
             logger.warning(warn_msg_template, "Description")
-        if "ModelPackageName" in request_dict:
-            request_dict.pop("ModelPackageName")
-            logger.warning(warn_msg_template, "ModelPackageName")
+
+        # Continue to pop job name if not explicitly opted-in via config
+        request_dict = trim_request_dict(request_dict, "ModelPackageName", _pipeline_config)
 
         return request_dict
 

src/sagemaker/workflow/automl_step.py

@@ -23,7 +23,7 @@
 from sagemaker.workflow.properties import Properties
 from sagemaker.workflow.retry import RetryPolicy
 from sagemaker.workflow.steps import ConfigurableRetryStep, CacheConfig, Step, StepTypeEnum
-from sagemaker.workflow.utilities import validate_step_args_input
+from sagemaker.workflow.utilities import validate_step_args_input, trim_request_dict
 from sagemaker.workflow.step_collections import StepCollection
 
 
@@ -89,10 +89,11 @@ def arguments(self) -> RequestType:
         NOTE: The `CreateAutoMLJob` request is not quite the
             args list that workflow needs.
 
-        The `AutoMLJobName`, `ModelDeployConfig` and `GenerateCandidateDefinitionsOnly`
+        `ModelDeployConfig` and `GenerateCandidateDefinitionsOnly`
             attribute cannot be included.
         """
         from sagemaker.workflow.utilities import execute_job_functions
+        from sagemaker.workflow.utilities import _pipeline_config
 
         # execute fit function in AutoML with saved parameters,
         # and store args in PipelineSession's _context
@@ -114,7 +115,10 @@ def arguments(self) -> RequestType:
             request_dict.pop("ModelDeployConfig", None)
         if "GenerateCandidateDefinitionsOnly" in request_dict:
             request_dict.pop("GenerateCandidateDefinitionsOnly", None)
-        request_dict.pop("AutoMLJobName", None)
+        # Continue to pop job name if not explicitly opted-in via config
+        # AutoML Trims to AutoMLJo-2023-06-23-22-57-39-083
+        request_dict = trim_request_dict(request_dict, "AutoMLJobName", _pipeline_config)
+
         return request_dict
 
     @property

src/sagemaker/workflow/clarify_check_step.py

@@ -44,6 +44,7 @@
 from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import Step, StepTypeEnum, CacheConfig
 from sagemaker.workflow.check_job_config import CheckJobConfig
+from sagemaker.workflow.utilities import trim_request_dict
 
 _DATA_BIAS_TYPE = "DATA_BIAS"
 _MODEL_BIAS_TYPE = "MODEL_BIAS"
@@ -253,6 +254,8 @@ def __init__(
     @property
     def arguments(self) -> RequestType:
         """The arguments dict that is used to define the ClarifyCheck step."""
+        from sagemaker.workflow.utilities import _pipeline_config
+
         normalized_inputs, normalized_outputs = self._baselining_processor._normalize_args(
             inputs=[self._processing_params["config_input"], self._processing_params["data_input"]],
             outputs=[self._processing_params["result_output"]],
@@ -266,8 +269,8 @@ def arguments(self) -> RequestType:
         request_dict = self._baselining_processor.sagemaker_session._get_process_request(
             **process_args
         )
-        if "ProcessingJobName" in request_dict:
-            request_dict.pop("ProcessingJobName")
+        # Continue to pop job name if not explicitly opted-in via config
+        request_dict = trim_request_dict(request_dict, "ProcessingJobName", _pipeline_config)
 
         return request_dict
 

src/sagemaker/workflow/pipeline.py

@@ -37,6 +37,7 @@
 )
 from sagemaker.workflow.execution_variables import ExecutionVariables
 from sagemaker.workflow.parameters import Parameter
+from sagemaker.workflow.pipeline_definition_config import PipelineDefinitionConfig
 from sagemaker.workflow.pipeline_experiment_config import PipelineExperimentConfig
 from sagemaker.workflow.parallelism_config import ParallelismConfiguration
 from sagemaker.workflow.properties import Properties
@@ -52,6 +53,8 @@
     ExecutionVariables.PIPELINE_NAME, ExecutionVariables.PIPELINE_EXECUTION_ID
 )
 
+_DEFAULT_DEFINITION_CFG = PipelineDefinitionConfig(use_custom_job_prefix=False)
+
 
 class Pipeline(Entity):
     """Pipeline for workflow."""
@@ -63,6 +66,7 @@ def __init__(
         pipeline_experiment_config: Optional[PipelineExperimentConfig] = _DEFAULT_EXPERIMENT_CFG,
         steps: Optional[Sequence[Union[Step, StepCollection]]] = None,
         sagemaker_session: Optional[Session] = None,
+        pipeline_definition_config: Optional[PipelineDefinitionConfig] = _DEFAULT_DEFINITION_CFG,
     ):
         """Initialize a Pipeline
 
@@ -84,12 +88,16 @@ def __init__(
             sagemaker_session (sagemaker.session.Session): Session object that manages interactions
                 with Amazon SageMaker APIs and any other AWS services needed. If not specified, the
                 pipeline creates one using the default AWS configuration chain.
+            pipeline_definition_config (Optional[PipelineDefinitionConfig]): If set,
+                the workflow customizes the pipeline definition using the configurations
+                specified. By default, custom job-prefixing is turned off.
         """
         self.name = name
         self.parameters = parameters if parameters else []
         self.pipeline_experiment_config = pipeline_experiment_config
         self.steps = steps if steps else []
         self.sagemaker_session = sagemaker_session if sagemaker_session else Session()
+        self.pipeline_definition_config = pipeline_definition_config
 
         self._version = "2020-12-01"
         self._metadata = dict()
@@ -105,7 +113,11 @@ def to_request(self) -> RequestType:
             "PipelineExperimentConfig": self.pipeline_experiment_config.to_request()
             if self.pipeline_experiment_config is not None
             else None,
-            "Steps": build_steps(self.steps, self.name),
+            "Steps": build_steps(
+                self.steps,
+                self.name,
+                self.pipeline_definition_config,
+            ),
         }
 
     def create(

src/sagemaker/workflow/pipeline_context.py

@@ -20,6 +20,7 @@
 
 from sagemaker.session import Session, SessionSettings
 from sagemaker.local import LocalSession
+from sagemaker.workflow.pipeline_definition_config import PipelineDefinitionConfig
 
 
 class _StepArguments:
@@ -87,13 +88,23 @@ class _PipelineConfig:
         step_name (str): step name
         code_hash (str): a hash of the code artifact for the particular step
         config_hash (str): a hash of the config artifact for the particular step (Processing)
+        pipeline_definition_config (PipelineDefinitionConfig): a configuration used to toggle
+            feature flags persistent in a pipeline definition
     """
 
-    def __init__(self, pipeline_name, step_name, code_hash, config_hash):
+    def __init__(
+        self,
+        pipeline_name: str,
+        step_name: str,
+        code_hash: str,
+        config_hash: str,
+        pipeline_definition_config: PipelineDefinitionConfig,
+    ):
         self.pipeline_name = pipeline_name
         self.step_name = step_name
         self.code_hash = code_hash
         self.config_hash = config_hash
+        self.pipeline_definition_config = pipeline_definition_config
 
 
 class PipelineSession(Session):

src/sagemaker/workflow/pipeline_definition_config.py

@@ -0,0 +1,31 @@
+# Copyright 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
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Pipeline experiment config for SageMaker pipeline."""
+from __future__ import absolute_import
+
+
+class PipelineDefinitionConfig:
+    """Pipeline Definition Configuration for SageMaker pipeline."""
+
+    def __init__(self, use_custom_job_prefix: bool):
+        """Create a `PipelineDefinitionConfig`.
+
+        Examples: Use a `PipelineDefinitionConfig` to turn on custom job prefixing::
+
+            PipelineDefinitionConfig(use_custom_job_prefix=True)
+
+        Args:
+            use_custom_job_prefix (bool): A feature flag to toggle on/off custom name prefixing
+                during pipeline orchestration.
+        """
+        self.use_custom_job_prefix = use_custom_job_prefix

src/sagemaker/workflow/quality_check_step.py

@@ -31,6 +31,7 @@
 from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import Step, StepTypeEnum, CacheConfig
 from sagemaker.workflow.check_job_config import CheckJobConfig
+from sagemaker.workflow.utilities import trim_request_dict
 
 _CONTAINER_BASE_PATH = "/opt/ml/processing"
 _CONTAINER_INPUT_PATH = "input"
@@ -225,6 +226,8 @@ def __init__(
     @property
     def arguments(self) -> RequestType:
         """The arguments dict that is used to define the QualityCheck step."""
+        from sagemaker.workflow.utilities import _pipeline_config
+
         normalized_inputs, normalized_outputs = self._baselining_processor._normalize_args(
             inputs=self._baseline_job_inputs,
             outputs=[self._baseline_output],
@@ -238,8 +241,8 @@ def arguments(self) -> RequestType:
         request_dict = self._baselining_processor.sagemaker_session._get_process_request(
             **process_args
         )
-        if "ProcessingJobName" in request_dict:
-            request_dict.pop("ProcessingJobName")
+        # Continue to pop job name if not explicitly opted-in via config
+        request_dict = trim_request_dict(request_dict, "ProcessingJobName", _pipeline_config)
 
         return request_dict
 

src/sagemaker/workflow/steps.py

@@ -50,6 +50,7 @@
 from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.functions import Join, JsonGet
 from sagemaker.workflow.retry import RetryPolicy
+from sagemaker.workflow.utilities import trim_request_dict
 
 if TYPE_CHECKING:
     from sagemaker.workflow.step_collections import StepCollection
@@ -471,9 +472,10 @@ def arguments(self) -> RequestType:
         """The arguments dictionary that is used to call `create_training_job`.
 
         NOTE: The `CreateTrainingJob` request is not quite the args list that workflow needs.
-        The `TrainingJobName` and `ExperimentConfig` attributes cannot be included.
+        `ExperimentConfig` attribute cannot be included.
         """
         from sagemaker.workflow.utilities import execute_job_functions
+        from sagemaker.workflow.utilities import _pipeline_config
 
         if self.step_args:
             # execute fit function with saved parameters,
@@ -493,7 +495,9 @@ def arguments(self) -> RequestType:
         if "HyperParameters" in request_dict:
             request_dict["HyperParameters"].pop("sagemaker_job_name", None)
 
-        request_dict.pop("TrainingJobName", None)
+        # Continue to pop job name if not explicitly opted-in via config
+        request_dict = trim_request_dict(request_dict, "TrainingJobName", _pipeline_config)
+
         Step._trim_experiment_config(request_dict)
 
         return request_dict
@@ -595,8 +599,8 @@ def arguments(self) -> RequestType:
         """The arguments dictionary that is used to call `create_model`.
 
         NOTE: The `CreateModelRequest` is not quite the args list that workflow needs.
-        `ModelName` cannot be included in the arguments.
         """
+        from sagemaker.workflow.utilities import _pipeline_config
 
         if self.step_args:
             request_dict = self.step_args
@@ -620,7 +624,9 @@ def arguments(self) -> RequestType:
                     vpc_config=self.model.vpc_config,
                     enable_network_isolation=self.model.enable_network_isolation(),
                 )
-        request_dict.pop("ModelName", None)
+
+        # Continue to pop job name if not explicitly opted-in via config
+        request_dict = trim_request_dict(request_dict, "ModelName", _pipeline_config)
 
         return request_dict
 
@@ -702,9 +708,10 @@ def arguments(self) -> RequestType:
         """The arguments dictionary that is used to call `create_transform_job`.
 
         NOTE: The `CreateTransformJob` request is not quite the args list that workflow needs.
-        `TransformJobName` and `ExperimentConfig` cannot be included in the arguments.
+        `ExperimentConfig` cannot be included in the arguments.
         """
         from sagemaker.workflow.utilities import execute_job_functions
+        from sagemaker.workflow.utilities import _pipeline_config
 
         if self.step_args:
             # execute transform function with saved parameters,
@@ -733,7 +740,9 @@ def arguments(self) -> RequestType:
                 **transform_args
             )
 
-        request_dict.pop("TransformJobName", None)
+        # Continue to pop job name if not explicitly opted-in via config
+        request_dict = trim_request_dict(request_dict, "TransformJobName", _pipeline_config)
+
         Step._trim_experiment_config(request_dict)
 
         return request_dict
@@ -864,9 +873,10 @@ def arguments(self) -> RequestType:
         """The arguments dictionary that is used to call `create_processing_job`.
 
         NOTE: The `CreateProcessingJob` request is not quite the args list that workflow needs.
-        `ProcessingJobName` and `ExperimentConfig` cannot be included in the arguments.
+        `ExperimentConfig` cannot be included in the arguments.
         """
         from sagemaker.workflow.utilities import execute_job_functions
+        from sagemaker.workflow.utilities import _pipeline_config
 
         if self.step_args:
             # execute run function with saved parameters,
@@ -890,7 +900,9 @@ def arguments(self) -> RequestType:
             )
             request_dict = self.processor.sagemaker_session._get_process_request(**process_args)
 
-        request_dict.pop("ProcessingJobName", None)
+        # Continue to pop job name if not explicitly opted-in via config
+        request_dict = trim_request_dict(request_dict, "ProcessingJobName", _pipeline_config)
+
         Step._trim_experiment_config(request_dict)
 
         return request_dict
@@ -1025,9 +1037,9 @@ def arguments(self) -> RequestType:
 
         NOTE: The `CreateHyperParameterTuningJob` request is not quite the
             args list that workflow needs.
-        The `HyperParameterTuningJobName` attribute cannot be included.
         """
         from sagemaker.workflow.utilities import execute_job_functions
+        from sagemaker.workflow.utilities import _pipeline_config
 
         if self.step_args:
             # execute fit function with saved parameters,
@@ -1048,7 +1060,11 @@ def arguments(self) -> RequestType:
             tuner_args = _TuningJob._get_tuner_args(self.tuner, self.inputs)
             request_dict = self.tuner.sagemaker_session._get_tuning_request(**tuner_args)
 
-        request_dict.pop("HyperParameterTuningJobName", None)
+        # Continue to pop job name if not explicitly opted-in via config
+        request_dict = trim_request_dict(
+            request_dict, "HyperParameterTuningJobName", _pipeline_config
+        )
+
         return request_dict
 
     @property