From 5b5fb0b8a6582ad1cf02eb6b6f96b05f217a9602 Mon Sep 17 00:00:00 2001 From: allegroai <> Date: Sun, 21 May 2023 22:53:11 +0300 Subject: [PATCH] Add `agent.package_manager.pytorch_resolve` configuration setting with `pip` or `direct` values. `pip` sets extra index based on cuda and lets pip resolve, `direct` is the previous parsing algorithm that does the matching and downloading (default `pip`) --- .../backend_api/config/default/agent.conf | 10 +++++++ clearml_agent/helper/package/pytorch.py | 27 ++++++++++++++++--- clearml_agent/helper/package/requirements.py | 17 ++++++++++++ docs/clearml.conf | 10 +++++++ 4 files changed, 60 insertions(+), 4 deletions(-) diff --git a/clearml_agent/backend_api/config/default/agent.conf b/clearml_agent/backend_api/config/default/agent.conf index eef86ed..3eb4123 100644 --- a/clearml_agent/backend_api/config/default/agent.conf +++ b/clearml_agent/backend_api/config/default/agent.conf @@ -80,6 +80,16 @@ # 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" + # "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 + # + # pytorch_resolve: "pip" + # additional conda channels to use when installing with conda package manager conda_channels: ["pytorch", "conda-forge", "defaults", ] diff --git a/clearml_agent/helper/package/pytorch.py b/clearml_agent/helper/package/pytorch.py index 3abd1e8..1d5e64c 100644 --- a/clearml_agent/helper/package/pytorch.py +++ b/clearml_agent/helper/package/pytorch.py @@ -310,6 +310,12 @@ class PytorchRequirement(SimpleSubstitution): # yes this is for linux python 2.7 support, this is the only python 2.7 we support... if py_ver and py_ver[0] == '2' and len(parts) > 3 and not parts[3].endswith('u'): continue + + # check if this an actual match + if not req.compare_version(v) or \ + (last_v and SimpleVersion.compare_versions(last_v, '>', v, ignore_sub_versions=False)): + continue + # update the closest matched version (from above) if not closest_v: closest_v = v @@ -318,10 +324,6 @@ class PytorchRequirement(SimpleSubstitution): SimpleVersion.compare_versions( version_a=v, op='>=', version_b=req.specs[0][1], num_parts=3): closest_v = v - # check if this an actual match - if not req.compare_version(v) or \ - (last_v and SimpleVersion.compare_versions(last_v, '>', v, ignore_sub_versions=False)): - continue url = '/'.join(torch_url.split('/')[:-1] + l.split('/')) last_v = v @@ -475,6 +477,23 @@ class PytorchRequirement(SimpleSubstitution): return self.match_version(req, base).replace(" ", "\n") def replace(self, req): + # 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": + # noinspection PyBroadException + try: + new_req = self._replace(req) + if new_req: + self._original_req.append((req, new_req)) + return new_req + except Exception: + pass + elif resolve_algorithm not in ("direct", "pip"): + print("Warning: `agent.package_manager.pytorch_resolve={}` " + "unrecognized, default to `pip`".format(resolve_algorithm)) + # check if package is already installed with system packages self.validate_python_version() diff --git a/clearml_agent/helper/package/requirements.py b/clearml_agent/helper/package/requirements.py index f04cdda..67296a6 100644 --- a/clearml_agent/helper/package/requirements.py +++ b/clearml_agent/helper/package/requirements.py @@ -240,6 +240,23 @@ class SimpleVersion: if not version_b: return True + # remove trailing "*" in both + if "*" in version_a: + ignore_sub_versions = True + while version_a.endswith(".*"): + version_a = version_a[:-2] + if version_a == "*": + version_a = "" + num_parts = min(len(version_a.split('.')), len(version_b.split('.')), ) + + if "*" in version_b: + ignore_sub_versions = True + while version_b.endswith(".*"): + version_b = version_b[:-2] + if version_b == "*": + version_b = "" + num_parts = min(len(version_a.split('.')), len(version_b.split('.')), ) + if not num_parts: num_parts = max(len(version_a.split('.')), len(version_b.split('.')), ) diff --git a/docs/clearml.conf b/docs/clearml.conf index cba1c8a..8dc3c42 100644 --- a/docs/clearml.conf +++ b/docs/clearml.conf @@ -93,6 +93,16 @@ agent { # extra_index_url: ["https://allegroai.jfrog.io/clearml/api/pypi/public/simple"] extra_index_url: [] + # 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 + # "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 + # + # pytorch_resolve: "pip" + # additional conda channels to use when installing with conda package manager conda_channels: ["pytorch", "conda-forge", "defaults", ] # conda_full_env_update: false