From 73e8fc80fe5881b303ba3bdaed4de61f7464a7d4 Mon Sep 17 00:00:00 2001 From: Philip Meier Date: Tue, 29 Jun 2021 17:22:05 +0200 Subject: [PATCH] Support multiple CUDA versions (#33) * Support multiple CUDA versions * cleanup * exclude specific windows combinations from tests * fix cpu backend ordering * prioritize backend over version * fix docstring --- light_the_torch/_pip/find.py | 113 +++++++---- light_the_torch/cli/commands.py | 7 +- light_the_torch/computation_backend.py | 160 +++++++++------ tests/cli/test_install.py | 4 +- tests/unit/{pip => _pip}/__init__.py | 0 tests/unit/{pip => _pip}/test_common.py | 0 tests/unit/{pip => _pip}/test_extract.py | 0 tests/unit/_pip/test_find.py | 180 +++++++++++++++++ tests/unit/conftest.py | 17 ++ tests/unit/pip/test_find.py | 213 -------------------- tests/unit/test_computation_backend.py | 245 ++++++++++++++--------- 11 files changed, 524 insertions(+), 415 deletions(-) rename tests/unit/{pip => _pip}/__init__.py (100%) rename tests/unit/{pip => _pip}/test_common.py (100%) rename tests/unit/{pip => _pip}/test_extract.py (100%) create mode 100644 tests/unit/_pip/test_find.py create mode 100644 tests/unit/conftest.py delete mode 100644 tests/unit/pip/test_find.py diff --git a/light_the_torch/_pip/find.py b/light_the_torch/_pip/find.py index 84f4464..af54428 100644 --- a/light_the_torch/_pip/find.py +++ b/light_the_torch/_pip/find.py @@ -1,5 +1,17 @@ import re -from typing import Any, Iterable, List, NoReturn, Optional, Text, Tuple, Union, cast +from typing import ( + Any, + Collection, + Iterable, + List, + NoReturn, + Optional, + Set, + Text, + Tuple, + Union, + cast, +) from pip._internal.index.collector import LinkCollector from pip._internal.index.package_finder import ( @@ -13,8 +25,10 @@ from pip._internal.models.search_scope import SearchScope from pip._internal.req.req_install import InstallRequirement from pip._internal.req.req_set import RequirementSet +from pip._vendor.packaging.version import Version + +import light_the_torch.computation_backend as cb -from ..computation_backend import ComputationBackend, detect_computation_backend from .common import ( InternalLTTError, PatchedInstallCommand, @@ -29,7 +43,9 @@ def find_links( pip_install_args: List[str], - computation_backend: Optional[Union[str, ComputationBackend]] = None, + computation_backends: Optional[ + Union[cb.ComputationBackend, Collection[cb.ComputationBackend]] + ] = None, channel: str = "stable", platform: Optional[str] = None, python_version: Optional[str] = None, @@ -41,9 +57,9 @@ def find_links( Args: pip_install_args: Arguments passed to ``pip install`` that will be searched for required PyTorch distributions - computation_backend: Computation backend, for example ``"cpu"`` or ``"cu102"``. - Defaults to the available hardware of the running system preferring CUDA - over CPU. + computation_backends: Collection of supported computation backends, for example + ``"cpu"`` or ``"cu102"``. Defaults to the available hardware of the running + system. channel: Channel of the PyTorch wheels. Can be one of ``"stable"`` (default), ``"test"``, and ``"nightly"``. platform: Platform, for example ``"linux_x86_64"`` or ``"win_amd64"``. Defaults @@ -55,10 +71,12 @@ def find_links( Returns: Wheel links with given properties for all required PyTorch distributions. """ - if computation_backend is None: - computation_backend = detect_computation_backend() - elif isinstance(computation_backend, str): - computation_backend = ComputationBackend.from_str(computation_backend) + if computation_backends is None: + computation_backends = cb.detect_compatible_computation_backends() + elif isinstance(computation_backends, cb.ComputationBackend): + computation_backends = {computation_backends} + else: + computation_backends = set(computation_backends) if channel not in ("stable", "test", "nightly"): raise ValueError( @@ -69,7 +87,7 @@ def find_links( dists = extract_dists(pip_install_args) cmd = StopAfterPytorchLinksFoundCommand( - computation_backend=computation_backend, channel=channel + computation_backends=computation_backends, channel=channel ) pip_install_args = adjust_pip_install_args(dists, platform, python_version) options, args = cmd.parser.parse_args(pip_install_args) @@ -172,37 +190,43 @@ def extract_computation_backend_from_link(self, link: Link) -> Optional[str]: class PytorchCandidatePreferences(CandidatePreferences): def __init__( - self, *args: Any, computation_backend: ComputationBackend, **kwargs: Any, + self, + *args: Any, + computation_backends: Set[cb.ComputationBackend], + **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) - self.computation_backend = computation_backend + self.computation_backends = computation_backends @classmethod def from_candidate_preferences( cls, candidate_preferences: CandidatePreferences, - computation_backend: ComputationBackend, + computation_backends: Set[cb.ComputationBackend], ) -> "PytorchCandidatePreferences": return new_from_similar( cls, candidate_preferences, ("prefer_binary", "allow_all_prereleases",), - computation_backend=computation_backend, + computation_backends=computation_backends, ) class PytorchCandidateEvaluator(CandidateEvaluator): def __init__( - self, *args: Any, computation_backend: ComputationBackend, **kwargs: Any, + self, + *args: Any, + computation_backends: Set[cb.ComputationBackend], + **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) - self.computation_backend = computation_backend + self.computation_backends = {cb.AnyBackend(), *computation_backends} @classmethod def from_candidate_evaluator( cls, candidate_evaluator: CandidateEvaluator, - computation_backend: ComputationBackend, + computation_backends: Set[cb.ComputationBackend], ) -> "PytorchCandidateEvaluator": return new_from_similar( cls, @@ -215,8 +239,19 @@ def from_candidate_evaluator( "allow_all_prereleases", "hashes", ), - computation_backend=computation_backend, + computation_backends=computation_backends, + ) + + def _sort_key( + self, candidate: InstallationCandidate + ) -> Tuple[cb.ComputationBackend, Version]: + version = Version( + f"{candidate.version.major}" + f".{candidate.version.minor}" + f".{candidate.version.micro}" ) + computation_backend = cb.ComputationBackend.from_str(candidate.version.local) + return computation_backend, version def get_applicable_candidates( self, candidates: List[InstallationCandidate] @@ -224,8 +259,7 @@ def get_applicable_candidates( return [ candidate for candidate in super().get_applicable_candidates(candidates) - if candidate.version.local == "any" - or candidate.version.local == self.computation_backend + if candidate.version.local in self.computation_backends ] @@ -233,33 +267,34 @@ class PytorchLinkCollector(LinkCollector): def __init__( self, *args: Any, - computation_backend: ComputationBackend, + computation_backends: Set[cb.ComputationBackend], channel: str = "stable", **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) if channel == "stable": - url = "https://download.pytorch.org/whl/torch_stable.html" + urls = ["https://download.pytorch.org/whl/torch_stable.html"] else: - url = ( + urls = [ f"https://download.pytorch.org/whl/" - f"{channel}/{computation_backend}/torch_{channel}.html" - ) - self.search_scope = SearchScope.create(find_links=[url], index_urls=[]) + f"{channel}/{backend}/torch_{channel}.html" + for backend in sorted(computation_backends, key=str) + ] + self.search_scope = SearchScope.create(find_links=urls, index_urls=[]) @classmethod def from_link_collector( cls, link_collector: LinkCollector, - computation_backend: ComputationBackend, + computation_backends: Set[cb.ComputationBackend], channel: str = "stable", ) -> "PytorchLinkCollector": return new_from_similar( cls, link_collector, - ("session", "search_scope",), + ("session", "search_scope"), channel=channel, - computation_backend=computation_backend, + computation_backends=computation_backends, ) @@ -270,18 +305,18 @@ class PytorchPackageFinder(PackageFinder): def __init__( self, *args: Any, - computation_backend: ComputationBackend, + computation_backends: Set[cb.ComputationBackend], channel: str = "stable", **kwargs: Any, ) -> None: super().__init__(*args, **kwargs) self._candidate_prefs = PytorchCandidatePreferences.from_candidate_preferences( - self._candidate_prefs, computation_backend=computation_backend + self._candidate_prefs, computation_backends=computation_backends ) self._link_collector = PytorchLinkCollector.from_link_collector( self._link_collector, channel=channel, - computation_backend=computation_backend, + computation_backends=computation_backends, ) def make_candidate_evaluator( @@ -290,7 +325,7 @@ def make_candidate_evaluator( candidate_evaluator = super().make_candidate_evaluator(*args, **kwargs) return PytorchCandidateEvaluator.from_candidate_evaluator( candidate_evaluator, - computation_backend=self._candidate_prefs.computation_backend, + computation_backends=self._candidate_prefs.computation_backends, ) def make_link_evaluator(self, *args: Any, **kwargs: Any) -> PytorchLinkEvaluator: @@ -301,7 +336,7 @@ def make_link_evaluator(self, *args: Any, **kwargs: Any) -> PytorchLinkEvaluator def from_package_finder( cls, package_finder: PackageFinder, - computation_backend: ComputationBackend, + computation_backends: Set[cb.ComputationBackend], channel: str = "stable", ) -> "PytorchPackageFinder": return new_from_similar( @@ -315,7 +350,7 @@ def from_package_finder( "candidate_prefs", "ignore_requires_python", ), - computation_backend=computation_backend, + computation_backends=computation_backends, channel=channel, ) @@ -338,19 +373,19 @@ def resolve( class StopAfterPytorchLinksFoundCommand(PatchedInstallCommand): def __init__( self, - computation_backend: ComputationBackend, + computation_backends: Set[cb.ComputationBackend], channel: str = "stable", **kwargs: Any, ) -> None: super().__init__(**kwargs) - self.computation_backend = computation_backend + self.computation_backends = computation_backends self.channel = channel def _build_package_finder(self, *args: Any, **kwargs: Any) -> PytorchPackageFinder: package_finder = super()._build_package_finder(*args, **kwargs) return PytorchPackageFinder.from_package_finder( package_finder, - computation_backend=self.computation_backend, + computation_backends=self.computation_backends, channel=self.channel, ) diff --git a/light_the_torch/cli/commands.py b/light_the_torch/cli/commands.py index c491e1b..c01628b 100644 --- a/light_the_torch/cli/commands.py +++ b/light_the_torch/cli/commands.py @@ -54,7 +54,8 @@ def _run(self, pip_install_args: List[str]) -> None: class FindCommand(Command): def __init__(self, args: argparse.Namespace) -> None: - self.computation_backend = args.computation_backend + # TODO split by comma + self.computation_backends = args.computation_backend self.channel = args.channel self.platform = args.platform self.python_version = args.python_version @@ -63,7 +64,7 @@ def __init__(self, args: argparse.Namespace) -> None: def _run(self, pip_install_args: List[str]) -> None: links = ltt.find_links( pip_install_args, - computation_backend=self.computation_backend, + computation_backends=self.computation_backends, channel=self.channel, platform=self.platform, python_version=self.python_version, @@ -88,7 +89,7 @@ def __init__(self, args: argparse.Namespace) -> None: def _run(self, pip_install_args: List[str]) -> None: links = ltt.find_links( pip_install_args, - computation_backend=CPUBackend() if self.force_cpu else None, + computation_backends={CPUBackend()} if self.force_cpu else None, channel=self.channel, verbose=self.verbose, ) diff --git a/light_the_torch/computation_backend.py b/light_the_torch/computation_backend.py index 38e916c..edb6e21 100644 --- a/light_the_torch/computation_backend.py +++ b/light_the_torch/computation_backend.py @@ -1,14 +1,17 @@ -import os +import platform import re import subprocess from abc import ABC, abstractmethod -from typing import Any, List, Optional +from typing import Any, Optional, Set + +from pip._vendor.packaging.version import InvalidVersion, Version __all__ = [ "ComputationBackend", + "AnyBackend", "CPUBackend", "CUDABackend", - "detect_computation_backend", + "detect_compatible_computation_backends", ] @@ -21,7 +24,11 @@ class ComputationBackend(ABC): @property @abstractmethod def local_specifier(self) -> str: - ... + pass + + @abstractmethod + def __lt__(self, other: Any) -> bool: + pass def __eq__(self, other: Any) -> bool: if isinstance(other, ComputationBackend): @@ -41,7 +48,9 @@ def __repr__(self) -> str: def from_str(cls, string: str) -> "ComputationBackend": parse_error = ParseError(string) string = string.lower() - if string == "cpu": + if string == "any": + return AnyBackend() + elif string == "cpu": return CPUBackend() elif string.startswith("cu"): match = re.match(r"^cu(da)?(?P[\d.]+)$", string) @@ -60,11 +69,29 @@ def from_str(cls, string: str) -> "ComputationBackend": raise parse_error +class AnyBackend(ComputationBackend): + @property + def local_specifier(self) -> str: + return "any" + + def __lt__(self, other: Any) -> bool: + if not isinstance(other, ComputationBackend): + return NotImplemented + + return False + + class CPUBackend(ComputationBackend): @property def local_specifier(self) -> str: return "cpu" + def __lt__(self, other: Any) -> bool: + if not isinstance(other, ComputationBackend): + return NotImplemented + + return True + class CUDABackend(ComputationBackend): def __init__(self, major: int, minor: int) -> None: @@ -75,65 +102,74 @@ def __init__(self, major: int, minor: int) -> None: def local_specifier(self) -> str: return f"cu{self.major}{self.minor}" + def __lt__(self, other: Any) -> bool: + if isinstance(other, AnyBackend): + return True + elif isinstance(other, CPUBackend): + return False + elif not isinstance(other, CUDABackend): + return NotImplemented -def detect_nvidia_driver() -> Optional[str]: - driver: Optional[str] - try: - output = subprocess.check_output( - "nvidia-smi --query-gpu=driver_version --format=csv", - shell=True, - stderr=subprocess.DEVNULL, - ) - driver = output.decode("utf-8").splitlines()[-1] - pattern = re.compile(r"(\d+\.\d+)") # match at least major and minor - if not pattern.match(driver): - driver = None - except subprocess.CalledProcessError: - driver = None - return driver - + return (self.major, self.minor) < (other.major, other.minor) -def get_supported_cuda_version() -> Optional[str]: - def split(version_string: str) -> List[int]: - return [int(n) for n in version_string.split(".")] - nvidia_driver = detect_nvidia_driver() - if nvidia_driver is None: +def _detect_nvidia_driver_version() -> Optional[Version]: + cmd = "nvidia-smi --query-gpu=driver_version --format=csv" + try: + output = ( + subprocess.check_output(cmd, shell=True, stderr=subprocess.DEVNULL) + .decode("utf-8") + .strip() + ) + return Version(output.splitlines()[-1]) + except (subprocess.CalledProcessError, InvalidVersion): return None - nvidia_driver = split(nvidia_driver) - cuda_version = None - if os.name == "nt": # windows - # Table 3 from https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html - if nvidia_driver >= split("456.38"): - cuda_version = "11.1" - elif nvidia_driver >= split("451.22"): - cuda_version = "11.0" - elif nvidia_driver >= split("441.22"): - cuda_version = "10.2" - elif nvidia_driver >= split("418.96"): - cuda_version = "10.1" - elif nvidia_driver >= split("398.26"): - cuda_version = "9.2" - else: # linux - # Table 1 from https://docs.nvidia.com/deploy/cuda-compatibility/index.html - if nvidia_driver >= split("450.80.02"): - cuda_version = "11.1" - elif nvidia_driver >= split("450.36.06"): - cuda_version = "11.0" - elif nvidia_driver >= split("440.33"): - cuda_version = "10.2" - elif nvidia_driver >= split("418.39"): - cuda_version = "10.1" - elif nvidia_driver >= split("396.26"): - cuda_version = "9.2" - return cuda_version - - -def detect_computation_backend() -> ComputationBackend: - cuda_version = get_supported_cuda_version() - if cuda_version is None: - return CPUBackend() - else: - major, minor = cuda_version.split(".") - return CUDABackend(int(major), int(minor)) + +# Table 3 from https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html +_MINIMUM_DRIVER_VERSIONS = { + "Linux": { + Version("11.1"): Version("455.32"), + Version("11.0"): Version("450.51.06"), + Version("10.2"): Version("440.33"), + Version("10.1"): Version("418.39"), + Version("10.0"): Version("410.48"), + Version("9.2"): Version("396.26"), + Version("9.1"): Version("390.46"), + Version("9.0"): Version("384.81"), + Version("8.0"): Version("375.26"), + Version("7.5"): Version("352.31"), + }, + "Windows": { + Version("11.1"): Version("456.81"), + Version("11.0"): Version("451.82"), + Version("10.2"): Version("441.22"), + Version("10.1"): Version("418.96"), + Version("10.0"): Version("411.31"), + Version("9.2"): Version("398.26"), + Version("9.1"): Version("391.29"), + Version("9.0"): Version("385.54"), + Version("8.0"): Version("376.51"), + Version("7.5"): Version("353.66"), + }, +} + + +def _detect_compatible_cuda_backends() -> Set[CUDABackend]: + driver_version = _detect_nvidia_driver_version() + if not driver_version: + return set() + + minimum_driver_versions = _MINIMUM_DRIVER_VERSIONS.get(platform.system()) + if not minimum_driver_versions: + return set() + + return { + CUDABackend(cuda_version.major, cuda_version.minor) + for cuda_version, minimum_driver_version in minimum_driver_versions.items() + if driver_version >= minimum_driver_version + } + + +def detect_compatible_computation_backends() -> Set[ComputationBackend]: + return {CPUBackend(), *_detect_compatible_cuda_backends()} diff --git a/tests/cli/test_install.py b/tests/cli/test_install.py index fa0c832..9c36b82 100644 --- a/tests/cli/test_install.py +++ b/tests/cli/test_install.py @@ -74,8 +74,8 @@ def test_ltt_install_force_cpu( cli.main() _, kwargs = find_links.call_args - assert "computation_backend" in kwargs - assert kwargs["computation_backend"] == CPUBackend() + assert "computation_backends" in kwargs + assert set(kwargs["computation_backends"]) == {CPUBackend()} def test_ltt_install_pytorch_only( diff --git a/tests/unit/pip/__init__.py b/tests/unit/_pip/__init__.py similarity index 100% rename from tests/unit/pip/__init__.py rename to tests/unit/_pip/__init__.py diff --git a/tests/unit/pip/test_common.py b/tests/unit/_pip/test_common.py similarity index 100% rename from tests/unit/pip/test_common.py rename to tests/unit/_pip/test_common.py diff --git a/tests/unit/pip/test_extract.py b/tests/unit/_pip/test_extract.py similarity index 100% rename from tests/unit/pip/test_extract.py rename to tests/unit/_pip/test_extract.py diff --git a/tests/unit/_pip/test_find.py b/tests/unit/_pip/test_find.py new file mode 100644 index 0000000..17faf94 --- /dev/null +++ b/tests/unit/_pip/test_find.py @@ -0,0 +1,180 @@ +import itertools + +import pytest +from pip._vendor.packaging.version import Version + +import light_the_torch as ltt +import light_the_torch.computation_backend as cb +from light_the_torch._pip.common import InternalLTTError +from light_the_torch._pip.find import maybe_add_option + + +@pytest.fixture +def patch_extract_dists(mocker): + def patch_extract_dists_(return_value=None): + if return_value is None: + return_value = [] + return mocker.patch( + "light_the_torch._pip.find.extract_dists", return_value=return_value + ) + + return patch_extract_dists_ + + +@pytest.fixture +def patch_run(mocker): + def patch_run_(): + return mocker.patch("light_the_torch._pip.find.run") + + return patch_run_ + + +CHANNELS = ("stable", "test", "nightly") +PLATFORMS = ("linux_x86_64", "macosx_10_9_x86_64", "win_amd64") +PLATFORM_MAP = dict(zip(PLATFORMS, ("Linux", "Darwin", "Windows"))) + + +SUPPORTED_PYTHON_VERSIONS = { + Version("11.1"): tuple(f"3.{minor}" for minor in (6, 7, 8, 9)), + Version("11.0"): tuple(f"3.{minor}" for minor in (6, 7, 8, 9)), + Version("10.2"): tuple(f"3.{minor}" for minor in (6, 7, 8, 9)), + Version("10.1"): tuple(f"3.{minor}" for minor in (6, 7, 8, 9)), + Version("10.0"): tuple(f"3.{minor}" for minor in (6, 7, 8)), + Version("9.2"): tuple(f"3.{minor}" for minor in (6, 7, 8, 9)), + Version("9.1"): tuple(f"3.{minor}" for minor in (6,)), + Version("9.0"): tuple(f"3.{minor}" for minor in (6, 7)), + Version("8.0"): tuple(f"3.{minor}" for minor in (6, 7)), + Version("7.5"): tuple(f"3.{minor}" for minor in (6,)), +} +PYTHON_VERSIONS = set(itertools.chain(*SUPPORTED_PYTHON_VERSIONS.values())) + + +def test_maybe_add_option_already_set(): + args = ["--foo", "bar"] + assert maybe_add_option(args, "--foo",) == args + assert maybe_add_option(args, "-f", aliases=("--foo",)) == args + + +def test_find_links_internal_error(patch_extract_dists, patch_run): + patch_extract_dists() + patch_run() + + with pytest.raises(InternalLTTError): + ltt.find_links([]) + + +def test_find_links_computation_backend_detect( + mocker, patch_extract_dists, patch_run, generic_backend +): + computation_backends = {generic_backend} + mocker.patch( + "light_the_torch.computation_backend.detect_compatible_computation_backends", + return_value=computation_backends, + ) + + patch_extract_dists() + run = patch_run() + + with pytest.raises(InternalLTTError): + ltt.find_links([], computation_backends=None) + + args, _ = run.call_args + cmd = args[0] + assert cmd.computation_backends == computation_backends + + +def test_find_links_unknown_channel(): + with pytest.raises(ValueError): + ltt.find_links([], channel="channel") + + +@pytest.mark.parametrize("platform", PLATFORMS) +def test_find_links_platform(patch_extract_dists, patch_run, platform): + patch_extract_dists() + run = patch_run() + + with pytest.raises(InternalLTTError): + ltt.find_links([], platform=platform) + + args, _ = run.call_args + options = args[2] + assert options.platform == platform + + +@pytest.mark.parametrize("python_version", PYTHON_VERSIONS) +def test_find_links_python_version(patch_extract_dists, patch_run, python_version): + patch_extract_dists() + run = patch_run() + + python_version_tuple = tuple(int(v) for v in python_version.split(".")) + + with pytest.raises(InternalLTTError): + ltt.find_links([], python_version=python_version) + + args, _ = run.call_args + options = args[2] + assert options.python_version == python_version_tuple + + +def wheel_properties(): + params = [] + for platform in PLATFORMS: + params.extend( + [ + (platform, cb.CPUBackend(), python_version) + for python_version in PYTHON_VERSIONS + ] + ) + + system = PLATFORM_MAP[platform] + cuda_versions = cb._MINIMUM_DRIVER_VERSIONS.get(system, {}).keys() + if not cuda_versions: + continue + + params.extend( + [ + ( + platform, + cb.CUDABackend(cuda_version.major, cuda_version.minor), + python_version, + ) + for cuda_version in cuda_versions + for python_version in SUPPORTED_PYTHON_VERSIONS[cuda_version] + if not ( + platform == "win_amd64" + and ( + (cuda_version == Version("7.5") and python_version == "3.6") + or (cuda_version == Version("9.2") and python_version == "3.9") + or (cuda_version == Version("10.0") and python_version == "3.8") + ) + ) + ] + ) + + return pytest.mark.parametrize( + ("platform", "computation_backend", "python_version"), params, ids=str, + ) + + +@pytest.mark.slow +@pytest.mark.parametrize( + "pytorch_dist", ["torch", "torchaudio", "torchtext", "torchvision"] +) +@wheel_properties() +def test_find_links_stable_smoke( + pytorch_dist, platform, computation_backend, python_version +): + assert ltt.find_links( + [pytorch_dist], + computation_backends=computation_backend, + platform=platform, + python_version=python_version, + ) + + +@pytest.mark.slow +@pytest.mark.parametrize("channel", CHANNELS) +def test_find_links_channel_smoke(channel): + assert ltt.find_links( + ["torch"], computation_backends={cb.CPUBackend()}, channel=channel + ) diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py new file mode 100644 index 0000000..b63c805 --- /dev/null +++ b/tests/unit/conftest.py @@ -0,0 +1,17 @@ +import pytest + +import light_the_torch.computation_backend as cb + + +class GenericComputationBackend(cb.ComputationBackend): + @property + def local_specifier(self): + return "generic" + + def __lt__(self, other): + return NotImplemented + + +@pytest.fixture +def generic_backend(): + return GenericComputationBackend() diff --git a/tests/unit/pip/test_find.py b/tests/unit/pip/test_find.py deleted file mode 100644 index 10c67ea..0000000 --- a/tests/unit/pip/test_find.py +++ /dev/null @@ -1,213 +0,0 @@ -import itertools - -import pytest - -import light_the_torch as ltt -from light_the_torch._pip.common import InternalLTTError -from light_the_torch._pip.find import maybe_add_option -from light_the_torch.computation_backend import ComputationBackend - - -@pytest.fixture -def patch_extract_dists(mocker): - def patch_extract_dists_(return_value=None): - if return_value is None: - return_value = [] - return mocker.patch( - "light_the_torch._pip.find.extract_dists", return_value=return_value - ) - - return patch_extract_dists_ - - -@pytest.fixture -def patch_run(mocker): - def patch_run_(): - return mocker.patch("light_the_torch._pip.find.run") - - return patch_run_ - - -@pytest.fixture -def computation_backends(): - return ("cpu", "cu92", "cu101", "cu102") - - -@pytest.fixture -def channels(): - return ("stable", "test", "nightly") - - -@pytest.fixture -def platforms(): - return ("linux_x86_64", "macosx_10_9_x86_64", "win_amd64") - - -@pytest.fixture -def python_versions(): - return ("3.6", "3.7", "3.8") - - -@pytest.fixture -def wheel_properties(computation_backends, platforms, python_versions): - properties = [] - for properties_ in itertools.product( - computation_backends, platforms, python_versions - ): - # macOS binaries don't support CUDA - computation_backend, platform, _ = properties_ - if platform.startswith("macosx") and computation_backend != "cpu": - continue - - properties.append( - dict( - zip(("computation_backend", "platform", "python_version"), properties_) - ) - ) - return tuple(properties) - - -def test_maybe_add_option_already_set(subtests): - args = ["--foo", "bar"] - assert maybe_add_option(args, "--foo",) == args - assert maybe_add_option(args, "-f", aliases=("--foo",)) == args - - -def test_find_links_internal_error(patch_extract_dists, patch_run): - patch_extract_dists() - patch_run() - - with pytest.raises(InternalLTTError): - ltt.find_links([]) - - -def test_find_links_computation_backend_detect(mocker, patch_extract_dists, patch_run): - class GenericComputationBackend(ComputationBackend): - @property - def local_specifier(self): - return "generic" - - computation_backend = GenericComputationBackend() - mocker.patch( - "light_the_torch._pip.find.detect_computation_backend", - return_value=computation_backend, - ) - - patch_extract_dists() - run = patch_run() - - with pytest.raises(InternalLTTError): - ltt.find_links([], computation_backend=None) - - args, _ = run.call_args - cmd = args[0] - assert cmd.computation_backend == computation_backend - - -def test_find_links_computation_backend_str( - subtests, patch_extract_dists, patch_run, computation_backends -): - patch_extract_dists() - run = patch_run() - - for computation_backend in computation_backends: - with subtests.test(computation_backend=computation_backend): - run.reset() - - with pytest.raises(InternalLTTError): - ltt.find_links([], computation_backend=computation_backend) - - args, _ = run.call_args - cmd = args[0] - assert cmd.computation_backend == ComputationBackend.from_str( - computation_backend - ) - - -def test_find_links_unknown_channel(): - with pytest.raises(ValueError): - ltt.find_links([], channel="channel") - - -def test_find_links_platform(subtests, patch_extract_dists, patch_run, platforms): - patch_extract_dists() - run = patch_run() - - for platform in platforms: - with subtests.test(platform=platform): - run.reset() - - with pytest.raises(InternalLTTError): - ltt.find_links([], platform=platform) - - args, _ = run.call_args - options = args[2] - assert options.platform == platform - - -def test_find_links_python_version( - subtests, patch_extract_dists, patch_run, python_versions -): - patch_extract_dists() - run = patch_run() - - for python_version in python_versions: - python_version_tuple = tuple(int(v) for v in python_version.split(".")) - with subtests.test(python_version=python_version): - run.reset() - - with pytest.raises(InternalLTTError): - ltt.find_links([], python_version=python_version) - - args, _ = run.call_args - options = args[2] - assert options.python_version == python_version_tuple - - -@pytest.mark.slow -def test_find_links_torch_smoke(subtests, wheel_properties): - dist = "torch" - - for properties in wheel_properties: - with subtests.test(**properties): - assert ltt.find_links([dist], **properties) - - -@pytest.mark.slow -def test_find_links_torchaudio_smoke(subtests, wheel_properties): - dist = "torchaudio" - - for properties in wheel_properties: - # torchaudio has no published releases for Windows - if properties["platform"].startswith("win"): - continue - with subtests.test(**properties): - a = ltt.find_links([dist], **properties) - assert a - - -@pytest.mark.slow -def test_find_links_torchtext_smoke(subtests, wheel_properties): - dist = "torchtext" - - for properties in wheel_properties: - with subtests.test(**properties): - assert ltt.find_links([dist], **properties) - - -@pytest.mark.slow -def test_find_links_torchvision_smoke(subtests, wheel_properties): - dist = "torchvision" - - for properties in wheel_properties: - with subtests.test(**properties): - assert ltt.find_links([dist], **properties) - - -@pytest.mark.slow -def test_find_links_torch_channel_smoke(subtests, channels): - dist = "torch" - - for channel in channels: - with subtests.test(channel=channel): - assert ltt.find_links([dist], computation_backend="cpu", channel=channel) diff --git a/tests/unit/test_computation_backend.py b/tests/unit/test_computation_backend.py index e70c100..f7952d7 100644 --- a/tests/unit/test_computation_backend.py +++ b/tests/unit/test_computation_backend.py @@ -5,122 +5,175 @@ from light_the_torch import computation_backend as cb -@pytest.fixture -def generic_backend(): - class GenericComputationBackend(cb.ComputationBackend): - @property - def local_specifier(self): - return "generic" - - return GenericComputationBackend() - - -def test_ComputationBackend_eq(generic_backend): - assert generic_backend == generic_backend - assert generic_backend == generic_backend.local_specifier - assert generic_backend != 0 - - -def test_ComputationBackend_hash_smoke(generic_backend): - assert isinstance(hash(generic_backend), int) - - -def test_ComputationBackend_repr_smoke(generic_backend): - assert isinstance(repr(generic_backend), str) - +class TestComputationBackend: + def test_eq(self, generic_backend): + assert generic_backend == generic_backend + assert generic_backend == generic_backend.local_specifier + assert generic_backend != 0 + + def test_hash_smoke(self, generic_backend): + assert isinstance(hash(generic_backend), int) + + def test_repr_smoke(self, generic_backend): + assert isinstance(repr(generic_backend), str) + + def test_from_str_cpu(self): + string = "cpu" + backend = cb.ComputationBackend.from_str(string) + assert isinstance(backend, cb.CPUBackend) + + @pytest.mark.parametrize( + ("major", "minor", "string"), + [ + pytest.param(major, minor, string, id=string) + for major, minor, string in ( + (12, 3, "cu123"), + (12, 3, "cu12.3"), + (12, 3, "cuda123"), + (12, 3, "cuda12.3"), + ) + ], + ) + def test_from_str_cuda(self, major, minor, string): + backend = cb.ComputationBackend.from_str(string) + assert isinstance(backend, cb.CUDABackend) + assert backend.major == major + assert backend.minor == minor -def test_ComputationBackend_from_str_cpu(): - string = "cpu" - backend = cb.ComputationBackend.from_str(string) - assert isinstance(backend, cb.CPUBackend) + @pytest.mark.parametrize("string", (("unknown", "cudnn"))) + def test_from_str_unknown(self, string): + with pytest.raises(cb.ParseError): + cb.ComputationBackend.from_str(string) -def test_ComputationBackend_from_str_cuda(subtests): - major, minor = 12, 3 - strings = ( - f"cu{major}{minor}", - f"cu{major}.{minor}", - f"cuda{major}{minor}", - f"cuda{major}.{minor}", - ) - for string in strings: - with subtests.test(string=string): - backend = cb.ComputationBackend.from_str(string) - assert isinstance(backend, cb.CUDABackend) - assert backend.major == major - assert backend.minor == minor +class TestCPUBackend: + def test_eq(self): + backend = cb.CPUBackend() + assert backend == "cpu" -def test_ComputationBackend_from_str_unknown(subtests): - strings = ("unknown", "cudnn") - for string in strings: - with subtests.test(string=string): - with pytest.raises(cb.ParseError): - cb.ComputationBackend.from_str(string) +class TestCUDABackend: + def test_eq(self): + major = 42 + minor = 21 + backend = cb.CUDABackend(major, minor) + assert backend == f"cu{major}{minor}" -def test_CPUBackend(): - backend = cb.CPUBackend() - assert backend == "cpu" +class TestOrdering: + def test_cpu(self): + assert cb.CPUBackend() < cb.CUDABackend(0, 0) + assert cb.CPUBackend() < cb.AnyBackend() + def test_any(self): + assert cb.AnyBackend() > cb.CUDABackend(99, 99) + assert cb.AnyBackend() > cb.CPUBackend() -def test_CUDABackend(): - major = 42 - minor = 21 - backend = cb.CUDABackend(major, minor) - assert backend == f"cu{major}{minor}" + def test_cuda(self): + assert cb.CUDABackend(0, 0) > cb.CPUBackend() + assert cb.CUDABackend(99, 99) < cb.AnyBackend() + assert cb.CUDABackend(1, 2) < cb.CUDABackend(2, 1) + assert cb.CUDABackend(2, 1) < cb.CUDABackend(10, 0) try: subprocess.check_call( - "nvcc --version", - shell=True, - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL, + "nvidia-smi", shell=True, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, ) - CUDA_AVAILABLE = True + NVIDIA_DRIVER_AVAILABLE = True except subprocess.CalledProcessError: - CUDA_AVAILABLE = False + NVIDIA_DRIVER_AVAILABLE = False -skip_if_cuda_unavailable = pytest.mark.skipif( - not CUDA_AVAILABLE, reason="Requires CUDA." +skip_if_nvidia_driver_unavailable = pytest.mark.skipif( + not NVIDIA_DRIVER_AVAILABLE, reason="Requires nVidia driver." ) -def test_detect_computation_backend_no_nvcc(mocker): - mocker.patch( - "light_the_torch.computation_backend.subprocess.check_output", - side_effect=subprocess.CalledProcessError(1, ""), - ) - - assert isinstance(cb.detect_computation_backend(), cb.CPUBackend) - - -def test_detect_computation_backend_unknown_release(mocker): - mocker.patch( - "light_the_torch.computation_backend.subprocess.check_output", - return_value="release unknown".encode("utf-8"), - ) - - assert isinstance(cb.detect_computation_backend(), cb.CPUBackend) - - -def test_detect_computation_backend_cuda(mocker): - major = 460 - minor = 20 - patch = 30 - mocker.patch( - "light_the_torch.computation_backend.subprocess.check_output", - return_value=f"foo\n{major}.{minor}.{patch}".encode("utf-8"), +@pytest.fixture +def patch_nvidia_driver_version(mocker): + def factory(version): + return mocker.patch( + "light_the_torch.computation_backend.subprocess.check_output", + return_value=f"driver_version\n{version}".encode("utf-8"), + ) + + return factory + + +def cuda_backends_params(): + params = [] + for system, minimum_driver_versions in cb._MINIMUM_DRIVER_VERSIONS.items(): + cuda_versions, driver_versions = zip(*sorted(minimum_driver_versions.items())) + cuda_backends = tuple( + cb.CUDABackend(version.major, version.minor) for version in cuda_versions + ) + + # latest driver supports every backend + params.append( + pytest.param( + system, + str(driver_versions[-1]), + set(cuda_backends), + id=f"{system.lower()}-latest", + ) + ) + + # outdated driver supports no backend + params.append( + pytest.param( + system, + str(driver_versions[0].major - 1), + {}, + id=f"{system.lower()}-outdated", + ) + ) + + # "normal" driver supports some backends + idx = len(cuda_versions) // 2 + params.append( + pytest.param( + system, + str(driver_versions[idx]), + set(cuda_backends[: idx + 1],), + id=f"{system.lower()}-normal", + ) + ) + + return pytest.mark.parametrize( + ("system", "nvidia_driver_version", "compatible_cuda_backends"), params ) - backend = cb.detect_computation_backend() - assert isinstance(backend, cb.CUDABackend) - assert backend.major == 11 - assert backend.minor == 1 - -@skip_if_cuda_unavailable -def test_detect_computation_backend_cuda_smoke(): - assert isinstance(cb.detect_computation_backend(), cb.CUDABackend) +class TestDetectCompatibleComputationBackends: + def test_no_nvidia_driver(self, mocker): + mocker.patch( + "light_the_torch.computation_backend.subprocess.check_output", + side_effect=subprocess.CalledProcessError(1, ""), + ) + + assert cb.detect_compatible_computation_backends() == {cb.CPUBackend()} + + @cuda_backends_params() + def test_cuda_backends( + self, + mocker, + patch_nvidia_driver_version, + system, + nvidia_driver_version, + compatible_cuda_backends, + ): + mocker.patch( + "light_the_torch.computation_backend.platform.system", return_value=system + ) + patch_nvidia_driver_version(nvidia_driver_version) + + backends = cb.detect_compatible_computation_backends() + assert backends == {cb.CPUBackend(), *compatible_cuda_backends} + + @skip_if_nvidia_driver_unavailable + def test_cuda_backend(self): + backend_types = { + type(backend) for backend in cb.detect_compatible_computation_backends() + } + assert cb.CUDABackend in backend_types