From 1260e3d942bc1ebdc984c35358b7d8ab9843a593 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Thu, 11 Feb 2021 14:47:26 +0200
Subject: [PATCH] Update cache entries on conda package manager

---
 clearml_agent/commands/worker.py          | 14 ++++++----
 clearml_agent/helper/package/conda_api.py | 31 +++++++++++++----------
 2 files changed, 27 insertions(+), 18 deletions(-)

diff --git a/clearml_agent/commands/worker.py b/clearml_agent/commands/worker.py
index d6f00d5..09b6231 100644
--- a/clearml_agent/commands/worker.py
+++ b/clearml_agent/commands/worker.py
@@ -848,8 +848,6 @@ class Worker(ServiceCommandSection):
         return None
 
     def set_runtime_properties(self, key, value):
-        ##
-        return True
         if self._runtime_props_support is not True:
             # either not supported or never tested
             if self._runtime_props_support == self._session.api_version:
@@ -1463,7 +1461,8 @@ class Worker(ServiceCommandSection):
                 cwd=vcs.location if vcs and vcs.location else directory,
                 package_api=self.global_package_api if install_globally else None,
             )
-        freeze = self.freeze_task_environment(requirements_manager=requirements_manager)
+        freeze = self.freeze_task_environment(
+            task_id=task_id, requirements_manager=requirements_manager, update_requirements=False)
         script_dir = directory
 
         # Summary
@@ -1711,10 +1710,11 @@ class Worker(ServiceCommandSection):
             "agent.package_manager.conda_full_env_update", False)
 
         freeze = self.freeze_task_environment(
-            current_task.id if not skip_freeze_update else None,
+            task_id=current_task.id,
             requirements_manager=requirements_manager,
             add_venv_folder_cache=venv_folder,
             execution_info=execution,
+            update_requirements=not skip_freeze_update,
         )
         script_dir = (directory if isinstance(directory, Path) else Path(directory)).absolute().as_posix()
 
@@ -2020,7 +2020,7 @@ class Worker(ServiceCommandSection):
             )
 
     def freeze_task_environment(self, task_id=None, requirements_manager=None,
-                                add_venv_folder_cache=None, execution_info=None):
+                                add_venv_folder_cache=None, execution_info=None, update_requirements=False):
         try:
             freeze = self.package_api.freeze()
         except Exception as e:
@@ -2061,6 +2061,10 @@ class Worker(ServiceCommandSection):
                 source_folder=add_venv_folder_cache,
                 exclude_sub_folders=['task_repository', 'code'])
 
+        # If do not update back requirements
+        if not update_requirements:
+            return freeze
+
         request = tasks_api.SetRequirementsRequest(task=task_id, requirements=requirements)
         try:
             self._session.send_api(request)
diff --git a/clearml_agent/helper/package/conda_api.py b/clearml_agent/helper/package/conda_api.py
index a9fe45c..29ab8c1 100644
--- a/clearml_agent/helper/package/conda_api.py
+++ b/clearml_agent/helper/package/conda_api.py
@@ -141,19 +141,7 @@ class CondaAPI(PackageManager):
         """
         if self.conda_env_as_base_docker and self.conda_pre_build_env_path:
             if Path(self.conda_pre_build_env_path).is_dir():
-                print("Using pre-existing Conda environment from {}".format(self.conda_pre_build_env_path))
-                self.path = Path(self.conda_pre_build_env_path)
-                self.source = ("conda", "activate", self.path.as_posix())
-                self.pip = CondaPip(
-                    session=self.session,
-                    source=self.source,
-                    python=self.python,
-                    requirements_manager=self.requirements_manager,
-                    path=self.path,
-                )
-                conda_env = self._get_conda_sh()
-                self.source = self.pip.source = CommandSequence(('source', conda_env.as_posix()), self.source)
-                self.env_read_only = True
+                self._init_existing_environment(self.conda_pre_build_env_path)
                 return self
             elif Path(self.conda_pre_build_env_path).is_file():
                 print("Restoring Conda environment from {}".format(self.conda_pre_build_env_path))
@@ -211,6 +199,21 @@ class CondaAPI(PackageManager):
             pass
         return self
 
+    def _init_existing_environment(self, conda_pre_build_env_path):
+        print("Using pre-existing Conda environment from {}".format(conda_pre_build_env_path))
+        self.path = Path(conda_pre_build_env_path)
+        self.source = ("conda", "activate", self.path.as_posix())
+        self.pip = CondaPip(
+            session=self.session,
+            source=self.source,
+            python=self.python,
+            requirements_manager=self.requirements_manager,
+            path=self.path,
+        )
+        conda_env = self._get_conda_sh()
+        self.source = self.pip.source = CommandSequence(('source', conda_env.as_posix()), self.source)
+        self.env_read_only = True
+
     def remove(self):
         """
         Delete a conda environment.
@@ -666,6 +669,8 @@ class CondaAPI(PackageManager):
         return result
 
     def get_python_command(self, extra=()):
+        if not self.source:
+            self._init_existing_environment(self.path)
         return CommandSequence(self.source, self.pip.get_python_command(extra=extra))
 
     def _get_conda_sh(self):