Add agent.enable_task_env set the OS environment based on the Environment section of the Task.

This commit is contained in:
allegroai 2021-02-28 19:47:44 +02:00
parent f019905720
commit 9dfa1294e2
3 changed files with 26 additions and 0 deletions

View File

@ -149,6 +149,9 @@
# arguments: ["--ipc=host", ] # arguments: ["--ipc=host", ]
} }
# set the OS environments based on the Task's Environment section before launching the Task process.
enable_task_env: false
# set the initial bash script to execute at the startup of any docker. # set the initial bash script to execute at the startup of any docker.
# all lines will be executed regardless of their exit code. # all lines will be executed regardless of their exit code.
# {python_single_digit} is translated to 'python3' or 'python2' according to requested python version # {python_single_digit} is translated to 'python3' or 'python2' according to requested python version

View File

@ -1823,6 +1823,12 @@ class Worker(ServiceCommandSection):
if repo_info: if repo_info:
self._update_commit_id(current_task.id, execution, repo_info) self._update_commit_id(current_task.id, execution, repo_info)
# get Task Environments and update the process
if self._session.config.get('agent.enable_task_env', None):
hyper_params = self._get_task_os_env(current_task)
if hyper_params:
os.environ.update(hyper_params)
# Add the script CWD to the python path # Add the script CWD to the python path
python_path = get_python_path(script_dir, execution.entry_point, self.package_api, is_conda_env=self.is_conda) python_path = get_python_path(script_dir, execution.entry_point, self.package_api, is_conda_env=self.is_conda)
if ENV_TASK_EXTRA_PYTHON_PATH.get(): if ENV_TASK_EXTRA_PYTHON_PATH.get():
@ -1900,6 +1906,20 @@ class Worker(ServiceCommandSection):
return 1 if exit_code is None else exit_code return 1 if exit_code is None else exit_code
def _get_task_os_env(self, current_task):
if not self._session.check_min_api_version('2.9'):
return None
# noinspection PyBroadException
try:
hyper_params = self._session.get(
service="tasks", action="get_hyper_params", tasks=[current_task.id])
hyper_params = {
str(p['name']): str(p['value'])
for p in hyper_params['params'][0]['hyperparams'] if p['section'] == 'Environment'}
return hyper_params
except Exception:
return None
def set_docker_variables(self, docker): def set_docker_variables(self, docker):
temp_config, docker_image_func = self.get_docker_config_cmd(docker) temp_config, docker_image_func = self.get_docker_config_cmd(docker)
self.dump_config(self.temp_config_path, config=temp_config) self.dump_config(self.temp_config_path, config=temp_config)

View File

@ -147,6 +147,9 @@ agent {
# arguments: ["--ipc=host"] # arguments: ["--ipc=host"]
} }
# set the OS environments based on the Task's Environment section before launching the Task process.
enable_task_env: false
# CUDA versions used for Conda setup & solving PyTorch wheel packages # CUDA versions used for Conda setup & solving PyTorch wheel packages
# it Should be detected automatically. Override with os environment CUDA_VERSION / CUDNN_VERSION # it Should be detected automatically. Override with os environment CUDA_VERSION / CUDNN_VERSION
# cuda_version: 10.1 # cuda_version: 10.1