Add sdk.development.log_os_environments to match TRAINS_LOG_ENVIRONMENT

Support logging OS environments into Environment section
This commit is contained in:
allegroai 2020-08-08 12:37:42 +03:00
parent f89ec8cbac
commit f4f53902ed
4 changed files with 48 additions and 20 deletions

View File

@ -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 {}

View File

@ -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

View File

@ -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

View File

@ -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