ClearML/clearml
1
0
Fork 0
mirror of https://github.com/clearml/clearml synced 2025-03-03 10:42:00 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add CLEARML_AUTH_TOKEN support for directly passing a token (no need for user/pass)

Browse Source
This commit is contained in:
allegroai 2021-05-08 23:57:50 +03:00
parent 05139411b7
commit 8c4fa4373a
2 changed files with 45 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
clearml/backend_api/session/defs.py
View File

@ -7,6 +7,7 @@ ENV_WEB_HOST = EnvEntry("CLEARML_WEB_HOST", "TRAINS_WEB_HOST")
ENV_FILES_HOST = EnvEntry("CLEARML_FILES_HOST", "TRAINS_FILES_HOST")
ENV_ACCESS_KEY = EnvEntry("CLEARML_API_ACCESS_KEY", "TRAINS_API_ACCESS_KEY")
ENV_SECRET_KEY = EnvEntry("CLEARML_API_SECRET_KEY", "TRAINS_API_SECRET_KEY")
ENV_AUTH_TOKEN = EnvEntry("CLEARML_AUTH_TOKEN")
ENV_VERBOSE = EnvEntry("CLEARML_API_VERBOSE", "TRAINS_API_VERBOSE", type=bool, default=False)
ENV_HOST_VERIFY_CERT = EnvEntry("CLEARML_API_HOST_VERIFY_CERT", "TRAINS_API_HOST_VERIFY_CERT",
type=bool, default=True)

66
clearml/backend_api/session/session.py
View File

@ -11,8 +11,9 @@ from requests.auth import HTTPBasicAuth
from six.moves.urllib.parse import urlparse, urlunparse
from .callresult import CallResult
from .defs import ENV_VERBOSE, ENV_HOST, ENV_ACCESS_KEY, ENV_SECRET_KEY, ENV_WEB_HOST, \
ENV_FILES_HOST, ENV_OFFLINE_MODE, ENV_TRAINS_NO_DEFAULT_SERVER
from .defs import (
ENV_VERBOSE, ENV_HOST, ENV_ACCESS_KEY, ENV_SECRET_KEY, ENV_WEB_HOST,
ENV_FILES_HOST, ENV_OFFLINE_MODE, ENV_TRAINS_NO_DEFAULT_SERVER, ENV_AUTH_TOKEN, )
from .request import Request, BatchRequest # noqa: F401
from .token_manager import TokenManager
from ..config import load
@ -115,29 +116,38 @@ class Session(TokenManager):
"auth.token_expiration_threshold_sec", 60
)
self._verbose = verbose if verbose is not None else ENV_VERBOSE.get()
self._logger = logger
self.__auth_token = None
if ENV_AUTH_TOKEN.get():
self.__access_key = self.__secret_key = None
self.__auth_token = ENV_AUTH_TOKEN.get()
# if we use a token we override make sure we are at least 3600 seconds (1 hour)
# away from the token expiration date, ask for a new one.
token_expiration_threshold_sec = max(token_expiration_threshold_sec, 3600)
else:
self.__access_key = api_key or ENV_ACCESS_KEY.get(
default=(self.config.get("api.credentials.access_key", None) or self.default_key)
)
if not self.access_key:
raise ValueError(
"Missing access_key. Please set in configuration file or pass in session init."
)
self.__secret_key = secret_key or ENV_SECRET_KEY.get(
default=(self.config.get("api.credentials.secret_key", None) or self.default_secret)
)
if not self.secret_key:
raise ValueError(
"Missing secret_key. Please set in configuration file or pass in session init."
)
# init the token manager
super(Session, self).__init__(
token_expiration_threshold_sec=token_expiration_threshold_sec, **kwargs
)
self._verbose = verbose if verbose is not None else ENV_VERBOSE.get()
self._logger = logger
self.__access_key = api_key or ENV_ACCESS_KEY.get(
default=(self.config.get("api.credentials.access_key", None) or self.default_key)
)
if not self.access_key:
raise ValueError(
"Missing access_key. Please set in configuration file or pass in session init."
)
self.__secret_key = secret_key or ENV_SECRET_KEY.get(
default=(self.config.get("api.credentials.secret_key", None) or self.default_secret)
)
if not self.secret_key:
raise ValueError(
"Missing secret_key. Please set in configuration file or pass in session init."
)
host = host or self.get_api_server_host(config=self.config)
if not host:
raise ValueError("host is required in init or config")
@ -611,7 +621,13 @@ class Session(TokenManager):
)
)
auth = HTTPBasicAuth(self.access_key, self.secret_key)
headers = None
# use token only once (the second time the token is already built into the http session)
if self.__auth_token:
headers = dict(Authorization="Bearer {}".format(self.__auth_token))
self.__auth_token = None
auth = HTTPBasicAuth(self.access_key, self.secret_key) if self.access_key and self.secret_key else None
res = None
try:
data = {"expiration_sec": exp} if exp else {}
@ -620,6 +636,7 @@ class Session(TokenManager):
action="login",
auth=auth,
json=data,
headers=headers,
refresh_token_if_unauthorized=False,
)
try:
@ -635,6 +652,11 @@ class Session(TokenManager):
)
if verbose:
self._logger.info("Received new token")
# make sure we keep the token updated on the OS environment, so that child processes will have access.
if ENV_AUTH_TOKEN.get():
ENV_AUTH_TOKEN.set(resp["data"]["token"])
return resp["data"]["token"]
except LoginError:
six.reraise(*sys.exc_info())