mirror of
https://github.com/clearml/clearml-agent
synced 2025-01-31 17:16:51 +00:00
91 lines
2.9 KiB
Python
91 lines
2.9 KiB
Python
from typing import Text
|
|
|
|
from furl import furl
|
|
from pathlib2 import Path
|
|
|
|
from clearml_agent.config import Config
|
|
from .pip_api.system import SystemPip
|
|
|
|
|
|
class RequirementsTranslator(object):
|
|
|
|
"""
|
|
Translate explicit URLs to local URLs after downloading them to cache
|
|
"""
|
|
|
|
SUPPORTED_SCHEMES = ["http", "https", "ftp"]
|
|
|
|
def __init__(self, session, interpreter=None, cache_dir=None):
|
|
self._session = session
|
|
config = session.config
|
|
self.cache_dir = cache_dir or Path(config["agent.pip_download_cache.path"]).expanduser().as_posix()
|
|
self.enabled = config["agent.pip_download_cache.enabled"]
|
|
Path(self.cache_dir).mkdir(parents=True, exist_ok=True)
|
|
self.config = Config()
|
|
self.pip = SystemPip(interpreter=interpreter, session=self._session)
|
|
self._translate_back = {}
|
|
|
|
def download(self, url):
|
|
self.pip.download_package(url, cache_dir=self.cache_dir)
|
|
|
|
@classmethod
|
|
def is_supported_link(cls, line):
|
|
# type: (Text) -> bool
|
|
"""
|
|
Return whether requirement is a link that should be downloaded to cache
|
|
"""
|
|
url = furl(line)
|
|
return (
|
|
url.scheme
|
|
and url.scheme.lower() in cls.SUPPORTED_SCHEMES
|
|
and line.lstrip().lower().startswith(url.scheme.lower())
|
|
)
|
|
|
|
def translate(self, line):
|
|
"""
|
|
If requirement is supported, download it to cache and return the download path
|
|
"""
|
|
if not (self.enabled and self.is_supported_link(line)):
|
|
return line
|
|
command = self.config.command
|
|
command.log('Downloading "{}" to pip cache'.format(line))
|
|
url = furl(line)
|
|
try:
|
|
wheel_name = url.path.segments[-1]
|
|
except IndexError:
|
|
command.error('Could not parse wheel name of "{}"'.format(line))
|
|
return line
|
|
try:
|
|
self.download(line)
|
|
downloaded = Path(self.cache_dir, wheel_name).expanduser().as_uri()
|
|
except Exception:
|
|
command.error('Could not download wheel name of "{}"'.format(line))
|
|
return line
|
|
|
|
self._translate_back[str(downloaded)] = line
|
|
|
|
return downloaded
|
|
|
|
def replace_back(self, requirements):
|
|
if not requirements:
|
|
return requirements
|
|
|
|
for k in requirements:
|
|
# k is either pip/conda
|
|
if k not in ('pip', 'conda'):
|
|
continue
|
|
|
|
original_requirements = requirements[k]
|
|
new_requirements = []
|
|
for line in original_requirements:
|
|
local_file = [d for d in self._translate_back.keys() if d in line]
|
|
if local_file:
|
|
local_file = local_file[0]
|
|
new_requirements.append(line.replace(local_file, self._translate_back[local_file]))
|
|
else:
|
|
new_requirements.append(line)
|
|
|
|
requirements[k] = new_requirements
|
|
|
|
return requirements
|