feat: add experiment_config for clarify processing job (#2287)

* feat: add experiment_config for clarify processing job

* chore: add comment

Author: xiaoyi-cheng

src/sagemaker/clarify.py

@@ -403,6 +403,7 @@ def _run(
         logs,
         job_name,
         kms_key,
+        experiment_config,
     ):
         """Runs a ProcessingJob with the Sagemaker Clarify container and an analysis config.
 
@@ -415,6 +416,9 @@ def _run(
             job_name (str): Processing job name.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
                 user code file (default: None).
+            experiment_config (dict[str, str]): Experiment management configuration.
+                Dictionary contains three optional keys:
+                'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
         """
         analysis_config["methods"]["report"] = {"name": "report", "title": "Analysis Report"}
         with tempfile.TemporaryDirectory() as tmpdirname:
@@ -457,6 +461,7 @@ def _run(
                 logs=logs,
                 job_name=job_name,
                 kms_key=kms_key,
+                experiment_config=experiment_config,
             )
 
     def run_pre_training_bias(
@@ -468,6 +473,7 @@ def run_pre_training_bias(
         logs=True,
         job_name=None,
         kms_key=None,
+        experiment_config=None,
     ):
         """Runs a ProcessingJob to compute the requested bias 'methods' of the input data.
 
@@ -487,13 +493,16 @@ def run_pre_training_bias(
                 "Clarify-Pretraining-Bias" and current timestamp.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
                 user code file (default: None).
+            experiment_config (dict[str, str]): Experiment management configuration.
+                Dictionary contains three optional keys:
+                'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
         """
         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:
             job_name = utils.name_from_base("Clarify-Pretraining-Bias")
-        self._run(data_config, analysis_config, wait, logs, job_name, kms_key)
+        self._run(data_config, analysis_config, wait, logs, job_name, kms_key, experiment_config)
 
     def run_post_training_bias(
         self,
@@ -506,6 +515,7 @@ def run_post_training_bias(
         logs=True,
         job_name=None,
         kms_key=None,
+        experiment_config=None,
     ):
         """Runs a ProcessingJob to compute the requested bias 'methods' of the model predictions.
 
@@ -532,6 +542,9 @@ def run_post_training_bias(
                 "Clarify-Posttraining-Bias" and current timestamp.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
                 user code file (default: None).
+            experiment_config (dict[str, str]): Experiment management configuration.
+                Dictionary contains three optional keys:
+                'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
         """
         analysis_config = data_config.get_config()
         analysis_config.update(data_bias_config.get_config())
@@ -545,7 +558,7 @@ def run_post_training_bias(
         _set(probability_threshold, "probability_threshold", analysis_config)
         if job_name is None:
             job_name = utils.name_from_base("Clarify-Posttraining-Bias")
-        self._run(data_config, analysis_config, wait, logs, job_name, kms_key)
+        self._run(data_config, analysis_config, wait, logs, job_name, kms_key, experiment_config)
 
     def run_bias(
         self,
@@ -559,6 +572,7 @@ def run_bias(
         logs=True,
         job_name=None,
         kms_key=None,
+        experiment_config=None,
     ):
         """Runs a ProcessingJob to compute the requested bias 'methods' of the model predictions.
 
@@ -589,6 +603,9 @@ def run_bias(
                 "Clarify-Bias" and current timestamp.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
                 user code file (default: None).
+            experiment_config (dict[str, str]): Experiment management configuration.
+                Dictionary contains three optional keys:
+                'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
         """
         analysis_config = data_config.get_config()
         analysis_config.update(bias_config.get_config())
@@ -609,7 +626,7 @@ def run_bias(
         }
         if job_name is None:
             job_name = utils.name_from_base("Clarify-Bias")
-        self._run(data_config, analysis_config, wait, logs, job_name, kms_key)
+        self._run(data_config, analysis_config, wait, logs, job_name, kms_key, experiment_config)
 
     def run_explainability(
         self,
@@ -621,6 +638,7 @@ def run_explainability(
         logs=True,
         job_name=None,
         kms_key=None,
+        experiment_config=None,
     ):
         """Runs a ProcessingJob computing for each example in the input the feature importance.
 
@@ -649,6 +667,9 @@ def run_explainability(
                 "Clarify-Explainability" and current timestamp.
             kms_key (str): The ARN of the KMS key that is used to encrypt the
                 user code file (default: None).
+            experiment_config (dict[str, str]): Experiment management configuration.
+                Dictionary contains three optional keys:
+                'ExperimentName', 'TrialName', and 'TrialComponentDisplayName'.
         """
         analysis_config = data_config.get_config()
         predictor_config = model_config.get_predictor_config()
@@ -657,7 +678,7 @@ def run_explainability(
         analysis_config["predictor"] = predictor_config
         if job_name is None:
             job_name = utils.name_from_base("Clarify-Explainability")
-        self._run(data_config, analysis_config, wait, logs, job_name, kms_key)
+        self._run(data_config, analysis_config, wait, logs, job_name, kms_key, experiment_config)
 
 
 def _upload_analysis_config(analysis_config_file, s3_output_path, sagemaker_session, kms_key):

tests/unit/test_clarify.py

@@ -287,7 +287,11 @@ def shap_config():
 def test_pre_training_bias(clarify_processor, data_config, data_bias_config):
     with patch.object(SageMakerClarifyProcessor, "_run", return_value=None) as mock_method:
         clarify_processor.run_pre_training_bias(
-            data_config, data_bias_config, wait=True, job_name="test"
+            data_config,
+            data_bias_config,
+            wait=True,
+            job_name="test",
+            experiment_config={"ExperimentName": "AnExperiment"},
         )
         expected_analysis_config = {
             "dataset_type": "text/csv",
@@ -304,7 +308,13 @@ def test_pre_training_bias(clarify_processor, data_config, data_bias_config):
             "methods": {"pre_training_bias": {"methods": "all"}},
         }
         mock_method.assert_called_once_with(
-            data_config, expected_analysis_config, True, True, "test", None
+            data_config,
+            expected_analysis_config,
+            True,
+            True,
+            "test",
+            None,
+            {"ExperimentName": "AnExperiment"},
         )
 
 
@@ -319,6 +329,7 @@ def test_post_training_bias(
             model_predicted_label_config,
             wait=True,
             job_name="test",
+            experiment_config={"ExperimentName": "AnExperiment"},
         )
         expected_analysis_config = {
             "dataset_type": "text/csv",
@@ -340,14 +351,26 @@ def test_post_training_bias(
             },
         }
         mock_method.assert_called_once_with(
-            data_config, expected_analysis_config, True, True, "test", None
+            data_config,
+            expected_analysis_config,
+            True,
+            True,
+            "test",
+            None,
+            {"ExperimentName": "AnExperiment"},
         )
 
 
 def test_shap(clarify_processor, data_config, model_config, shap_config):
     with patch.object(SageMakerClarifyProcessor, "_run", return_value=None) as mock_method:
         clarify_processor.run_explainability(
-            data_config, model_config, shap_config, model_scores=None, wait=True, job_name="test"
+            data_config,
+            model_config,
+            shap_config,
+            model_scores=None,
+            wait=True,
+            job_name="test",
+            experiment_config={"ExperimentName": "AnExperiment"},
         )
         expected_analysis_config = {
             "dataset_type": "text/csv",
@@ -380,5 +403,11 @@ def test_shap(clarify_processor, data_config, model_config, shap_config):
             },
         }
         mock_method.assert_called_once_with(
-            data_config, expected_analysis_config, True, True, "test", None
+            data_config,
+            expected_analysis_config,
+            True,
+            True,
+            "test",
+            None,
+            {"ExperimentName": "AnExperiment"},
         )