diff --git a/clearml_agent/definitions.py b/clearml_agent/definitions.py index 63093ee..f73c18b 100644 --- a/clearml_agent/definitions.py +++ b/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") diff --git a/clearml_agent/helper/package/base.py b/clearml_agent/helper/package/base.py index 504e1f4..c4306e2 100644 --- a/clearml_agent/helper/package/base.py +++ b/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 [] diff --git a/clearml_agent/helper/package/pip_api/system.py b/clearml_agent/helper/package/pip_api/system.py index 99ee6a8..8db6105 100644 --- a/clearml_agent/helper/package/pip_api/system.py +++ b/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 diff --git a/clearml_agent/helper/package/poetry_api.py b/clearml_agent/helper/package/poetry_api.py index 5687a1a..def9387 100644 --- a/clearml_agent/helper/package/poetry_api.py +++ b/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]} diff --git a/clearml_agent/helper/package/priority_req.py b/clearml_agent/helper/package/priority_req.py index fe2e6be..2288207 100644 --- a/clearml_agent/helper/package/priority_req.py +++ b/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,31 +50,39 @@ 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 - try: - for k, lines in list_of_requirements.items(): - # k is either pip/conda - if k not in ('pip', 'conda'): - continue - for i, line in enumerate(lines): - if not line or line.lstrip().startswith('#'): - continue - parts = [p for p in re.split(r'\s|=|\.|<|>|~|!|@|#', line) if p] - if not parts: - continue - # if we found setuptools, do nothing - if parts[0] == "setuptools": - 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 we are here it means we have not found setuptools - # we should add it: - if "pip" in list_of_requirements: - list_of_requirements["pip"] = [self._replaced_packages["setuptools"]] + list_of_requirements["pip"] + if "setuptools" in self._replaced_packages: + try: + for k, lines in list_of_requirements.items(): + # k is either pip/conda + if k not in ('pip', 'conda'): + continue + for i, line in enumerate(lines): + if not line or line.lstrip().startswith('#'): + continue + parts = [p for p in re.split(r'\s|=|\.|<|>|~|!|@|#', line) if p] + if not parts: + continue + # if we found setuptools, do nothing + if parts[0] == "setuptools": + return list_of_requirements - except Exception as ex: # noqa - return list_of_requirements + # if we are here it means we have not found setuptools + # we should add it: + if "pip" in list_of_requirements: + list_of_requirements["pip"] = [self._replaced_packages["setuptools"]] + list_of_requirements["pip"] + + except Exception as ex: # noqa + return list_of_requirements return list_of_requirements diff --git a/docs/clearml.conf b/docs/clearml.conf index 57413bf..20749a2 100644 --- a/docs/clearml.conf +++ b/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