From 3cedc104dfdd6d87a772c891c8d5c9801a0fdd01 Mon Sep 17 00:00:00 2001
From: Niels ten Boom <nielstenboom@gmail.com>
Date: Sun, 5 Mar 2023 13:19:57 +0100
Subject: [PATCH] Add poetry cwd support (#142)

Closes #138
---
 .../backend_api/config/default/agent.conf        |  4 ++++
 clearml_agent/commands/worker.py                 | 16 +++++++++++-----
 2 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 2aa8766..e8475dc 100644
--- a/clearml_agent/backend_api/config/default/agent.conf
+++ b/clearml_agent/backend_api/config/default/agent.conf
@@ -105,6 +105,10 @@
         # set to True to support torch nightly build installation,
         # notice: torch nightly builds are ephemeral and are deleted from time to time
         torch_nightly: false,
+
+        # if set to true, the agent will look for the "poetry.lock" file 
+        # in the passed current working directory instead of the repository's root directory.
+        poetry_files_from_repo_working_dir: false
     },
 
     # target folder for virtual environments builds, created when executing experiment
diff --git a/clearml_agent/commands/worker.py b/clearml_agent/commands/worker.py
index 1366776..6dbf446 100644
--- a/clearml_agent/commands/worker.py
+++ b/clearml_agent/commands/worker.py
@@ -2876,19 +2876,22 @@ class Worker(ServiceCommandSection):
             self.log_traceback(e)
         return freeze
 
-    def _install_poetry_requirements(self, repo_info):
-        # type: (Optional[RepoInfo]) -> Optional[PoetryAPI]
+    def _install_poetry_requirements(self, repo_info, lockfile_path):
+        # type: (Optional[RepoInfo], Path) -> Optional[PoetryAPI]
         if not repo_info:
             return None
         try:
             if not self.poetry.enabled:
                 return None
-            self.poetry.initialize(cwd=repo_info.root)
-            api = self.poetry.get_api(repo_info.root)
+
+            self.poetry.initialize(cwd=lockfile_path)
+            api = self.poetry.get_api(lockfile_path)
             if api.enabled:
                 print('Poetry Enabled: Ignoring requested python packages, using repository poetry lock file!')
                 api.install()
                 return api
+            
+            print(f"Could not find pyproject.toml or poetry.lock file in {lockfile_path} \n")
         except Exception as ex:
             self.log.error("failed installing poetry requirements: {}".format(ex))
         return None
@@ -2919,7 +2922,10 @@ class Worker(ServiceCommandSection):
          """
         if package_api:
             package_api.cwd = cwd
-        api = self._install_poetry_requirements(repo_info)
+
+        files_from_working_dir = package_api.session.config.get("agent.package_manager.poetry_files_from_repo_working_dir", False)
+        lockfile_path = Path(repo_info.root) / (execution.working_dir if files_from_working_dir else "")
+        api = self._install_poetry_requirements(repo_info, lockfile_path)
         if api:
             # update back the package manager, this hack should be fixed
             if package_api == self.package_api: