From f4f53902edc37cacf2c6f77b0b79aa1b9e67e039 Mon Sep 17 00:00:00 2001 From: allegroai <> Date: Sat, 8 Aug 2020 12:37:42 +0300 Subject: [PATCH] Add sdk.development.log_os_environments to match TRAINS_LOG_ENVIRONMENT Support logging OS environments into Environment section --- trains/backend_interface/task/task.py | 4 +-- trains/binding/environ_bind.py | 52 ++++++++++++++++++--------- trains/config/default/sdk.conf | 8 +++++ trains/task.py | 4 +-- 4 files changed, 48 insertions(+), 20 deletions(-) diff --git a/trains/backend_interface/task/task.py b/trains/backend_interface/task/task.py index 05e4fd60..96ba7150 100644 --- a/trains/backend_interface/task/task.py +++ b/trains/backend_interface/task/task.py @@ -1504,7 +1504,7 @@ class Task(IdObjectBase, AccessMixin, SetupUploadMixin): mutually_exclusive(config_dict=config_dict, config_text=config_text) if not Session.check_min_api_version('2.9'): - raise ValueError("Multiple configurations are not supported with the current 'trains-server', " + raise ValueError("Multiple configurations is not supported with the current 'trains-server', " "please upgrade to the latest version") if description: @@ -1527,7 +1527,7 @@ class Task(IdObjectBase, AccessMixin, SetupUploadMixin): return None if configuration name is not valid. """ if not Session.check_min_api_version('2.9'): - raise ValueError("Multiple configurations are not supported with the current 'trains-server', " + raise ValueError("Multiple configurations is not supported with the current 'trains-server', " "please upgrade to the latest version") configuration = self.data.configuration or {} diff --git a/trains/binding/environ_bind.py b/trains/binding/environ_bind.py index 3f5bdacc..b2902b7f 100644 --- a/trains/binding/environ_bind.py +++ b/trains/binding/environ_bind.py @@ -2,11 +2,12 @@ import os import six -from ..config import TASK_LOG_ENVIRONMENT, running_remotely +from ..config import TASK_LOG_ENVIRONMENT, running_remotely, config class EnvironmentBind(object): _task = None + _environment_section = 'Environment' @classmethod def update_current_task(cls, current_task): @@ -21,21 +22,36 @@ class EnvironmentBind(object): def _bind_environment(cls): if not cls._task: return - environ_log = str(TASK_LOG_ENVIRONMENT.get() or '').strip() + + # get ENVIRONMENT and put it into the OS environment + if running_remotely(): + params = cls._task.get_parameters_as_dict() + if params and cls._environment_section in params: + # put back into os: + os.environ.update(params[cls._environment_section]) + return + + environ_log = \ + str(TASK_LOG_ENVIRONMENT.get() or '').strip() or config.get('development.log_os_environments', []) + if environ_log and isinstance(environ_log, str): + environ_log = [e.strip() for e in environ_log.split(',')] + 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) + env_param = dict() + for match in (environ_log or []): + match = match.strip() + if match == '*': + env_param.update({k: os.environ.get(k) for k in os.environ + if not k.startswith('TRAINS_') and not k.startswith('ALG_')}) + elif match.endswith('*'): + match = match.strip('*') + env_param.update({k: os.environ.get(k) for k in os.environ if k.startswith(match)}) + elif match in os.environ: + env_param.update({match: os.environ.get(match)}) + # store os environments + cls._task.connect(env_param, cls._environment_section) class PatchOsFork(object): @@ -43,6 +59,7 @@ class PatchOsFork(object): @classmethod def patch_fork(cls): + # noinspection PyBroadException try: # only once if cls._original_fork: @@ -67,18 +84,21 @@ class PatchOsFork(object): task.get_logger().flush() # Hack: now make sure we setup the reporter thread + + # noinspection PyProtectedMember task._setup_reporter() # TODO: Check if the signal handler method is enough, for the time being, we have both # # if we got here patch the os._exit of our instance to call us - def _at_exit_callback(*args, **kwargs): + def _at_exit_callback(*a_args, **a_kwargs): # call at exit manually # noinspection PyProtectedMember task._at_exit() - # noinspection PyProtectedMember - return os._org_exit(*args, **kwargs) + # noinspection PyProtectedMember, PyUnresolvedReferences + return os._org_exit(*a_args, **a_kwargs) if not hasattr(os, '_org_exit'): + # noinspection PyProtectedMember, PyUnresolvedReferences os._org_exit = os._exit os._exit = _at_exit_callback diff --git a/trains/config/default/sdk.conf b/trains/config/default/sdk.conf index 51fdb1b8..17255551 100644 --- a/trains/config/default/sdk.conf +++ b/trains/config/default/sdk.conf @@ -154,6 +154,14 @@ # If this flag is true (default is false), instead of analyzing the code with Pigar, analyze with `pip freeze` detect_with_pip_freeze: false + # Log specific environment variables. OS environments are enlisted in the "Environment" section + # of the Hyper-Parameters. + # multiple selected variables are supported including the suffix '*'. + # For example: "AWS_*" will log any OS environment variable starting with 'AWS_'. + # This value can be overwritten with os environment variable TRAINS_LOG_ENVIRONMENT="[AWS_*, CUDA_VERSION]" + # Example: log_os_environments: ["AWS_*", "CUDA_VERSION"] + log_os_environments: [] + # Development mode worker worker { # Status report period in seconds diff --git a/trains/task.py b/trains/task.py index fa54cbc6..6ae6d590 100644 --- a/trains/task.py +++ b/trains/task.py @@ -943,7 +943,7 @@ class Task(_Task): name = self.__default_configuration_name if not multi_config_support and name and name != self.__default_configuration_name: - raise ValueError("Multiple configurations are not supported with the current 'trains-server', " + raise ValueError("Multiple configurations is not supported with the current 'trains-server', " "please upgrade to the latest version") for mutable_type, method in dispatch: @@ -1006,7 +1006,7 @@ class Task(_Task): name = self.__default_configuration_name if not multi_config_support and name and name != self.__default_configuration_name: - raise ValueError("Multiple configurations are not supported with the current 'trains-server', " + raise ValueError("Multiple configurations is not supported with the current 'trains-server', " "please upgrade to the latest version") # parameter dictionary