Merge branch 'master' into mm-batch-support-on-demand

Author: jerrypeng7773

doc/amazon_sagemaker_model_building_pipeline.rst

@@ -741,6 +741,8 @@ There are a number of properties for a pipeline execution that can only be resol
 - :class:`sagemaker.workflow.execution_variables.ExecutionVariables.PIPELINE_EXECUTION_ARN`: The execution ARN for an execution.
 - :class:`sagemaker.workflow.execution_variables.ExecutionVariables.PIPELINE_NAME`: The name of the pipeline.
 - :class:`sagemaker.workflow.execution_variables.ExecutionVariables.PIPELINE_ARN`: The ARN of the pipeline.
+- :class:`sagemaker.workflow.execution_variables.ExecutionVariables.TRAINING_JOB_NAME`: The name of the training job launched by the training step.
+- :class:`sagemaker.workflow.execution_variables.ExecutionVariables.PROCESSING_JOB_NAME`: The name of the processing job launched by the processing step.
 
 You can use these execution variables as you see fit. The following example uses the :code:`START_DATETIME` execution variable to construct a processing output path:
 

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

@@ -52,7 +52,7 @@ Execution Variables
 .. autoclass:: sagemaker.workflow.execution_variables.ExecutionVariable
 
 .. autoclass:: sagemaker.workflow.execution_variables.ExecutionVariables
-    :members: START_DATETIME, CURRENT_DATETIME, PIPELINE_EXECUTION_ID, PIPELINE_EXECUTION_ARN, PIPELINE_NAME, PIPELINE_ARN
+    :members: START_DATETIME, CURRENT_DATETIME, PIPELINE_EXECUTION_ID, PIPELINE_EXECUTION_ARN, PIPELINE_NAME, PIPELINE_ARN, TRAINING_JOB_NAME, PROCESSING_JOB_NAME
 
 Functions
 ---------

src/sagemaker/amazon/amazon_estimator.py

@@ -16,6 +16,7 @@
 import json
 import logging
 import tempfile
+from typing import Union
 
 from six.moves.urllib.parse import urlparse
 
@@ -27,6 +28,7 @@
 from sagemaker.estimator import EstimatorBase, _TrainingJob
 from sagemaker.inputs import FileSystemInput, TrainingInput
 from sagemaker.utils import sagemaker_timestamp
+from sagemaker.workflow.entities import PipelineVariable
 from sagemaker.workflow.pipeline_context import runnable_by_pipeline
 
 logger = logging.getLogger(__name__)
@@ -304,7 +306,12 @@ class RecordSet(object):
     """Placeholder docstring"""
 
     def __init__(
-        self, s3_data, num_records, feature_dim, s3_data_type="ManifestFile", channel="train"
+        self,
+        s3_data: Union[str, PipelineVariable],
+        num_records: int,
+        feature_dim: int,
+        s3_data_type: Union[str, PipelineVariable] = "ManifestFile",
+        channel: Union[str, PipelineVariable] = "train",
     ):
         """A collection of Amazon :class:~`Record` objects serialized and stored in S3.
 

src/sagemaker/clarify.py

@@ -25,6 +25,8 @@
 
 import tempfile
 from abc import ABC, abstractmethod
+from typing import List, Union
+
 from sagemaker import image_uris, s3, utils
 from sagemaker.processing import ProcessingInput, ProcessingOutput, Processor
 
@@ -63,7 +65,6 @@ def __init__(
             label (str): Target attribute of the model required by bias metrics.
                 Specified as column name or index for CSV dataset or as JSONPath for JSONLines.
                 *Required parameter* except for when the input dataset does not contain the label.
-                Cannot be used at the same time as ``predicted_label``.
             features (str): JSONPath for locating the feature columns for bias metrics if the
                 dataset format is JSONLines.
             dataset_type (str): Format of the dataset. Valid values are ``"text/csv"`` for CSV,
@@ -103,7 +104,7 @@ def __init__(
             predicted_label (str or int): Predicted label of the target attribute of the model
                 required for running bias analysis. Specified as column name or index for CSV data.
                 Clarify uses the predicted labels directly instead of making model inference API
-                calls. Cannot be used at the same time as ``label``.
+                calls.
             excluded_columns (list[int] or list[str]): A list of names or indices of the columns
                 which are to be excluded from making model inference API calls.
 
@@ -922,6 +923,7 @@ def __init__(
             version (str): Clarify version to use.
         """  # noqa E501  # pylint: disable=c0301
         container_uri = image_uris.retrieve("clarify", sagemaker_session.boto_region_name, version)
