ClearML/clearml
1
0
Fork 0
mirror of https://github.com/clearml/clearml synced 2025-06-26 18:16:07 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Support AWS profile when specifying bucket storage credentials

Browse Source
This commit is contained in:
allegroai
2024-02-04 19:31:30 +02:00
parent 122b308646
commit 233f94f741
3 changed files with 48 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

38
clearml/backend_config/bucket_config.py
View File

@@ -34,9 +34,18 @@ class S3BucketConfig(object):
verify = attrib(type=bool, default=None) verify = attrib(type=bool, default=None)
use_credentials_chain = attrib(type=bool, default=False) use_credentials_chain = attrib(type=bool, default=False)
extra_args = attrib(type=dict, default=None) extra_args = attrib(type=dict, default=None)
profile = attrib(type=str, default="")
def update( def update(
self, key, secret, multipart=True, region=None, use_credentials_chain=False, token="", extra_args=None self,
key="",
secret="",
multipart=True,
region=None,
use_credentials_chain=False,
token="",
extra_args=None,
profile=""
): ):
self.key = key self.key = key
self.secret = secret self.secret = secret
@@ -45,6 +54,7 @@ class S3BucketConfig(object):
self.region = region self.region = region
self.use_credentials_chain = use_credentials_chain self.use_credentials_chain = use_credentials_chain
self.extra_args = extra_args self.extra_args = extra_args
self.profile = profile
def is_valid(self): def is_valid(self):
return (self.key and self.secret) or self.use_credentials_chain return (self.key and self.secret) or self.use_credentials_chain
@@ -107,6 +117,8 @@ class S3BucketConfigurations(BaseBucketConfigurations):
default_token="", default_token="",
default_extra_args=None, default_extra_args=None,
default_verify=None, default_verify=None,
default_profile="",
default_secure=True
): ):
super(S3BucketConfigurations, self).__init__() super(S3BucketConfigurations, self).__init__()
self._buckets = buckets if buckets else list() self._buckets = buckets if buckets else list()
@@ -118,12 +130,12 @@ class S3BucketConfigurations(BaseBucketConfigurations):
self._default_use_credentials_chain = default_use_credentials_chain self._default_use_credentials_chain = default_use_credentials_chain
self._default_extra_args = default_extra_args self._default_extra_args = default_extra_args
self._default_verify = default_verify self._default_verify = default_verify
self._default_profile = default_profile
self._default_secure = default_secure
@classmethod @classmethod
def from_config(cls, s3_configuration): def from_config(cls, s3_configuration):
config_list = S3BucketConfig.from_list( config_list = S3BucketConfig.from_list(s3_configuration.get("credentials", []))
s3_configuration.get("credentials", [])
)
default_key = s3_configuration.get("key", "") or getenv("AWS_ACCESS_KEY_ID", "") default_key = s3_configuration.get("key", "") or getenv("AWS_ACCESS_KEY_ID", "")
default_secret = s3_configuration.get("secret", "") or getenv("AWS_SECRET_ACCESS_KEY", "") default_secret = s3_configuration.get("secret", "") or getenv("AWS_SECRET_ACCESS_KEY", "")
@@ -132,11 +144,14 @@ class S3BucketConfigurations(BaseBucketConfigurations):
default_use_credentials_chain = s3_configuration.get("use_credentials_chain") or False default_use_credentials_chain = s3_configuration.get("use_credentials_chain") or False
default_extra_args = s3_configuration.get("extra_args") default_extra_args = s3_configuration.get("extra_args")
default_verify = s3_configuration.get("verify", None) default_verify = s3_configuration.get("verify", None)
default_profile = s3_configuration.get("profile", "") or getenv("AWS_PROFILE", "")
default_secure = s3_configuration.get("secure", True)
default_key = _none_to_empty_string(default_key) default_key = _none_to_empty_string(default_key).strip()
default_secret = _none_to_empty_string(default_secret) default_secret = _none_to_empty_string(default_secret).strip()
default_token = _none_to_empty_string(default_token) default_token = _none_to_empty_string(default_token).strip()
default_region = _none_to_empty_string(default_region) default_region = _none_to_empty_string(default_region).strip()
default_profile = _none_to_empty_string(default_profile).strip()
return cls( return cls(
config_list, config_list,
@@ -147,6 +162,8 @@ class S3BucketConfigurations(BaseBucketConfigurations):
default_token, default_token,
default_extra_args, default_extra_args,
default_verify, default_verify,
default_profile,
default_secure
) )
def add_config(self, bucket_config): def add_config(self, bucket_config):
@@ -178,6 +195,8 @@ class S3BucketConfigurations(BaseBucketConfigurations):
use_credentials_chain=self._default_use_credentials_chain, use_credentials_chain=self._default_use_credentials_chain,
token=self._default_token, token=self._default_token,
extra_args=self._default_extra_args, extra_args=self._default_extra_args,
profile=self._default_profile,
secure=self._default_secure
) )
def _get_prefix_from_bucket_config(self, config): def _get_prefix_from_bucket_config(self, config):
@@ -204,7 +223,6 @@ class S3BucketConfigurations(BaseBucketConfigurations):
:param uri: URI of bucket, directory or file :param uri: URI of bucket, directory or file
:return: S3BucketConfig: bucket config :return: S3BucketConfig: bucket config
""" """
def find_match(uri): def find_match(uri):
self._update_prefixes(refresh=False) self._update_prefixes(refresh=False)
uri = uri.lower() uri = uri.lower()
@@ -244,6 +262,8 @@ class S3BucketConfigurations(BaseBucketConfigurations):
host=host, host=host,
token=self._default_token, token=self._default_token,
extra_args=self._default_extra_args, extra_args=self._default_extra_args,
profile=self._default_profile,
secure=self._default_secure
) )

