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

View File

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

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.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():