change: add s3_analysis_config_output_path field in DataConfig constructor (#2698)

Co-authored-by: Dewen Qi <qidewen@amazon.com>
Co-authored-by: Shreya Pandit <shreya.pandit25@gmail.com>
Co-authored-by: Navin Soni <navinsoni89@gmail.com>
Co-authored-by: Ahsan Khan <ahsan.al.zaki@gmail.com>

Author: 5 people

src/sagemaker/clarify.py

@@ -31,6 +31,7 @@ def __init__(
         self,
         s3_data_input_path,
         s3_output_path,
+        s3_analysis_config_output_path=None,
         label=None,
         headers=None,
         features=None,
@@ -43,6 +44,9 @@ def __init__(
         Args:
             s3_data_input_path (str): Dataset S3 prefix/object URI.
             s3_output_path (str): S3 prefix to store the output.
+            s3_analysis_config_output_path (str): S3 prefix to store the analysis_config output
+                If this field is None, then the s3_output_path will be used
+                to store the analysis_config output
             label (str): Target attribute of the model required by bias metrics (optional for SHAP)
                 Specified as column name or index for CSV dataset, or as JSONPath for JSONLines.
             headers (list[str]): A list of column names in the input dataset.
@@ -61,6 +65,7 @@ def __init__(
             )
         self.s3_data_input_path = s3_data_input_path
         self.s3_output_path = s3_output_path
+        self.s3_analysis_config_output_path = s3_analysis_config_output_path
         self.s3_data_distribution_type = s3_data_distribution_type
         self.s3_compression_type = s3_compression_type
         self.label = label
@@ -473,7 +478,7 @@ def _run(
                 json.dump(analysis_config, f)
             s3_analysis_config_file = _upload_analysis_config(
                 analysis_config_file,
-                data_config.s3_output_path,
+                data_config.s3_analysis_config_output_path or data_config.s3_output_path,
                 self.sagemaker_session,
                 kms_key,
             )

tests/unit/test_clarify.py

@@ -24,7 +24,7 @@
     ModelPredictedLabelConfig,
     SHAPConfig,
 )
-from sagemaker import image_uris
+from sagemaker import image_uris, Processor
 
 JOB_NAME_PREFIX = "my-prefix"
 TIMESTAMP = "2021-06-17-22-29-54-685"
@@ -499,6 +499,59 @@ def test_post_training_bias(
         )
 
 
+@patch.object(Processor, "run")
+def test_run_on_s3_analysis_config_file(
+    processor_run, sagemaker_session, clarify_processor, data_config
+):
+    analysis_config = {
+        "methods": {"post_training_bias": {"methods": "all"}},
+    }
+    with patch("sagemaker.clarify._upload_analysis_config", return_value=None) as mock_method:
+        clarify_processor._run(
+            data_config,
+            analysis_config,
+            True,
+            True,
+            "test",
+            None,
+            {"ExperimentName": "AnExperiment"},
+        )
+        analysis_config_file = mock_method.call_args[0][0]
+        mock_method.assert_called_with(
+            analysis_config_file, data_config.s3_output_path, sagemaker_session, None
+        )
+
+        data_config_with_analysis_config_output = DataConfig(
+            s3_data_input_path="s3://input/train.csv",
+            s3_output_path="s3://output/analysis_test_result",
+            s3_analysis_config_output_path="s3://analysis_config_output",
+            label="Label",
+            headers=[
+                "Label",
+                "F1",
+                "F2",
+                "F3",
+            ],
+            dataset_type="text/csv",
+        )
+        clarify_processor._run(
+            data_config_with_analysis_config_output,
+            analysis_config,
+            True,
+            True,
+            "test",
+            None,
+            {"ExperimentName": "AnExperiment"},
+        )
+        analysis_config_file = mock_method.call_args[0][0]
+        mock_method.assert_called_with(
+            analysis_config_file,
+            data_config_with_analysis_config_output.s3_analysis_config_output_path,
+            sagemaker_session,
+            None,
+        )
+
+
 def _run_test_shap(
     name_from_base,
     clarify_processor,