Add TRAINS_AGENT_EXTRA_PYTHON_PATH to allow adding additional python path for task execution (helpful when using extra untracked modules)

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)

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'
 
 

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 = {}