feat: Support uncompressed model upload (#3862)

Co-authored-by: AwesomeLeo Zhang <[email protected]>

Author: zhangluoma

src/sagemaker/estimator.py

@@ -173,6 +173,7 @@ def __init__(
         training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None,
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
+        disable_output_compression: bool = False,
         **kwargs,
     ):
         """Initialize an ``EstimatorBase`` instance.
@@ -531,6 +532,8 @@ def __init__(
                 the default train processing instructions.
             container_arguments (List[str]): Optional. The arguments for a container used to run
                 a training job.
+            disable_output_compression (bool): Optional. When set to true, Model is uploaded
+                to Amazon S3 without compression after training finishes.
         """
         instance_count = renamed_kwargs(
             "train_instance_count", "instance_count", instance_count, kwargs
@@ -712,7 +715,7 @@ def __init__(
         self.profiler_rule_configs = None
         self.profiler_rules = None
         self.debugger_rules = None
-
+        self.disable_output_compression = disable_output_compression
         validate_source_code_input_against_pipeline_variables(
             entry_point=entry_point,
             source_dir=source_dir,
@@ -2507,6 +2510,7 @@ def __init__(
         training_repository_credentials_provider_arn: Optional[Union[str, PipelineVariable]] = None,
         container_entry_point: Optional[List[str]] = None,
         container_arguments: Optional[List[str]] = None,
+        disable_output_compression: bool = False,
         **kwargs,
     ):
         """Initialize an ``Estimator`` instance.
@@ -2864,6 +2868,8 @@ def __init__(
                 the default train processing instructions.
             container_arguments (List[str]): Optional. The arguments for a container used to run
                 a training job.
+            disable_output_compression (bool): Optional. When set to true, Model is uploaded
+                to Amazon S3 without compression after training finishes.
         """
         self.image_uri = image_uri
         self._hyperparameters = hyperparameters.copy() if hyperparameters else {}
@@ -2913,6 +2919,7 @@ def __init__(
             training_repository_credentials_provider_arn=training_repository_credentials_provider_arn,  # noqa: E501 # pylint: disable=line-too-long
             container_entry_point=container_entry_point,
             container_arguments=container_arguments,
+            disable_output_compression=disable_output_compression,
             **kwargs,
         )
 

src/sagemaker/job.py

@@ -71,7 +71,11 @@ def _load_config(inputs, estimator, expand_role=True, validate_uri=True):
             if (expand_role and not is_pipeline_variable(estimator.role))
             else estimator.role
         )
-        output_config = _Job._prepare_output_config(estimator.output_path, estimator.output_kms_key)
+        output_config = _Job._prepare_output_config(
+            estimator.output_path,
+            estimator.output_kms_key,
+            disable_output_compression=estimator.disable_output_compression,
+        )
         resource_config = _Job._prepare_resource_config(
             estimator.instance_count,
             estimator.instance_type,
@@ -273,11 +277,13 @@ def _format_record_set_list_input(inputs):
         return input_dict
 
     @staticmethod
-    def _prepare_output_config(s3_path, kms_key_id):
+    def _prepare_output_config(s3_path, kms_key_id, disable_output_compression=False):
         """Placeholder docstring"""
         config = {"S3OutputPath": s3_path}
         if kms_key_id is not None:
             config["KmsKeyId"] = kms_key_id
+        if disable_output_compression:
+            config["CompressionType"] = "NONE"
         return config
 
     @staticmethod

tests/unit/sagemaker/jumpstart/estimator/test_estimator.py

@@ -522,7 +522,7 @@ def test_jumpstart_estimator_kwargs_match_parent_class(self):
         and cut a ticket sev-3 to JumpStart team: AWS > SageMaker > JumpStart"""
 
         init_args_to_skip: Set[str] = set(
-            ["container_entry_point", "container_arguments", "kwargs"]
+            ["container_entry_point", "container_arguments", "disable_output_compression", "kwargs"]
         )
         fit_args_to_skip: Set[str] = set()
         deploy_args_to_skip: Set[str] = set(["kwargs"])

tests/unit/test_estimator.py

@@ -1819,6 +1819,33 @@ def test_get_instance_type_gpu(sagemaker_session):
     assert "ml.p3.16xlarge" == estimator._get_instance_type()
 
 
+def test_estimator_with_output_compression_disabled(sagemaker_session):
+    estimator = Estimator(
+        image_uri="some-image",
+        role="some_image",
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        sagemaker_session=sagemaker_session,
+        base_job_name="base_job_name",
+        disable_output_compression=True,
+    )
+
+    assert estimator.disable_output_compression
+
+
+def test_estimator_with_output_compression_as_default(sagemaker_session):
+    estimator = Estimator(
+        image_uri="some-image",
+        role="some_image",
+        instance_count=INSTANCE_COUNT,
+        instance_type=INSTANCE_TYPE,
+        sagemaker_session=sagemaker_session,
+        base_job_name="base_job_name",
+    )
+
+    assert not estimator.disable_output_compression
+
+
 def test_get_instance_type_cpu(sagemaker_session):
     estimator = Estimator(
         image_uri="some-image",

tests/unit/test_job.py

@@ -38,6 +38,7 @@
 REGION = "us-west-2"
 IMAGE_NAME = "fakeimage"
 SCRIPT_NAME = "script.py"
+NONE_COMPRESSION_TYPE = "NONE"
 JOB_NAME = "fakejob"
 VOLUME_KMS_KEY = "volkmskey"
 MODEL_CHANNEL_NAME = "testModelChannel"
@@ -144,6 +145,23 @@ def test_load_config(estimator):
     assert config["role"] == ROLE
     assert config["output_config"]["S3OutputPath"] == S3_OUTPUT_PATH
     assert "KmsKeyId" not in config["output_config"]
+    assert "CompressionType" not in config["output_config"]
+    assert config["resource_config"]["InstanceCount"] == INSTANCE_COUNT
+    assert config["resource_config"]["InstanceType"] == INSTANCE_TYPE
+    assert config["resource_config"]["VolumeSizeInGB"] == VOLUME_SIZE
+    assert config["stop_condition"]["MaxRuntimeInSeconds"] == MAX_RUNTIME
+
+
+def test_load_config_with_output_compression_disabled(estimator):
+    inputs = TrainingInput(BUCKET_NAME)
+    estimator.disable_output_compression = True
+    config = _Job._load_config(inputs, estimator)
+
+    assert config["input_config"][0]["DataSource"]["S3DataSource"]["S3Uri"] == BUCKET_NAME
+    assert config["role"] == ROLE
+    assert config["output_config"]["S3OutputPath"] == S3_OUTPUT_PATH
+    assert "KmsKeyId" not in config["output_config"]
+    assert config["output_config"]["CompressionType"] == NONE_COMPRESSION_TYPE
     assert config["resource_config"]["InstanceCount"] == INSTANCE_COUNT
     assert config["resource_config"]["InstanceType"] == INSTANCE_TYPE
     assert config["resource_config"]["VolumeSizeInGB"] == VOLUME_SIZE