ClearML/clearml-agent
1
0
Fork 0
mirror of https://github.com/clearml/clearml-agent synced 2025-01-31 09:06:52 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix PyTorch aarch64 and windows support

Browse Source
This commit is contained in:
allegroai 2022-03-16 17:40:21 +02:00
parent a51f9bed49
commit 8712c5e636
1 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
clearml_agent/helper/package/pytorch.py
View File

@ -2,6 +2,7 @@ from __future__ import unicode_literals
import re import re
import sys import sys
import platform
from furl import furl from furl import furl
import urllib.parse import urllib.parse
from operator import itemgetter from operator import itemgetter
@ -245,10 +246,15 @@ class PytorchRequirement(SimpleSubstitution):
return "macos" return "macos"
raise RuntimeError("unrecognized OS") raise RuntimeError("unrecognized OS")
@staticmethod
def get_arch():
return str(platform.machine()).lower()
def _get_link_from_torch_page(self, req, torch_url): def _get_link_from_torch_page(self, req, torch_url):
links_parser = LinksHTMLParser() links_parser = LinksHTMLParser()
links_parser.feed(requests.get(torch_url, timeout=10).text) links_parser.feed(requests.get(torch_url, timeout=10).text)
platform_wheel = "win" if self.get_platform() == "windows" else self.get_platform() platform_wheel = "win" if self.get_platform() == "windows" else self.get_platform()
arch_wheel = self.get_arch()
py_ver = self.python_major_minor_str.replace('.', '') py_ver = self.python_major_minor_str.replace('.', '')
url = None url = None
last_v = None last_v = None
@ -269,8 +275,11 @@ class PytorchRequirement(SimpleSubstitution):
continue continue
if len(parts) < 3 or not parts[2].endswith(py_ver): if len(parts) < 3 or not parts[2].endswith(py_ver):
continue continue
if len(parts) < 5 or platform_wheel not in parts[4]: if len(parts) < 5 or platform_wheel not in parts[4].lower():
continue continue
if len(parts) < 5 or arch_wheel not in parts[4].lower():
continue
# yes this is for linux python 2.7 support, this is the only python 2.7 we support... # 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'): if py_ver and py_ver[0] == '2' and len(parts) > 3 and not parts[3].endswith('u'):
continue continue