承担IAM角色以将MLFlow Artifact存储在另一个帐户的S3存储桶上

Question

我必须将MLFlow工件（使用Databricks Unified Analytics）保存到S3存储桶中，并使用KMS密钥在服务端进行加密。

我的实例进入一个AWS账户A，我的S3存储桶，我的KMS密钥进入账户B。我无法将KMS密钥进入账户A。

出于安全考虑，我不想使用DBFS来挂载S3存储桶（存储桶可能包含敏感数据，并且我不想在用户之间共享此信息）。

我必须承担IAM角色才能访问存储桶，就像我通过s3a（使用spark.hadoop.fs.s3a.credentialsType和spark.hadoop.fs.s3a.stsAssumeRole.arn参数访问存储桶一样）。

当我使用s3创建实验并尝试记录这样的模型时：

import mlflow
import mlflow.sklearn
id_exp = mlflow.create_experiment("/Users/first.last@company.org/Experiment",'s3://s3-bucket-name/')
with mlflow.start_run(experiment_id=id_exp):
  clf_mlf = tree.DecisionTreeClassifier()
  clf_mlf = clf_mlf.fit(X_train, y_train)
  y_pred = clf_mlf.predict(X_test)
  mlflow.sklearn.log_model(clf_mlf, "model", serialization_format='pickle')

我遇到此错误：

S3UploadFailedError: Failed to upload /tmp/tmp2yl2olhi/model/conda.yaml to s3-bucket-name//05c17a33a33d46a5ad3cc811a9faf35a/artifacts/model/conda.yaml: An error occurred (KMS.NotFoundException) when calling the PutObject operation: Key 'arn:aws:kms:eu-central-1:account_a_id:key/key_id' does not exist

在访问S3之前如何告诉MLFlow扮演角色？