Add TRAINS_AGENT_EXEC_USER support for multiple daemon instances

This commit is contained in:
allegroai 2020-03-05 12:46:53 +02:00
parent 81008ee00e
commit 8f95eecf2e
3 changed files with 20 additions and 6 deletions

View File

@ -17,7 +17,7 @@ from datetime import datetime
from distutils.spawn import find_executable from distutils.spawn import find_executable
from functools import partial from functools import partial
from itertools import chain from itertools import chain
from tempfile import gettempdir, mkdtemp from tempfile import mkdtemp
from time import sleep, time from time import sleep, time
from typing import Text, Optional, Any, Tuple 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 trains_agent.backend_api.services import tasks as tasks_api
from pathlib2 import Path from pathlib2 import Path
from pyhocon import ConfigTree, ConfigFactory 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 six.moves.urllib.parse import quote
from trains_agent.helper.check_update import start_check_update_daemon from trains_agent.helper.check_update import start_check_update_daemon
@ -40,7 +37,9 @@ from trains_agent.definitions import (
ENVIRONMENT_SDK_PARAMS, ENVIRONMENT_SDK_PARAMS,
INVALID_WORKER_ID, INVALID_WORKER_ID,
PROGRAM_NAME, 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.definitions import WORKING_REPOSITORY_DIR, PIP_EXTRA_INDICES
from trains_agent.errors import APIError, CommandFailedError, Sigterm from trains_agent.errors import APIError, CommandFailedError, Sigterm
from trains_agent.helper.base import ( from trains_agent.helper.base import (
@ -1178,6 +1177,16 @@ class Worker(ServiceCommandSection):
if python_path: if python_path:
os.environ['PYTHONPATH'] = 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)) print("Starting Task Execution:\n".format(task_id))
exit_code = -1 exit_code = -1
try: try:

View File

@ -114,6 +114,7 @@ DEFAULT_VCS_CACHE = normalize_path(CONFIG_DIR, "vcs-cache")
PIP_EXTRA_INDICES = [ PIP_EXTRA_INDICES = [
] ]
DEFAULT_PIP_DOWNLOAD_CACHE = normalize_path(CONFIG_DIR, "pip-download-cache") DEFAULT_PIP_DOWNLOAD_CACHE = normalize_path(CONFIG_DIR, "pip-download-cache")
ENV_TASK_EXECUTE_AS_USER = 'TRAINS_AGENT_EXEC_USER'
class FileBuffering(IntEnum): class FileBuffering(IntEnum):

View File

@ -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 import Session as _Session, Request
from trains_agent.backend_api.session.client import APIClient 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.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.errors import APIError
from trains_agent.helper.base import HOCONEncoder from trains_agent.helper.base import HOCONEncoder
from trains_agent.helper.process import Argv from trains_agent.helper.process import Argv
@ -167,6 +167,10 @@ class Session(_Session):
'agent.docker_pip_cache', 'agent.docker_apt_cache') 'agent.docker_pip_cache', 'agent.docker_apt_cache')
singleton_folders = ('agent.venvs_dir', 'agent.vcs_cache.path',) 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: for key in folder_keys:
folder_key = ConfigValue(self.config, key) folder_key = ConfigValue(self.config, key)
if not folder_key.get(): if not folder_key.get():