From b3e8be6296c10b0cdaf01e02f6fd8738273370ca Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Fri, 7 Jan 2022 15:11:59 +0200
Subject: [PATCH] Add agent.force_git_root_python_path configuration setting to
 force adding the git repository root folder to the PYTHONPATH (if set working
 directory is not added to the PYHTONPATH)

---
 clearml_agent/backend_api/config/default/agent.conf |  4 ++++
 clearml_agent/commands/worker.py                    |  5 ++++-
 clearml_agent/helper/base.py                        | 11 +++++++----
 3 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/clearml_agent/backend_api/config/default/agent.conf b/clearml_agent/backend_api/config/default/agent.conf
index a2d5aab..9d07cdc 100644
--- a/clearml_agent/backend_api/config/default/agent.conf
+++ b/clearml_agent/backend_api/config/default/agent.conf
@@ -30,6 +30,10 @@
     # specific python version and the system supports multiple python the agent will use the requested python version)
     # ignore_requested_python_version: true
 
+    # Force the root folder of the git repository (instead of the working directory) into the PYHTONPATH
+    # default false, only the working directory will be added to the PYHTONPATH
+    # force_git_root_python_path: false
+
     # select python package manager:
     # currently supported: pip, conda and poetry
     # if "pip" or "conda" are used, the agent installs the required packages
diff --git a/clearml_agent/commands/worker.py b/clearml_agent/commands/worker.py
index 6e9aa6a..18c9a10 100644
--- a/clearml_agent/commands/worker.py
+++ b/clearml_agent/commands/worker.py
@@ -2230,7 +2230,10 @@ class Worker(ServiceCommandSection):
                 os.environ.update(hyper_params)
 
         # Add the script CWD to the python path
-        python_path = get_python_path(script_dir, execution.entry_point, self.package_api, is_conda_env=self.is_conda)
+        if repo_info and repo_info.root and self._session.config.get('agent.force_git_root_python_path', None):
+            python_path = get_python_path(repo_info.root, None, self.package_api, is_conda_env=self.is_conda)
+        else:
+            python_path = get_python_path(script_dir, execution.entry_point, self.package_api, is_conda_env=self.is_conda)
         if ENV_TASK_EXTRA_PYTHON_PATH.get():
             python_path = add_python_path(python_path, ENV_TASK_EXTRA_PYTHON_PATH.get())
         if python_path:
diff --git a/clearml_agent/helper/base.py b/clearml_agent/helper/base.py
index fd7107b..aa881a1 100644
--- a/clearml_agent/helper/base.py
+++ b/clearml_agent/helper/base.py
@@ -204,10 +204,13 @@ def get_python_path(script_dir, entry_point, package_api, is_conda_env=False):
             ["-c", "import sys; print('{}'.join(sys.path))".format(python_path_sep)])
         org_python_path = python_path_cmd.get_output(cwd=script_dir)
         # Add path of the script directory and executable directory
-        python_path = '{}{python_path_sep}{}{python_path_sep}'.format(
-            Path(script_dir).absolute().as_posix(),
-            (Path(script_dir) / Path(entry_point)).parent.absolute().as_posix(),
-            python_path_sep=python_path_sep)
+        python_path = '{}{python_path_sep}'.format(
+            Path(script_dir).absolute().as_posix(), python_path_sep=python_path_sep)
+        if entry_point:
+            python_path += '{}{python_path_sep}'.format(
+                (Path(script_dir) / Path(entry_point)).parent.absolute().as_posix(),
+                python_path_sep=python_path_sep)
+
         if is_windows_platform():
             python_path = python_path.replace('/', '\\')