From d16825029de37e3f9727d2e350737615bbff03e2 Mon Sep 17 00:00:00 2001 From: allegroai <> Date: Sat, 2 Sep 2023 17:45:10 +0300 Subject: [PATCH] Add new pytorch no resolver mode and CLEARML_AGENT_PACKAGE_PYTORCH_RESOLVE to change resolver on a Task basis, now supports "pip", "direct", "none" --- .../backend_api/config/default/agent.conf | 5 +-- clearml_agent/definitions.py | 2 ++ clearml_agent/helper/package/pytorch.py | 31 +++++++++++++++---- docs/clearml.conf | 5 +-- 4 files changed, 33 insertions(+), 10 deletions(-) diff --git a/clearml_agent/backend_api/config/default/agent.conf b/clearml_agent/backend_api/config/default/agent.conf index b6cf157..ec85567 100644 --- a/clearml_agent/backend_api/config/default/agent.conf +++ b/clearml_agent/backend_api/config/default/agent.conf @@ -80,14 +80,15 @@ # additional artifact repositories to use when installing python packages # extra_index_url: ["https://allegroai.jfrog.io/clearml/api/pypi/public/simple"] - # control the pytorch wheel resolving algorithm, options are: "pip", "direct" + # control the pytorch wheel resolving algorithm, options are: "pip", "direct", "none" + # Override with environment variable CLEARML_AGENT_PACKAGE_PYTORCH_RESOLVE # "pip" (default): would automatically detect the cuda version, and supply pip with the correct # extra-index-url, based on pytorch.org tables # "direct": would resolve a direct link to the pytorch wheel by parsing the pytorch.org pip repository # and matching the automatically detected cuda version with the required pytorch wheel. # if the exact cuda version is not found for the required pytorch wheel, it will try # a lower cuda version until a match is found - # + # "none": No resolver used, install pytorch like any other package # pytorch_resolve: "pip" # additional conda channels to use when installing with conda package manager diff --git a/clearml_agent/definitions.py b/clearml_agent/definitions.py index f73c18b..8ae8071 100644 --- a/clearml_agent/definitions.py +++ b/clearml_agent/definitions.py @@ -238,6 +238,8 @@ ENV_CUSTOM_BUILD_SCRIPT = EnvironmentConfig("CLEARML_AGENT_CUSTOM_BUILD_SCRIPT") standard flow. """ +ENV_PACKAGE_PYTORCH_RESOLVE = EnvironmentConfig("CLEARML_AGENT_PACKAGE_PYTORCH_RESOLVE") + class FileBuffering(IntEnum): """ diff --git a/clearml_agent/helper/package/pytorch.py b/clearml_agent/helper/package/pytorch.py index 3b90602..15e8d4c 100644 --- a/clearml_agent/helper/package/pytorch.py +++ b/clearml_agent/helper/package/pytorch.py @@ -16,6 +16,7 @@ import six from .requirements import ( SimpleSubstitution, FatalSpecsResolutionError, SimpleVersion, MarkerRequirement, compare_version_rules, ) +from ...definitions import ENV_PACKAGE_PYTORCH_RESOLVE from ...external.requirements_parser.requirement import Requirement OS_TO_WHEEL_NAME = {"linux": "linux_x86_64", "windows": "win_amd64"} @@ -174,6 +175,7 @@ class PytorchRequirement(SimpleSubstitution): extra_index_url_template = 'https://download.pytorch.org/whl/cu{}/' nightly_extra_index_url_template = 'https://download.pytorch.org/whl/nightly/cu{}/' torch_index_url_lookup = {} + resolver_types = ("pip", "direct", "none") def __init__(self, *args, **kwargs): os_name = kwargs.pop("os_override", None) @@ -208,6 +210,13 @@ class PytorchRequirement(SimpleSubstitution): if self.config.get("agent.package_manager.torch_url_template", None): PytorchWheel.url_template = \ self.config.get("agent.package_manager.torch_url_template", None) + self.resolve_algorithm = str( + ENV_PACKAGE_PYTORCH_RESOLVE.get() or + self.config.get("agent.package_manager.pytorch_resolve", "pip")).lower() + if self.resolve_algorithm not in self.resolver_types: + print("WARNING: agent.package_manager.pytorch_resolve=={} not in {} reverting to '{}'".format( + self.resolve_algorithm, self.resolver_types, self.resolver_types[0])) + self.resolve_algorithm = self.resolver_types[0] def _init_python_ver_cuda_ver(self): if self.cuda is None: @@ -261,6 +270,10 @@ class PytorchRequirement(SimpleSubstitution): ) def match(self, req): + if self.resolve_algorithm == "none": + # skipping resolver + return False + return req.name in self.packages @staticmethod @@ -347,8 +360,10 @@ class PytorchRequirement(SimpleSubstitution): from pip._internal.commands.show import search_packages_info installed_torch = list(search_packages_info([req.name])) # notice the comparison order, the first part will make sure we have a valid installed package - installed_torch_version = (getattr(installed_torch[0], 'version', None) or installed_torch[0]['version']) \ - if installed_torch else None + installed_torch_version = \ + (getattr(installed_torch[0], 'version', None) or + installed_torch[0]['version']) if installed_torch else None + if installed_torch and installed_torch_version and \ req.compare_version(installed_torch_version): print('PyTorch: requested "{}" version {}, using pre-installed version {}'.format( @@ -356,6 +371,7 @@ class PytorchRequirement(SimpleSubstitution): # package already installed, do nothing req.specs = [('==', str(installed_torch_version))] return '{} {} {}'.format(req.name, req.specs[0][0], req.specs[0][1]), True + except Exception: pass @@ -480,8 +496,11 @@ class PytorchRequirement(SimpleSubstitution): # we first try to resolve things ourselves because pytorch pip is not always picking the correct # versions from their pip repository - resolve_algorithm = str(self.config.get("agent.package_manager.pytorch_resolve", "pip")).lower() - if resolve_algorithm == "direct": + resolve_algorithm = self.resolve_algorithm + if resolve_algorithm == "none": + # skipping resolver + return None + elif resolve_algorithm == "direct": # noinspection PyBroadException try: new_req = self._replace(req) @@ -489,8 +508,8 @@ class PytorchRequirement(SimpleSubstitution): self._original_req.append((req, new_req)) return new_req except Exception: - pass - elif resolve_algorithm not in ("direct", "pip"): + print("Warning: Failed resolving using `pytorch_resolve=direct` reverting to `pytorch_resolve=pip`") + elif resolve_algorithm not in self.resolver_types: print("Warning: `agent.package_manager.pytorch_resolve={}` " "unrecognized, default to `pip`".format(resolve_algorithm)) diff --git a/docs/clearml.conf b/docs/clearml.conf index bca728b..7de6934 100644 --- a/docs/clearml.conf +++ b/docs/clearml.conf @@ -96,14 +96,15 @@ agent { # 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" + # control the pytorch wheel resolving algorithm, options are: "pip", "direct", "none" + # Override with environment variable CLEARML_AGENT_PACKAGE_PYTORCH_RESOLVE # "pip" (default): would automatically detect the cuda version, and supply pip with the correct # extra-index-url, based on pytorch.org tables # "direct": would resolve a direct link to the pytorch wheel by parsing the pytorch.org pip repository # and matching the automatically detected cuda version with the required pytorch wheel. # if the exact cuda version is not found for the required pytorch wheel, it will try # a lower cuda version until a match is found - # + # "none": No resolver used, install pytorch like any other package # pytorch_resolve: "pip" # additional conda channels to use when installing with conda package manager