diff --git a/trains_agent/commands/worker.py b/trains_agent/commands/worker.py index be8108d..0b79928 100644 --- a/trains_agent/commands/worker.py +++ b/trains_agent/commands/worker.py @@ -17,7 +17,7 @@ from datetime import datetime from distutils.spawn import find_executable from functools import partial from itertools import chain -from tempfile import gettempdir, mkdtemp +from tempfile import mkdtemp from time import sleep, time from typing import Text, Optional, Any, Tuple @@ -28,9 +28,6 @@ from trains_agent.backend_api.services import queues as queues_api from trains_agent.backend_api.services import tasks as tasks_api from pathlib2 import Path from pyhocon import ConfigTree, ConfigFactory -from requests import Session as HTTPSession -from requests.adapters import HTTPAdapter -from urllib3.util import Retry from six.moves.urllib.parse import quote from trains_agent.helper.check_update import start_check_update_daemon @@ -40,7 +37,9 @@ from trains_agent.definitions import ( ENVIRONMENT_SDK_PARAMS, INVALID_WORKER_ID, PROGRAM_NAME, - DEFAULT_VENV_UPDATE_URL) + DEFAULT_VENV_UPDATE_URL, + ENV_TASK_EXECUTE_AS_USER +) from trains_agent.definitions import WORKING_REPOSITORY_DIR, PIP_EXTRA_INDICES from trains_agent.errors import APIError, CommandFailedError, Sigterm from trains_agent.helper.base import ( @@ -1178,6 +1177,16 @@ class Worker(ServiceCommandSection): if python_path: os.environ['PYTHONPATH'] = python_path + # check if we want to run as another user, only supported on linux + if os.environ.get(ENV_TASK_EXECUTE_AS_USER, None) and is_linux_platform(): + command, script_dir = self._run_as_user_patch( + command, script_dir, venv_folder, + self._session.config.get('sdk.storage.cache.default_base_dir'), + os.environ.get(ENV_TASK_EXECUTE_AS_USER)) + use_execv = False + else: + use_execv = is_linux_platform() and not isinstance(self.package_api, (PoetryAPI, CondaAPI)) + print("Starting Task Execution:\n".format(task_id)) exit_code = -1 try: diff --git a/trains_agent/definitions.py b/trains_agent/definitions.py index 8456167..20b1f95 100644 --- a/trains_agent/definitions.py +++ b/trains_agent/definitions.py @@ -114,6 +114,7 @@ DEFAULT_VCS_CACHE = normalize_path(CONFIG_DIR, "vcs-cache") PIP_EXTRA_INDICES = [ ] DEFAULT_PIP_DOWNLOAD_CACHE = normalize_path(CONFIG_DIR, "pip-download-cache") +ENV_TASK_EXECUTE_AS_USER = 'TRAINS_AGENT_EXEC_USER' class FileBuffering(IntEnum): diff --git a/trains_agent/session.py b/trains_agent/session.py index ada691f..a6d1aba 100644 --- a/trains_agent/session.py +++ b/trains_agent/session.py @@ -15,7 +15,7 @@ from pyhocon import ConfigFactory, HOCONConverter, ConfigTree from trains_agent.backend_api.session import Session as _Session, Request from trains_agent.backend_api.session.client import APIClient from trains_agent.backend_config.defs import LOCAL_CONFIG_FILE_OVERRIDE_VAR, LOCAL_CONFIG_FILES -from trains_agent.definitions import ENVIRONMENT_CONFIG +from trains_agent.definitions import ENVIRONMENT_CONFIG, ENV_TASK_EXECUTE_AS_USER from trains_agent.errors import APIError from trains_agent.helper.base import HOCONEncoder from trains_agent.helper.process import Argv @@ -167,6 +167,10 @@ class Session(_Session): 'agent.docker_pip_cache', 'agent.docker_apt_cache') singleton_folders = ('agent.venvs_dir', 'agent.vcs_cache.path',) + if os.environ.get(ENV_TASK_EXECUTE_AS_USER): + folder_keys = tuple(list(folder_keys) + ['sdk.storage.cache.default_base_dir']) + singleton_folders = tuple(list(singleton_folders) + ['sdk.storage.cache.default_base_dir']) + for key in folder_keys: folder_key = ConfigValue(self.config, key) if not folder_key.get():