change: s3_upload mode for CV jobs in clarify processing output

Author: Satish615

src/sagemaker/clarify.py

@@ -26,7 +26,7 @@
 import tempfile
 from abc import ABC, abstractmethod
 from typing import List, Union, Dict, Optional, Any
-
+from enum import Enum
 from schema import Schema, And, Use, Or, Optional as SchemaOptional, Regex
 
 from sagemaker import image_uris, s3, utils
@@ -302,6 +302,16 @@
 )
 
 
+class DatasetType(Enum):
+    """Enum to store different dataset types supported in the Analysis config file"""
+
+    TEXTCSV = "text/csv"
+    JSONLINES = "application/jsonlines"
+    JSON = "application/json"
+    PARQUET = "application/x-parquet"
+    IMAGE = "application/x-image"
+
+
 class DataConfig:
     """Config object related to configurations of the input and output dataset."""
 
@@ -1363,7 +1373,7 @@ def _run(
                 source=self._CLARIFY_OUTPUT,
                 destination=data_config.s3_output_path,
                 output_name="analysis_result",
-                s3_upload_mode="EndOfJob",
+                s3_upload_mode=ProcessingOutputHandler.get_s3_upload_mode(analysis_config),
             )
 
             return super().run(
@@ -2083,6 +2093,30 @@ def _upload_analysis_config(analysis_config_file, s3_output_path, sagemaker_sess
     )
 
 
+class ProcessingOutputHandler:
+    """Handles the parameters sent in SagemakerProcessor.Processingoutput based on the dataset
+    type in analysis_config.
+    """
+
+    class S3UploadMode(Enum):
+        """Enum values for different uplaod modes to s3 bucket"""
+
+        CONTINUOUS = "Continuous"
+        ENDOFJOB = "EndOfJob"
+
+    @classmethod
+    def get_s3_upload_mode(cls, analysis_config: Dict[str, Any]) -> str:
+        """
+        returns the s3_upload mode based on the shap_config values
+        """
+        dataset_type = analysis_config["dataset_type"]
+        return (
+            ProcessingOutputHandler.S3UploadMode.CONTINUOUS.value
+            if dataset_type == DatasetType.IMAGE
+            else ProcessingOutputHandler.S3UploadMode.ENDOFJOB.value
+        )
+
+
 def _set(value, key, dictionary):
     """Sets dictionary[key] = value if value is not None."""
     if value is not None:

tests/unit/test_clarify.py

@@ -30,6 +30,8 @@
     TextConfig,
     ImageConfig,
     _AnalysisConfigGenerator,
+    DatasetType,
+    ProcessingOutputHandler,
 )
 
 JOB_NAME_PREFIX = "my-prefix"
@@ -1714,3 +1716,15 @@ def test_invalid_analysis_config(data_config, data_bias_config, model_config):
             pre_training_methods="all",
             post_training_methods="all",
         )
+
+
+class TestProcessingOutputHandler:
+    def test_get_s3_upload_mode_image(self):
+        analysis_config = {"dataset_type": DatasetType.IMAGE}
+        s3_upload_mode = ProcessingOutputHandler.get_s3_upload_mode(analysis_config)
+        assert s3_upload_mode == ProcessingOutputHandler.S3UploadMode.CONTINUOUS.value
+
+    def test_get_s3_upload_mode_text(self):
+        analysis_config = {"dataset_type": DatasetType.TEXTCSV}
+        s3_upload_mode = ProcessingOutputHandler.get_s3_upload_mode(analysis_config)
+        assert s3_upload_mode == ProcessingOutputHandler.S3UploadMode.ENDOFJOB.value