Change: Allow extra_args to be passed to uploader

Author: martinRenou

src/sagemaker/experiments/_helper.py

@@ -59,11 +59,15 @@ def __init__(
         self.artifact_prefix = artifact_prefix
         self._s3_client = self.sagemaker_session.boto_session.client("s3")
 
-    def upload_artifact(self, file_path):
+    def upload_artifact(self, file_path, extra_args=None):
         """Upload an artifact file to S3.
 
         Args:
             file_path (str): the file path of the artifact
+            extra_args (dict): Optional extra arguments that may be passed to the upload operation.
+                Similar to ExtraArgs parameter in S3 upload_file function. Please refer to the
+                ExtraArgs parameter documentation here:
+                https://boto3.amazonaws.com/v1/documentation/api/latest/guide/s3-uploading-files.html#the-extraargs-parameter
 
         Returns:
             (str, str): The s3 URI of the uploaded file and the etag of the file.
@@ -91,7 +95,12 @@ def upload_artifact(self, file_path):
         artifact_s3_key = "{}/{}/{}".format(
             self.artifact_prefix, self.trial_component_name, artifact_name
         )
-        self._s3_client.upload_file(file_path, self.artifact_bucket, artifact_s3_key)
+        self._s3_client.upload_file(
+            file_path,
+            self.artifact_bucket,
+            artifact_s3_key,
+            ExtraArgs=extra_args,
+        )
         etag = self._try_get_etag(artifact_s3_key)
         return "s3://{}/{}".format(self.artifact_bucket, artifact_s3_key), etag
 

src/sagemaker/experiments/run.py

@@ -503,7 +503,8 @@ def log_file(
         file_path: str,
         name: Optional[str] = None,
         media_type: Optional[str] = None,
-        is_output: bool = True,
+        is_output: Optional[bool] = True,
+        extra_args: Optional[dict] = None,
     ):
         """Upload a file to s3 and store it as an input/output artifact in this run.
 
@@ -516,11 +517,15 @@ def log_file(
             is_output (bool): Determines direction of association to the
                 run. Defaults to True (output artifact).
                 If set to False then represented as input association.
+            extra_args (dict): Optional extra arguments that may be passed to the upload operation.
+                Similar to ExtraArgs parameter in S3 upload_file function. Please refer to the
+                ExtraArgs parameter documentation here:
+                https://boto3.amazonaws.com/v1/documentation/api/latest/guide/s3-uploading-files.html#the-extraargs-parameter
         """
         self._verify_trial_component_artifacts_length(is_output)
         media_type = media_type or guess_media_type(file_path)
         name = name or resolve_artifact_name(file_path)
-        s3_uri, _ = self._artifact_uploader.upload_artifact(file_path)
+        s3_uri, _ = self._artifact_uploader.upload_artifact(file_path, extra_args=extra_args)
         if is_output:
             self._trial_component.output_artifacts[name] = TrialComponentArtifact(
                 value=s3_uri, media_type=media_type