def test_get_s3_file_upload_extra_args_invalid_json():
os.environ.setdefault(
"MLFLOW_S3_UPLOAD_EXTRA_ARGS", '"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": "123456"}'
)
with pytest.raises(json.decoder.JSONDecodeError, match=r".+"):
S3ArtifactRepository.get_s3_file_upload_extra_args()
def test_get_s3_file_upload_extra_args_invalid_json():
os.environ.setdefault(
'MLFLOW_S3_UPLOAD_EXTRA_ARGS',
'"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": "123456"}')
with pytest.raises(ValueError):
S3ArtifactRepository.get_s3_file_upload_extra_args()
def _upload_s3(local_model_path, bucket, prefix, region_name, s3_client):
"""
Upload dir to S3 as .tar.gz.
:param local_model_path: Local path to a dir.
:param bucket: S3 bucket where to store the data.
:param prefix: Path within the bucket.
:param region_name: The AWS region in which to upload data to S3.
:param s3_client: A boto3 client for S3.
:return: S3 path of the uploaded artifact.
"""
import boto3
sess = boto3.Session(region_name=region_name)
with TempDir() as tmp:
model_data_file = tmp.path("model.tar.gz")
_make_tarfile(model_data_file, local_model_path)
with open(model_data_file, "rb") as fobj:
key = os.path.join(prefix, "model.tar.gz")
obj = sess.resource("s3").Bucket(bucket).Object(key)
environ_extra_args = S3ArtifactRepository.get_s3_file_upload_extra_args(
)
obj.upload_fileobj(fobj, ExtraArgs=environ_extra_args)
response = s3_client.put_object_tagging(
Bucket=bucket,
Key=key,
Tagging={"TagSet": [{
"Key": "SageMaker",
"Value": "true"
}]})
_logger.info("tag response: %s", response)
return "{}/{}/{}".format(s3_client.meta.endpoint_url, bucket, key)
def test_get_s3_file_upload_extra_args():
os.environ.setdefault(
"MLFLOW_S3_UPLOAD_EXTRA_ARGS",
'{"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": "123456"}',
)
parsed_args = S3ArtifactRepository.get_s3_file_upload_extra_args()
assert parsed_args == {"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": "123456"}
def test_get_s3_file_upload_extra_args_env_var_not_present():
parsed_args = S3ArtifactRepository.get_s3_file_upload_extra_args()
assert parsed_args is None
