Add CLEARML_EXTRA_PIP_INSTALL_FLAGS / agent.package_manager.extra_pip_install_flags to control additional pip install flags

Fix pip version marking in "installed packages" is now preserved for and reinstalled

clearml_agent/definitions.py

@@ -152,6 +152,7 @@ WORKING_STANDALONE_DIR = "code"
 DEFAULT_VCS_CACHE = normalize_path(CONFIG_DIR, "vcs-cache")
 PIP_EXTRA_INDICES = []
 DEFAULT_PIP_DOWNLOAD_CACHE = normalize_path(CONFIG_DIR, "pip-download-cache")
+ENV_PIP_EXTRA_INSTALL_FLAGS = EnvironmentConfig("CLEARML_EXTRA_PIP_INSTALL_FLAGS", type=list)
 ENV_DOCKER_IMAGE = EnvironmentConfig("CLEARML_DOCKER_IMAGE", "TRAINS_DOCKER_IMAGE")
 ENV_WORKER_ID = EnvironmentConfig("CLEARML_WORKER_ID", "TRAINS_WORKER_ID")
 ENV_WORKER_TAGS = EnvironmentConfig("CLEARML_WORKER_TAGS")

clearml_agent/helper/package/base.py

@@ -50,7 +50,7 @@ class PackageManager(object):
         pass
 
     @abc.abstractmethod
-    def freeze(self):
+    def freeze(self, freeze_full_environment=False):
         pass
 
     @abc.abstractmethod
@@ -141,8 +141,9 @@ class PackageManager(object):
     @classmethod
     def out_of_scope_install_package(cls, package_name, *args):
         if PackageManager._selected_manager is not None:
+            # noinspection PyBroadException
             try:
-                result = PackageManager._selected_manager._install(package_name, *args)
+                result = PackageManager._selected_manager.install_packages(package_name, *args)
                 if result not in (0, None, True):
                     return False
             except Exception:
@@ -150,10 +151,11 @@ class PackageManager(object):
         return True
 
     @classmethod
-    def out_of_scope_freeze(cls):
+    def out_of_scope_freeze(cls, freeze_full_environment=False):
         if PackageManager._selected_manager is not None:
+            # noinspection PyBroadException
             try:
-                return PackageManager._selected_manager.freeze()
+                return PackageManager._selected_manager.freeze(freeze_full_environment)
             except Exception:
                 pass
         return []

clearml_agent/helper/package/pip_api/system.py

@@ -4,7 +4,7 @@ from itertools import chain
 from pathlib import Path
 from typing import Text, Optional
 
-from clearml_agent.definitions import PIP_EXTRA_INDICES, PROGRAM_NAME
+from clearml_agent.definitions import PIP_EXTRA_INDICES, PROGRAM_NAME, ENV_PIP_EXTRA_INSTALL_FLAGS
 from clearml_agent.helper.package.base import PackageManager
 from clearml_agent.helper.process import Argv, DEVNULL
 from clearml_agent.session import Session
@@ -52,7 +52,7 @@ class SystemPip(PackageManager):
                 package,
                 '--dest', cache_dir,
                 '--no-deps',
-            ) + self.install_flags()
+            ) + self.download_flags()
         )
 
     def load_requirements(self, requirements):
@@ -65,13 +65,14 @@ class SystemPip(PackageManager):
     def uninstall(self, package):
         self.run_with_env(('uninstall', '-y', package))
 
-    def freeze(self):
+    def freeze(self, freeze_full_environment=False):
         """
         pip freeze to all install packages except the running program
         :return: Dict contains pip as key and pip's packages to install
         :rtype: Dict[str: List[str]]
         """
-        packages = self.run_with_env(('freeze',), output=True).splitlines()
+        packages = self.run_with_env(
+            ('freeze',) if not freeze_full_environment else ('freeze', '--all'), output=True).splitlines()
         packages_without_program = [package for package in packages if PROGRAM_NAME not in package]
         return {'pip': packages_without_program}
 
@@ -87,6 +88,11 @@ class SystemPip(PackageManager):
         # make sure we are not running it with our own PYTHONPATH
         env = dict(**os.environ)
         env.pop('PYTHONPATH', None)
+
+        # Debug print
+        if self.session.debug_mode:
+            print(command)
+
         return (command.get_output if output else command.check_call)(stdin=DEVNULL, env=env, **kwargs)
 
     def _make_command(self, command):
@@ -97,4 +103,17 @@ class SystemPip(PackageManager):
             self.indices_args = tuple(
                 chain.from_iterable(('--extra-index-url', x) for x in PIP_EXTRA_INDICES)
             )
+
+        extra_pip_flags = \
+            ENV_PIP_EXTRA_INSTALL_FLAGS.get() or \
+            self.session.config.get("agent.package_manager.extra_pip_install_flags", None)
+
+        return (self.indices_args + tuple(extra_pip_flags)) if extra_pip_flags else self.indices_args
+
+    def download_flags(self):
+        if self.indices_args is None:
+            self.indices_args = tuple(
+                chain.from_iterable(('--extra-index-url', x) for x in PIP_EXTRA_INDICES)
+            )
+
         return self.indices_args

clearml_agent/helper/package/poetry_api.py

@@ -147,7 +147,7 @@ class PoetryAPI(object):
             any((self.path / indicator).exists() for indicator in self.INDICATOR_FILES)
         )
 
-    def freeze(self):
+    def freeze(self, freeze_full_environment=False):
         lines = self.config.run("show", cwd=str(self.path)).splitlines()
         lines = [[p for p in line.split(' ') if p] for line in lines]
         return {"pip": [parts[0]+'=='+parts[1]+' # '+' '.join(parts[2:]) for parts in lines]}

clearml_agent/helper/package/priority_req.py

@@ -7,7 +7,7 @@ from .requirements import SimpleSubstitution
 
 class PriorityPackageRequirement(SimpleSubstitution):
 
-    name = ("cython", "numpy", "setuptools", )
+    name = ("cython", "numpy", "setuptools", "pip", )
     optional_package_names = tuple()
 
     def __init__(self, *args, **kwargs):
@@ -50,9 +50,17 @@ class PriorityPackageRequirement(SimpleSubstitution):
         """
         # if we replaced setuptools, it means someone requested it, and since freeze will not contain it,
         # we need to add it manually
-        if not self._replaced_packages or "setuptools" not in self._replaced_packages:
+        if not self._replaced_packages:
             return list_of_requirements
 
+        if "pip" in self._replaced_packages:
+            full_freeze = PackageManager.out_of_scope_freeze(freeze_full_environment=True)
+            # now let's look for pip
+            pips = [line for line in full_freeze.get("pip", []) if line.split("==")[0] == "pip"]
+            if pips and "pip" in list_of_requirements:
+                list_of_requirements["pip"] = [pips[0]] + list_of_requirements["pip"]
+
+        if "setuptools" in self._replaced_packages:
             try:
                 for k, lines in list_of_requirements.items():
                     # k is either pip/conda

docs/clearml.conf

@@ -93,6 +93,9 @@ agent {
         # extra_index_url: ["https://allegroai.jfrog.io/clearml/api/pypi/public/simple"]
         extra_index_url: []
 
+        # additional flags to use when calling pip install, example: ["--use-deprecated=legacy-resolver", ]
+        # extra_pip_install_flags: []
+
         # control the pytorch wheel resolving algorithm, options are: "pip", "direct"
         # "pip" (default): would automatically detect the cuda version, and supply pip with the correct
         # extra-index-url, based on pytorch.org tables