23
clearml/storage/helper.py
View File

@@ -461,15 +461,15 @@ class _Boto3Driver(_Driver):
) )
} }
if not cfg.use_credentials_chain: if not cfg.use_credentials_chain:
boto_kwargs["aws_access_key_id"] = cfg.key boto_kwargs["aws_access_key_id"] = cfg.key or None
boto_kwargs["aws_secret_access_key"] = cfg.secret boto_kwargs["aws_secret_access_key"] = cfg.secret or None
if cfg.token: if cfg.token:
boto_kwargs["aws_session_token"] = cfg.token boto_kwargs["aws_session_token"] = cfg.token
self.resource = boto3.resource( boto_session = boto3.Session(
"s3", profile_name=cfg.profile or None,
**boto_kwargs
) )
self.resource = boto_session.resource("s3", **boto_kwargs)
self.config = cfg self.config = cfg
bucket_name = self.name[len(cfg.host) + 1:] if cfg.host else self.name bucket_name = self.name[len(cfg.host) + 1:] if cfg.host else self.name
@@ -683,7 +683,12 @@ class _Boto3Driver(_Driver):
'time': datetime.utcnow().isoformat() 'time': datetime.utcnow().isoformat()
} }
boto_session = boto3.Session(conf.key, conf.secret, aws_session_token=conf.token) boto_session = boto3.Session(
aws_access_key_id=conf.key or None,
aws_secret_access_key=conf.secret or None,
aws_session_token=conf.token or None,
profile_name=conf.profile or None
)
endpoint = (('https://' if conf.secure else 'http://') + conf.host) if conf.host else None endpoint = (('https://' if conf.secure else 'http://') + conf.host) if conf.host else None
boto_resource = boto_session.resource('s3', region_name=conf.region or None, endpoint_url=endpoint) boto_resource = boto_session.resource('s3', region_name=conf.region or None, endpoint_url=endpoint)
bucket = boto_resource.Bucket(bucket_name) bucket = boto_resource.Bucket(bucket_name)
@@ -738,7 +743,9 @@ class _Boto3Driver(_Driver):
cls._bucket_location_failure_reported.add(conf.get_bucket_host()) cls._bucket_location_failure_reported.add(conf.get_bucket_host())
try: try:
boto_session = boto3.Session(conf.key, conf.secret, aws_session_token=conf.token) boto_session = boto3.Session(
conf.key, conf.secret, aws_session_token=conf.token, profile_name=conf.profile_name or None
)
boto_resource = boto_session.resource('s3') boto_resource = boto_session.resource('s3')
return boto_resource.meta.client.get_bucket_location(Bucket=conf.bucket)["LocationConstraint"] return boto_resource.meta.client.get_bucket_location(Bucket=conf.bucket)["LocationConstraint"]
@@ -2018,6 +2025,7 @@ class StorageHelper(object):
logger=None, logger=None,
retries=5, retries=5,
token=None, token=None,
profile=None,
**kwargs **kwargs
): ):
level = config.get("storage.log.level", None) level = config.get("storage.log.level", None)
@@ -2072,6 +2080,7 @@ class StorageHelper(object):
region=final_region, region=final_region,
use_credentials_chain=self._conf.use_credentials_chain, use_credentials_chain=self._conf.use_credentials_chain,
token=token or self._conf.token, token=token or self._conf.token,
profile=profile or self._conf.profile,
extra_args=self._conf.extra_args, extra_args=self._conf.extra_args,
) )

3
docs/clearml.conf
View File

@@ -92,6 +92,8 @@ sdk {
# Specify explicit keys # Specify explicit keys
key: "" key: ""
secret: "" secret: ""
# Specify profile
profile: ""
# Or enable credentials chain to let Boto3 pick the right credentials. # Or enable credentials chain to let Boto3 pick the right credentials.
# This includes picking credentials from environment variables, # This includes picking credentials from environment variables,
# credential file and IAM role using metadata service. # credential file and IAM role using metadata service.
@@ -103,6 +105,7 @@ sdk {
credentials: [ credentials: [
# specifies key/secret credentials to use when handling s3 urls (read or write) # specifies key/secret credentials to use when handling s3 urls (read or write)
# Note that all all fields in the global s3 config section are supported here
# { # {
# bucket: "my-bucket-name" # bucket: "my-bucket-name"
# key: "my-access-key" # key: "my-access-key"