+        self._last_analysis_config = None
         self.job_name_prefix = job_name_prefix
         super(SageMakerClarifyProcessor, self).__init__(
             role,
@@ -983,10 +985,10 @@ def _run(
                   the Trial Component will be unassociated.
                 * ``'TrialComponentDisplayName'`` is used for display in Amazon SageMaker Studio.
         """
-        analysis_config["methods"]["report"] = {
-            "name": "report",
-            "title": "Analysis Report",
-        }
+        # for debugging: to access locally, i.e. without a need to look for it in an S3 bucket
+        self._last_analysis_config = analysis_config
+        logger.info("Analysis Config: %s", analysis_config)
+
         with tempfile.TemporaryDirectory() as tmpdirname:
             analysis_config_file = os.path.join(tmpdirname, "analysis_config.json")
             with open(analysis_config_file, "w") as f:
@@ -1083,14 +1085,13 @@ def run_pre_training_bias(
                   the Trial Component will be unassociated.
                 * ``'TrialComponentDisplayName'`` is used for display in Amazon SageMaker Studio.
         """  # noqa E501  # pylint: disable=c0301
-        analysis_config = data_config.get_config()
-        analysis_config.update(data_bias_config.get_config())
-        analysis_config["methods"] = {"pre_training_bias": {"methods": methods}}
-        if job_name is None:
-            if self.job_name_prefix:
-                job_name = utils.name_from_base(self.job_name_prefix)
-            else:
-                job_name = utils.name_from_base("Clarify-Pretraining-Bias")
+        analysis_config = _AnalysisConfigGenerator.bias_pre_training(
+            data_config, data_bias_config, methods
+        )
+        # when name is either not provided (is None) or an empty string ("")
+        job_name = job_name or utils.name_from_base(
+            self.job_name_prefix or "Clarify-Pretraining-Bias"
+        )
         return self._run(
             data_config,
             analysis_config,
@@ -1165,21 +1166,13 @@ def run_post_training_bias(
                   the Trial Component will be unassociated.
                 * ``'TrialComponentDisplayName'`` is used for display in Amazon SageMaker Studio.
         """  # noqa E501  # pylint: disable=c0301
-        analysis_config = data_config.get_config()
-        analysis_config.update(data_bias_config.get_config())
-        (
-            probability_threshold,
-            predictor_config,
-        ) = model_predicted_label_config.get_predictor_config()
-        predictor_config.update(model_config.get_predictor_config())
-        analysis_config["methods"] = {"post_training_bias": {"methods": methods}}
-        analysis_config["predictor"] = predictor_config
-        _set(probability_threshold, "probability_threshold", analysis_config)
-        if job_name is None:
-            if self.job_name_prefix:
-                job_name = utils.name_from_base(self.job_name_prefix)
-            else:
-                job_name = utils.name_from_base("Clarify-Posttraining-Bias")
+        analysis_config = _AnalysisConfigGenerator.bias_post_training(
+            data_config, data_bias_config, model_predicted_label_config, methods, model_config
+        )
+        # when name is either not provided (is None) or an empty string ("")
+        job_name = job_name or utils.name_from_base(
+            self.job_name_prefix or "Clarify-Posttraining-Bias"
+        )
         return self._run(
             data_config,
             analysis_config,
@@ -1264,28 +1257,16 @@ def run_bias(
                   the Trial Component will be unassociated.
                 * ``'TrialComponentDisplayName'`` is used for display in Amazon SageMaker Studio.
         """  # noqa E501  # pylint: disable=c0301
-        analysis_config = data_config.get_config()
-        analysis_config.update(bias_config.get_config())
-        analysis_config["predictor"] = model_config.get_predictor_config()
-        if model_predicted_label_config:
-            (
-                probability_threshold,
-                predictor_config,
-            ) = model_predicted_label_config.get_predictor_config()
-            if predictor_config:
-                analysis_config["predictor"].update(predictor_config)
-            if probability_threshold is not None:
-                analysis_config["probability_threshold"] = probability_threshold
-
-        analysis_config["methods"] = {
-            "pre_training_bias": {"methods": pre_training_methods},
-            "post_training_bias": {"methods": post_training_methods},
-        }
-        if job_name is None:
-            if self.job_name_prefix:
-                job_name = utils.name_from_base(self.job_name_prefix)
-            else:
-                job_name = utils.name_from_base("Clarify-Bias")
+        analysis_config = _AnalysisConfigGenerator.bias(
+            data_config,
+            bias_config,
+            model_config,
+            model_predicted_label_config,
+            pre_training_methods,
+            post_training_methods,
+        )
+        # when name is either not provided (is None) or an empty string ("")
+        job_name = job_name or utils.name_from_base(self.job_name_prefix or "Clarify-Bias")
         return self._run(
             data_config,
             analysis_config,
@@ -1370,6 +1351,36 @@ def run_explainability(
                   the Trial Component will be unassociated.
                 * ``'TrialComponentDisplayName'`` is used for display in Amazon SageMaker Studio.
         """  # noqa E501  # pylint: disable=c0301
+        analysis_config = _AnalysisConfigGenerator.explainability(
+            data_config, model_config, model_scores, explainability_config
+        )
+        # when name is either not provided (is None) or an empty string ("")
+        job_name = job_name or utils.name_from_base(
+            self.job_name_prefix or "Clarify-Explainability"
+        )
+        return self._run(
+            data_config,
+            analysis_config,
+            wait,
+            logs,
+            job_name,
+            kms_key,
+            experiment_config,
+        )
+
+
+class _AnalysisConfigGenerator:
+    """Creates analysis_config objects for different type of runs."""
+
+    @classmethod
+    def explainability(
+        cls,
+        data_config: DataConfig,
+        model_config: ModelConfig,
+        model_scores: ModelPredictedLabelConfig,
+        explainability_config: ExplainabilityConfig,
+    ):
+        """Generates a config for Explainability"""
         analysis_config = data_config.get_config()
         predictor_config = model_config.get_predictor_config()
         if isinstance(model_scores, ModelPredictedLabelConfig):
@@ -1406,20 +1417,84 @@ def run_explainability(
             explainability_methods = explainability_config.get_explainability_config()
         analysis_config["methods"] = explainability_methods
         analysis_config["predictor"] = predictor_config
-        if job_name is None:
-            if self.job_name_prefix:
-                job_name = utils.name_from_base(self.job_name_prefix)
-            else:
-                job_name = utils.name_from_base("Clarify-Explainability")
-        return self._run(
-            data_config,
-            analysis_config,
-            wait,
-            logs,
-            job_name,
-            kms_key,
-            experiment_config,
-        )
+        return cls._common(analysis_config)
+
+    @classmethod
+    def bias_pre_training(
+        cls, data_config: DataConfig, bias_config: BiasConfig, methods: Union[str, List[str]]
+    ):
+        """Generates a config for Bias Pre Training"""
+        analysis_config = {
+            **data_config.get_config(),
+            **bias_config.get_config(),
+            "methods": {"pre_training_bias": {"methods": methods}},
+        }
+        return cls._common(analysis_config)
+
+    @classmethod
+    def bias_post_training(
+        cls,
+        data_config: DataConfig,
+        bias_config: BiasConfig,
+        model_predicted_label_config: ModelPredictedLabelConfig,
+        methods: Union[str, List[str]],
+        model_config: ModelConfig,
+    ):
+        """Generates a config for Bias Post Training"""
+        analysis_config = {
+            **data_config.get_config(),
+            **bias_config.get_config(),
+            "predictor": {**model_config.get_predictor_config()},
+            "methods": {"post_training_bias": {"methods": methods}},
+        }
+        if model_predicted_label_config:
+            (
+                probability_threshold,
+                predictor_config,
+            ) = model_predicted_label_config.get_predictor_config()
+            if predictor_config:
+                analysis_config["predictor"].update(predictor_config)
+            _set(probability_threshold, "probability_threshold", analysis_config)
+        return cls._common(analysis_config)
+
+    @classmethod
+    def bias(
+        cls,
+        data_config: DataConfig,
+        bias_config: BiasConfig,
+        model_config: ModelConfig,
+        model_predicted_label_config: ModelPredictedLabelConfig,
+        pre_training_methods: Union[str, List[str]] = "all",
+        post_training_methods: Union[str, List[str]] = "all",
+    ):
+        """Generates a config for Bias"""
+        analysis_config = {
+            **data_config.get_config(),
+            **bias_config.get_config(),
+            "predictor": model_config.get_predictor_config(),
+            "methods": {
+                "pre_training_bias": {"methods": pre_training_methods},
+                "post_training_bias": {"methods": post_training_methods},
+            },
+        }
+        if model_predicted_label_config:
+            (
+                probability_threshold,
+                predictor_config,
+            ) = model_predicted_label_config.get_predictor_config()
+            if predictor_config:
+                analysis_config["predictor"].update(predictor_config)
+            _set(probability_threshold, "probability_threshold", analysis_config)
+        return cls._common(analysis_config)
+
+    @staticmethod
+    def _common(analysis_config):
+        """Extends analysis config with common values"""
+        analysis_config["methods"]["report"] = {
+            "name": "report",
+            "title": "Analysis Report",
+        }
+        return analysis_config
 
 
 def _upload_analysis_config(analysis_config_file, s3_output_path, sagemaker_session, kms_key):