From 0bb267115b00947a34116a0d8c74368d48eb7182 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Wed, 24 Jul 2024 17:59:18 +0300
Subject: [PATCH] Add venvs_cache.path mount override for non-root containers
 (use: agent.docker_internal_mounts.venvs_cache)

---
 clearml_agent/backend_api/config/default/agent.conf |  1 +
 clearml_agent/commands/worker.py                    | 12 +++++++++---
 clearml_agent/helper/package/base.py                |  2 ++
 docs/clearml.conf                                   |  1 +
 4 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/clearml_agent/backend_api/config/default/agent.conf b/clearml_agent/backend_api/config/default/agent.conf
index 037c7d5..51a7dc2 100644
--- a/clearml_agent/backend_api/config/default/agent.conf
+++ b/clearml_agent/backend_api/config/default/agent.conf
@@ -359,6 +359,7 @@
         pip_cache: "/root/.cache/pip"
         poetry_cache: "/root/.cache/pypoetry"
         vcs_cache: "/root/.clearml/vcs-cache"
+        venvs_cache: "/root/.clearml/venvs-cache"
         venv_build: "~/.clearml/venvs-builds"
         pip_download: "/root/.clearml/pip-download-cache"
     }
diff --git a/clearml_agent/commands/worker.py b/clearml_agent/commands/worker.py
index 6f9ba49..41c7a3f 100644
--- a/clearml_agent/commands/worker.py
+++ b/clearml_agent/commands/worker.py
@@ -3328,7 +3328,6 @@ class Worker(ServiceCommandSection):
         if not self.poetry.enabled:
             # add to cache
             if add_venv_folder_cache and not self._standalone_mode:
-                print('Adding venv into cache: {}'.format(add_venv_folder_cache))
                 self.package_api.add_cached_venv(
                     requirements=[freeze, previous_reqs],
                     docker_cmd=execution_info.docker_cmd if execution_info else None,
@@ -4036,8 +4035,15 @@ class Worker(ServiceCommandSection):
         if force_system_site_packages:
             temp_config.put("agent.package_manager.system_site_packages", True)
 
+        # if venvs_cache is NOT disabled,
+        # we have to update the mount path and the path inside the container
         if temp_config.get("agent.venvs_cache.path", None):
-            temp_config.put("agent.venvs_cache.path", '/root/.clearml/venvs-cache')
+            venvs_cache_path = temp_config.get("agent.docker_internal_mounts.venvs_cache", None)
+            if not venvs_cache_path:
+                venvs_cache_path = "/root/.clearml/venvs-cache"
+                temp_config.put("agent.docker_internal_mounts.venvs_cache", venvs_cache_path)
+            # update the venvs cache path to the container scope
+            temp_config.put("agent.venvs_cache.path", venvs_cache_path)
 
         if (ENV_SSH_AUTH_SOCK.get() or '').strip():
             self._host_ssh_cache = None
@@ -4157,7 +4163,7 @@ class Worker(ServiceCommandSection):
         mounted_cache_dir = temp_config.get("sdk.storage.cache.default_base_dir")
         mounted_pip_dl_dir = temp_config.get("agent.pip_download_cache.path")
         mounted_vcs_cache = temp_config.get("agent.vcs_cache.path")
-        mounted_venvs_cache = temp_config.get("agent.venvs_cache.path", "")
+        mounted_venvs_cache = temp_config.get("agent.docker_internal_mounts.venvs_cache", "")
         mount_ssh = temp_config.get("agent.docker_internal_mounts.ssh_folder", None)
         mount_ssh_ro = temp_config.get("agent.docker_internal_mounts.ssh_ro_folder", None)
         mount_apt_cache = temp_config.get("agent.docker_internal_mounts.apt_cache", None)
diff --git a/clearml_agent/helper/package/base.py b/clearml_agent/helper/package/base.py
index 09d8147..2ce929f 100644
--- a/clearml_agent/helper/package/base.py
+++ b/clearml_agent/helper/package/base.py
@@ -219,6 +219,8 @@ class PackageManager(object):
         if not self._get_cache_manager():
             return
 
+        print('Adding venv into cache: {}'.format(source_folder))
+
         try:
             keys = self._generate_reqs_hash_keys(requirements, docker_cmd, python_version, cuda_version)
             return self._get_cache_manager().add_entry(
diff --git a/docs/clearml.conf b/docs/clearml.conf
index b47410a..1ba3bcf 100644
--- a/docs/clearml.conf
+++ b/docs/clearml.conf
@@ -308,6 +308,7 @@ agent {
     #     pip_cache: "/root/.cache/pip"
     #     poetry_cache: "/root/.cache/pypoetry"
     #     vcs_cache: "/root/.clearml/vcs-cache"
+    #     venvs_cache: "/root/.clearml/venvs-cache"
     #     venv_build: "~/.clearml/venvs-builds"
     #     pip_download: "/root/.clearml/pip-download-cache"
     # }