Fix PyTorch extra index pip resolver

This commit is contained in:
allegroai 2023-09-02 17:43:41 +03:00
parent eefb94d1bc
commit fb639afcb9
2 changed files with 11 additions and 13 deletions

View File

@ -12,8 +12,6 @@ from clearml_agent.session import Session
class SystemPip(PackageManager):
indices_args = None
def __init__(self, interpreter=None, session=None):
# type: (Optional[Text], Optional[Session]) -> ()
"""
@ -99,21 +97,19 @@ class SystemPip(PackageManager):
return Argv(self.bin, '-m', 'pip', '--disable-pip-version-check', *command)
def install_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)
)
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
return (indices_args + tuple(extra_pip_flags)) if extra_pip_flags else 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)
)
indices_args = tuple(
chain.from_iterable(('--extra-index-url', x) for x in PIP_EXTRA_INDICES)
)
return self.indices_args
return indices_args

View File

@ -527,6 +527,8 @@ class PytorchRequirement(SimpleSubstitution):
# return the original line
line = req.line
print("PyTorch: Adding index `{}` and installing `{}`".format(extra_index_url[0], line))
return line
except Exception: # noqa
@ -681,7 +683,7 @@ class PytorchRequirement(SimpleSubstitution):
# noinspection PyBroadException
try:
if requests.get(torch_url, timeout=10).ok:
print('Torch CUDA {} index page found'.format(c))
print('Torch CUDA {} index page found, adding `{}`'.format(c, torch_url))
cls.torch_index_url_lookup[c] = torch_url
return cls.torch_index_url_lookup[c], c
except Exception: