add the ability to select PyTorch wheel channel (#27)

Author: pmeier

light_the_torch/_pip/find.py

@@ -30,6 +30,7 @@
 def find_links(
     pip_install_args: List[str],
     computation_backend: Optional[Union[str, ComputationBackend]] = None,
+    channel: str = "stable",
     platform: Optional[str] = None,
     python_version: Optional[str] = None,
     verbose: bool = False,
@@ -43,6 +44,8 @@ def find_links(
         computation_backend: Computation backend, for example ``"cpu"`` or ``"cu102"``.
             Defaults to the available hardware of the running system preferring CUDA
             over CPU.
+        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
             to the platform of the running system.
         python_version: Python version, for example ``"3"`` or ``"3.7"``. Defaults to
@@ -57,9 +60,17 @@ def find_links(
     elif isinstance(computation_backend, str):
         computation_backend = ComputationBackend.from_str(computation_backend)
 
+    if channel not in ("stable", "test", "nightly"):
+        raise ValueError(
+            f"channel can be one of 'stable', 'test', or 'nightly', "
+            f"but got {channel} instead."
+        )
+
     dists = extract_dists(pip_install_args)
 
-    cmd = StopAfterPytorchLinksFoundCommand(computation_backend=computation_backend)
+    cmd = StopAfterPytorchLinksFoundCommand(
+        computation_backend=computation_backend, channel=channel
+    )
     pip_install_args = adjust_pip_install_args(dists, platform, python_version)
     options, args = cmd.parser.parse_args(pip_install_args)
     try:
@@ -222,32 +233,55 @@ class PytorchLinkCollector(LinkCollector):
     def __init__(
         self,
         *args: Any,
-        url: str = "https://download.pytorch.org/whl/torch_stable.html",
+        computation_backend: ComputationBackend,
+        channel: str = "stable",
         **kwargs: Any,
     ) -> None:
         super().__init__(*args, **kwargs)
+        if channel == "stable":
+            url = "https://download.pytorch.org/whl/torch_stable.html"
+        else:
+            url = (
+                f"https://download.pytorch.org/whl/"
+                f"{channel}/{computation_backend}/torch_{channel}.html"
+            )
         self.search_scope = SearchScope.create(find_links=[url], index_urls=[])
 
     @classmethod
     def from_link_collector(
-        cls, link_collector: LinkCollector
+        cls,
+        link_collector: LinkCollector,
+        computation_backend: ComputationBackend,
+        channel: str = "stable",
     ) -> "PytorchLinkCollector":
-        return new_from_similar(cls, link_collector, ("session", "search_scope",),)
+        return new_from_similar(
+            cls,
+            link_collector,
+            ("session", "search_scope",),
+            channel=channel,
+            computation_backend=computation_backend,
+        )
 
 
 class PytorchPackageFinder(PackageFinder):
     _candidate_prefs: PytorchCandidatePreferences
     _link_collector: PytorchLinkCollector
 
     def __init__(
-        self, *args: Any, computation_backend: ComputationBackend, **kwargs: Any,
+        self,
+        *args: Any,
+        computation_backend: 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._link_collector = PytorchLinkCollector.from_link_collector(
-            self._link_collector
+            self._link_collector,
+            channel=channel,
+            computation_backend=computation_backend,
         )
 
     def make_candidate_evaluator(
@@ -265,7 +299,10 @@ def make_link_evaluator(self, *args: Any, **kwargs: Any) -> PytorchLinkEvaluator
 
     @classmethod
     def from_package_finder(
-        cls, package_finder: PackageFinder, computation_backend: ComputationBackend,
+        cls,
+        package_finder: PackageFinder,
+        computation_backend: ComputationBackend,
+        channel: str = "stable",
     ) -> "PytorchPackageFinder":
         return new_from_similar(
             cls,
@@ -279,6 +316,7 @@ def from_package_finder(
                 "ignore_requires_python",
             ),
             computation_backend=computation_backend,
+            channel=channel,
         )
 
 
@@ -298,14 +336,22 @@ def resolve(
 
 
 class StopAfterPytorchLinksFoundCommand(PatchedInstallCommand):
-    def __init__(self, computation_backend: ComputationBackend, **kwargs: Any) -> None:
+    def __init__(
+        self,
+        computation_backend: ComputationBackend,
+        channel: str = "stable",
+        **kwargs: Any,
+    ) -> None:
         super().__init__(**kwargs)
         self.computation_backend = computation_backend
+        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
+            package_finder,
+            computation_backend=self.computation_backend,
+            channel=self.channel,
         )
 
     def make_resolver(

light_the_torch/cli/__init__.py

@@ -111,6 +111,15 @@ def add_ltt_install_parser(subparsers: SubParsers) -> None:
         action="store_true",
         help="install only PyTorch distributions",
     )
+    parser.add_argument(
+        "--channel",
+        type=str,
+        default="stable",
+        help=(
+            "Channel of the PyTorch wheels. "
+            "Can be one of 'stable' (default), 'test', or 'nightly'"
+        ),
+    )
     parser.add_argument(
         "--install-cmd",
         type=str,
@@ -143,6 +152,15 @@ def add_ltt_find_parser(subparsers: SubParsers) -> None:
             "preferring CUDA over CPU."
         ),
     )
+    parser.add_argument(
+        "--channel",
+        type=str,
+        default="stable",
+        help=(
+            "Channel of the PyTorch wheels. "
+            "Can be one of 'stable' (default), 'test', or 'nightly'"
+        ),
+    )
     add_pip_install_arguments(parser, "platform", "python_version")
     LTTParser.add_common_arguments(parser)
 

light_the_torch/cli/commands.py

@@ -54,6 +54,7 @@ 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
+        self.channel = args.channel
         self.platform = args.platform
         self.python_version = args.python_version
         self.verbose = args.verbose
@@ -62,6 +63,7 @@ def _run(self, pip_install_args: List[str]) -> None:
         links = ltt.find_links(
             pip_install_args,
             computation_backend=self.computation_backend,
+            channel=self.channel,
             platform=self.platform,
             python_version=self.python_version,
             verbose=self.verbose,
@@ -73,6 +75,7 @@ class InstallCommand(Command):
     def __init__(self, args: argparse.Namespace) -> None:
         self.force_cpu = args.force_cpu
         self.pytorch_only = args.pytorch_only
+        self.channel = args.channel
 
         install_cmd = args.install_cmd
         if "{packages}" not in install_cmd:
@@ -85,6 +88,7 @@ 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,
+            channel=self.channel,
             verbose=self.verbose,
         )
         cmd = self.install_cmd.format(packages=" ".join(links))

tests/cli/test_install.py

@@ -92,6 +92,24 @@ def test_ltt_install_pytorch_only(
     collect_packages.assert_not_called()
 
 
+def test_ltt_install_channel(
+    patch_install_argv, patch_find_links, patch_subprocess_call, patch_collect_packages,
+):
+    channel = "channel"
+
+    patch_install_argv(f"--channel={channel}")
+    find_links = patch_find_links()
+    patch_subprocess_call()
+    patch_collect_packages()
+
+    with exits():
+        cli.main()
+
+    _, kwargs = find_links.call_args
+    assert "channel" in kwargs
+    assert kwargs["channel"] == channel
+
+
 def test_ltt_install_install_cmd(
     patch_install_argv, patch_find_links, patch_subprocess_call,
 ):

tests/unit/pip/test_find.py

@@ -33,6 +33,11 @@ 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")
@@ -119,6 +124,11 @@ def test_find_links_computation_backend_str(
             )
 
 
+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()
@@ -192,3 +202,12 @@ def test_find_links_torchvision_smoke(subtests, wheel_properties):
     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)