diff --git a/trains/binding/environ_bind.py b/trains/binding/environ_bind.py new file mode 100644 index 00000000..1238a406 --- /dev/null +++ b/trains/binding/environ_bind.py @@ -0,0 +1,36 @@ +import os + +from ..config import TASK_LOG_ENVIRONMENT, running_remotely + + +class EnvironmentBind(object): + _task = None + + @classmethod + def update_current_task(cls, current_task): + cls._task = current_task + # noinspection PyBroadException + try: + cls._bind_environment() + except Exception: + pass + + @classmethod + def _bind_environment(cls): + if not cls._task: + return + environ_log = str(TASK_LOG_ENVIRONMENT.get() or '').strip() + if not environ_log: + return + + if environ_log == '*': + env_param = {k: os.environ.get(k) for k in os.environ + if not k.startswith('TRAINS_') and not k.startswith('ALG_')} + else: + environ_log = [e.strip() for e in environ_log.split(',')] + env_param = {k: os.environ.get(k) for k in os.environ if k in environ_log} + + env_param = cls._task.connect(env_param) + if running_remotely(): + # put back into os: + os.environ.update(env_param) diff --git a/trains/config/defs.py b/trains/config/defs.py index c18553a7..b13f94bd 100644 --- a/trains/config/defs.py +++ b/trains/config/defs.py @@ -15,6 +15,7 @@ PROC_MASTER_ID_ENV_VAR = EnvEntry("TRAINS_PROC_MASTER_ID", "ALG_PROC_MASTER_ID", LOG_STDERR_REDIRECT_LEVEL = EnvEntry("TRAINS_LOG_STDERR_REDIRECT_LEVEL", "ALG_LOG_STDERR_REDIRECT_LEVEL") DEV_WORKER_NAME = EnvEntry("TRAINS_WORKER_NAME", "ALG_WORKER_NAME") DEV_TASK_NO_REUSE = EnvEntry("TRAINS_TASK_NO_REUSE", "ALG_TASK_NO_REUSE", type=bool) +TASK_LOG_ENVIRONMENT = EnvEntry("TRAINS_LOG_ENVIRONMENT", "ALG_LOG_ENVIRONMENT", type=str) LOG_LEVEL_ENV_VAR = EnvEntry("TRAINS_LOG_LEVEL", "ALG_LOG_LEVEL", converter=or_(int, str)) @@ -26,7 +27,3 @@ VCS_BRANCH = EnvEntry("TRAINS_VCS_BRANCH", "ALG_VCS_BRANCH") VCS_ROOT = EnvEntry("TRAINS_VCS_ROOT", "ALG_VCS_ROOT") VCS_STATUS = EnvEntry("TRAINS_VCS_STATUS", "ALG_VCS_STATUS", converter=base64_to_text) VCS_DIFF = EnvEntry("TRAINS_VCS_DIFF", "ALG_VCS_DIFF", converter=base64_to_text) - -# User credentials -API_ACCESS_KEY = EnvEntry("TRAINS_API_ACCESS_KEY", "ALG_API_ACCESS_KEY", help="API Access Key") -API_SECRET_KEY = EnvEntry("TRAINS_API_SECRET_KEY", "ALG_API_SECRET_KEY", help="API Secret Key")