mlflow · artjen · Jul 10, 2024 · Jun 26, 2024 · Jun 27, 2024 · Jun 27, 2024
diff --git a/mlflow/protos/databricks_uc_registry_messages.proto b/mlflow/protos/databricks_uc_registry_messages.proto
@@ -141,6 +141,8 @@ message TemporaryCredentials {
   optional int64 expiration_time = 1;
 
   optional StorageMode storage_mode = 6;
+
+  optional EncryptionDetails encryption_details = 7;
 }
 
 // AWS temporary credentials for API authentication.
@@ -190,6 +192,26 @@ enum StorageMode {
   DEFAULT_STORAGE = 2;
 }
 
+message EncryptionDetails {
+  oneof encryption_details_type {
+    // Details for CLOUD_MANAGED_SSE_KEYS:
+    SseEncryptionDetails sse_encryption_details = 1;
+  }
+}
+
+enum SseEncryptionAlgorithm {
+  SSE_ENCRYPTION_ALGORITHM_UNSPECIFIED = 0;
+  AWS_SSE_KMS = 2;   // "aws:kms" in x-amz-server-side-encryption' header
+}
+
+message SseEncryptionDetails {
+  optional SseEncryptionAlgorithm algorithm = 1;
+
+  // Optional. The ARN of the SSE-KMS key used with the S3 location, when algorithm = "SSE-KMS".
+  // Sets the value of the 'x-amz-server-side-encryption-aws-kms-key-id' header.
+  optional string aws_kms_key_arn = 2;
+}
+
 message CreateRegisteredModelRequest {
   option (scalapb.message).extends = "com.databricks.rpc.RPC[CreateRegisteredModelResponse]";
 

diff --git a/mlflow/protos/databricks_uc_registry_messages_pb2.py b/mlflow/protos/databricks_uc_registry_messages_pb2.py
diff --git a/mlflow/store/artifact/optimized_s3_artifact_repo.py b/mlflow/store/artifact/optimized_s3_artifact_repo.py
@@ -52,6 +52,7 @@ def __init__(
         credential_refresh_def=None,
         addressing_style=None,
         s3_endpoint_url=None,
+        s3_upload_extra_args=None,
     ):
         super().__init__(artifact_uri)
         self._access_key_id = access_key_id
@@ -62,6 +63,7 @@ def __init__(
         self._s3_endpoint_url = s3_endpoint_url
         self.bucket, self.bucket_path = self.parse_s3_compliant_uri(self.artifact_uri)
         self._region_name = self._get_region_name()
+        self._s3_upload_extra_args = s3_upload_extra_args if s3_upload_extra_args else {}
 
     def _refresh_credentials(self):
         if not self._credential_refresh_def:
@@ -70,6 +72,7 @@ def _refresh_credentials(self):
         self._access_key_id = new_creds["access_key_id"]
         self._secret_access_key = new_creds["secret_access_key"]
         self._session_token = new_creds["session_token"]
+        self._s3_upload_extra_args = new_creds["s3_upload_extra_args"]
         return self._get_s3_client()
 
     def _get_region_name(self):
@@ -146,6 +149,7 @@ def get_s3_file_upload_extra_args():
 
     def _upload_file(self, s3_client, local_file, bucket, key):
         extra_args = {}
+        extra_args.update(self._s3_upload_extra_args)
         guessed_type, guessed_encoding = guess_type(local_file)
         if guessed_type is not None:
             extra_args["ContentType"] = guessed_type

diff --git a/mlflow/store/artifact/r2_artifact_repo.py b/mlflow/store/artifact/r2_artifact_repo.py
@@ -14,6 +14,7 @@ def __init__(
         secret_access_key=None,
         session_token=None,
         credential_refresh_def=None,
+        s3_upload_extra_args=None,
     ):
         # setup Cloudflare R2 backend to be endpoint_url, otherwise all s3 requests
         # will go to AWS S3 by default
@@ -31,6 +32,7 @@ def __init__(
             credential_refresh_def=credential_refresh_def,
             addressing_style="virtual",
             s3_endpoint_url=s3_endpoint_url,
+            s3_upload_extra_args=s3_upload_extra_args,
         )
 
     # Cloudflare implementation of head_bucket is not the same as AWS's, so we

diff --git a/mlflow/utils/_unity_catalog_utils.py b/mlflow/utils/_unity_catalog_utils.py
@@ -23,7 +23,10 @@
 from mlflow.protos.databricks_uc_registry_messages_pb2 import (
     RegisteredModelTag as ProtoRegisteredModelTag,
 )
-from mlflow.protos.databricks_uc_registry_messages_pb2 import TemporaryCredentials
+from mlflow.protos.databricks_uc_registry_messages_pb2 import (
+    SseEncryptionAlgorithm,
+    TemporaryCredentials,
+)
 from mlflow.store.artifact.artifact_repo import ArtifactRepository
 
 _STRING_TO_STATUS = {k: ProtoModelVersionStatus.Value(k) for k in ProtoModelVersionStatus.keys()}
@@ -155,14 +158,17 @@ def _get_artifact_repo_from_storage_info(
         from mlflow.store.artifact.optimized_s3_artifact_repo import OptimizedS3ArtifactRepository
 
         aws_creds = scoped_token.aws_temp_credentials
+        s3_upload_extra_args = _parse_aws_sse_credential(scoped_token)
 
         def aws_credential_refresh():
             new_scoped_token = base_credential_refresh_def()
             new_aws_creds = new_scoped_token.aws_temp_credentials
+            new_s3_upload_extra_args = _parse_aws_sse_credential(new_scoped_token)
             return {
                 "access_key_id": new_aws_creds.access_key_id,
                 "secret_access_key": new_aws_creds.secret_access_key,
                 "session_token": new_aws_creds.session_token,
+                "s3_upload_extra_args": new_s3_upload_extra_args,
             }
 
         return OptimizedS3ArtifactRepository(
@@ -171,6 +177,7 @@ def aws_credential_refresh():
             secret_access_key=aws_creds.secret_access_key,
             session_token=aws_creds.session_token,
             credential_refresh_def=aws_credential_refresh,
+            s3_upload_extra_args=s3_upload_extra_args,
         )
     elif credential_type == "azure_user_delegation_sas":
         from azure.core.credentials import AzureSasCredential
@@ -232,6 +239,26 @@ def r2_credential_refresh():
         )
 
 
+def _parse_aws_sse_credential(scoped_token: TemporaryCredentials):
+    encryption_details = scoped_token.encryption_details
+    if not encryption_details:
+        return {}
+
+    if encryption_details.WhichOneof("encryption_details_type") != "sse_encryption_details":
+        return {}
+
+    sse_encryption_details = encryption_details.sse_encryption_details
+
+    if sse_encryption_details.algorithm != SseEncryptionAlgorithm.AWS_SSE_KMS:
+        return {}
+
+    key_id = sse_encryption_details.aws_kms_key_arn.split("/")[-1]
+    return {
+        "ServerSideEncryption": "aws:kms",
+        "SSEKMSKeyId": key_id,
+    }
+
+
 def get_full_name_from_sc(name, spark) -> str:
     """
     Constructs the full name of a registered model using the active catalog and schema in a spark

diff --git a/tests/store/_unity_catalog/model_registry/test_unity_catalog_rest_store.py b/tests/store/_unity_catalog/model_registry/test_unity_catalog_rest_store.py
@@ -314,6 +314,7 @@ def test_create_model_version_with_langchain_dependencies(store, langchain_local
             secret_access_key=secret_access_key,
             session_token=session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_artifact_repo.log_artifacts.assert_called_once_with(local_dir=ANY, artifact_path="")
         _assert_create_model_version_endpoints_called(
@@ -377,6 +378,7 @@ def test_create_model_version_with_resources(store, langchain_local_model_dir_wi
             secret_access_key=secret_access_key,
             session_token=session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_artifact_repo.log_artifacts.assert_called_once_with(local_dir=ANY, artifact_path="")
         _assert_create_model_version_endpoints_called(
@@ -434,6 +436,7 @@ def test_create_model_version_with_langchain_no_dependencies(
             secret_access_key=secret_access_key,
             session_token=session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_artifact_repo.log_artifacts.assert_called_once_with(local_dir=ANY, artifact_path="")
         _assert_create_model_version_endpoints_called(
@@ -1019,6 +1022,7 @@ def test_create_model_version_aws(store, local_model_dir):
             secret_access_key=secret_access_key,
             session_token=session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_artifact_repo.log_artifacts.assert_called_once_with(local_dir=ANY, artifact_path="")
         _assert_create_model_version_endpoints_called(

diff --git a/tests/store/artifact/test_optimized_s3_artifact_repo.py b/tests/store/artifact/test_optimized_s3_artifact_repo.py
@@ -267,6 +267,7 @@ def credential_refresh_def():
                 "access_key_id": "my-id-2",
                 "secret_access_key": "my-key-2",
                 "session_token": "my-session-2",
+                "s3_upload_extra_args": {},
             }
 
         repo = OptimizedS3ArtifactRepository(

diff --git a/tests/store/artifact/test_unity_catalog_models_artifact_repo.py b/tests/store/artifact/test_unity_catalog_models_artifact_repo.py
@@ -142,6 +142,7 @@ def test_uc_models_artifact_repo_download_artifacts_uses_temporary_creds_aws(mon
             secret_access_key=fake_secret_access_key,
             session_token=fake_session_token,
             credential_refresh_def=ANY,
+            s3_upload_extra_args={},
         )
         mock_s3_repo.download_artifacts.assert_called_once_with("artifact_path", "dst_path")
         request_mock.assert_called_with(

diff --git a/tests/utils/test_unity_catalog_utils.py b/tests/utils/test_unity_catalog_utils.py
@@ -9,6 +9,12 @@
 )
 from mlflow.entities.model_registry.model_version_search import ModelVersionSearch
 from mlflow.entities.model_registry.registered_model_search import RegisteredModelSearch
+from mlflow.protos.databricks_uc_registry_messages_pb2 import (
+    EncryptionDetails,
+    SseEncryptionAlgorithm,
+    SseEncryptionDetails,
+    TemporaryCredentials,
+)
 from mlflow.protos.databricks_uc_registry_messages_pb2 import ModelVersion as ProtoModelVersion
 from mlflow.protos.databricks_uc_registry_messages_pb2 import (
     ModelVersionStatus as ProtoModelVersionStatus,
@@ -26,6 +32,7 @@
     RegisteredModelTag as ProtoRegisteredModelTag,
 )
 from mlflow.utils._unity_catalog_utils import (
+    _parse_aws_sse_credential,
     model_version_from_uc_proto,
     model_version_search_from_uc_proto,
     registered_model_from_uc_proto,
@@ -245,3 +252,37 @@ def test_registered_model_and_registered_model_search_equality():
     registered_model_search_2 = RegisteredModelSearch(**kwargs)
 
     assert registered_model_2 == registered_model_search_2
+
+
+@pytest.mark.parametrize(
+    ("temp_credentials", "parsed"),
+    [
+        (TemporaryCredentials(), {}),
+        (
+            TemporaryCredentials(
+                encryption_details=EncryptionDetails(
+                    sse_encryption_details=SseEncryptionDetails(
+                        algorithm=SseEncryptionAlgorithm.SSE_ENCRYPTION_ALGORITHM_UNSPECIFIED
+                    )
+                )
+            ),
+            {},
+        ),
+        (
+            TemporaryCredentials(
+                encryption_details=EncryptionDetails(
+                    sse_encryption_details=SseEncryptionDetails(
+                        algorithm=SseEncryptionAlgorithm.AWS_SSE_KMS,
+                        aws_kms_key_arn="key_id",
+                    )
+                )
+            ),
+            {
+                "ServerSideEncryption": "aws:kms",
+                "SSEKMSKeyId": "key_id",
+            },
+        ),
+    ],
+)
+def test_parse_aws_sse_credential(temp_credentials, parsed):
+    assert _parse_aws_sse_credential(temp_credentials) == parsed