diff --git a/trains_agent/commands/worker.py b/trains_agent/commands/worker.py index 0bfd9b0..818fc2a 100644 --- a/trains_agent/commands/worker.py +++ b/trains_agent/commands/worker.py @@ -39,8 +39,8 @@ from trains_agent.definitions import ( PROGRAM_NAME, DEFAULT_VENV_UPDATE_URL, ENV_TASK_EXECUTE_AS_USER, - ENV_K8S_HOST_MOUNT -) + ENV_K8S_HOST_MOUNT, + ENV_TASK_EXTRA_PYTHON_PATH) 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 ( @@ -63,8 +63,8 @@ from trains_agent.helper.base import ( error, get_python_path, is_linux_platform, - rm_file -) + rm_file, + add_python_path) from trains_agent.helper.console import ensure_text, print_text, decode_binary_lines from trains_agent.helper.package.base import PackageManager from trains_agent.helper.package.conda_api import CondaAPI @@ -862,7 +862,7 @@ class Worker(ServiceCommandSection): """ if not lines: return 0 - print_text("".join(lines)) + print_text("".join(lines), newline=False) # remove backspaces from the text log, they look bad. for i, l in enumerate(lines): @@ -1209,6 +1209,8 @@ class Worker(ServiceCommandSection): # Add the script CWD to the python path python_path = get_python_path(script_dir, execution.entry_point, self.package_api) \ if not self.is_conda else None + if os.environ.get(ENV_TASK_EXTRA_PYTHON_PATH): + python_path = add_python_path(python_path, os.environ.get(ENV_TASK_EXTRA_PYTHON_PATH)) if python_path: os.environ['PYTHONPATH'] = python_path @@ -2030,7 +2032,7 @@ class Worker(ServiceCommandSection): return base_cmd - def _run_as_user_patch(self, command, script_dir, venv_folder, sdk_cache_folder, user_uid): + def _run_as_user_patch(self, command, trains_conf, script_dir, venv_folder, sdk_cache_folder, user_uid): class RunasArgv(Argv): def __init__(self, *args): super(RunasArgv, self).__init__(*args) diff --git a/trains_agent/definitions.py b/trains_agent/definitions.py index 2aac359..11b8e08 100644 --- a/trains_agent/definitions.py +++ b/trains_agent/definitions.py @@ -121,6 +121,7 @@ PIP_EXTRA_INDICES = [ ] DEFAULT_PIP_DOWNLOAD_CACHE = normalize_path(CONFIG_DIR, "pip-download-cache") ENV_TASK_EXECUTE_AS_USER = 'TRAINS_AGENT_EXEC_USER' +ENV_TASK_EXTRA_PYTHON_PATH = 'TRAINS_AGENT_EXTRA_PYTHON_PATH' ENV_K8S_HOST_MOUNT = 'TRAINS_AGENT_K8S_HOST_MOUNT' diff --git a/trains_agent/helper/base.py b/trains_agent/helper/base.py index 5170685..4aaedcc 100644 --- a/trains_agent/helper/base.py +++ b/trains_agent/helper/base.py @@ -199,6 +199,20 @@ def get_python_path(script_dir, entry_point, package_api): return None +def add_python_path(base_path, extra_path): + try: + if not extra_path: + return base_path + python_path_sep = ';' if is_windows_platform() else ':' + base_path = base_path or '' + if not base_path.endswith(python_path_sep): + base_path += python_path_sep + base_path += extra_path.replace(':', python_path_sep) + except: + pass + return base_path + + class Singleton(ABCMeta): _instances = {}