Install torch nightly from wheel URLs

Author: nevillelyh

src/monobase/monogen.py

@@ -28,11 +28,8 @@ class MonoGen:
         python={'3.12': '3.12.7'},
         torch=[
             '2.4.1',
-            # NOTE(meatballhat): This is turned off until we can figure out how to handle
-            # nightlies better since the torch package index only retains ~2 months of
-            # versions:
-            ## Nightly
-            #'2.6.0.dev20240918'
+            # Nightly
+            '2.6.0.dev20240918',
         ],
         pip_pkgs=SEED_PKGS,
     ),

src/monobase/torch.py

@@ -30,11 +30,8 @@ class TorchDeps:
 }
 
 torch_deps_dict = {
-    # NOTE(meatballhat): This is turned off until we can figure out how to handle
-    # nightlies better since the torch package index only retains ~2 months of
-    # versions:
-    ## Nightly
-    #'2.6.0.dev20240918': TorchDeps('2.5.0.dev20240918', '0.20.0.dev20240918'),
+    # Nightly
+    '2.6.0.dev20240918': TorchDeps('2.5.0.dev20240918', '0.20.0.dev20240918'),
     # Releases
     '2.4.1': TorchDeps('2.4.1', '0.19.1'),
     '2.4.0': TorchDeps('2.4.0', '0.19.0'),

src/monobase/uv.py

@@ -34,15 +34,24 @@ def index_args(torch_version: Version, cuda_version: str) -> list[str]:
 
 
 def pip_packages(
-    torch_version: Version, cuda_version: str, pip_pkgs: list[str]
+    torch_version: Version, python_version: str, cuda_version: str, pip_pkgs: list[str]
 ) -> list[str]:
     deps = torch_deps[torch_version]
     cu = cuda_suffix(cuda_version)
-    pkgs = [
-        f'torch=={torch_version}+{cu}',
-        f'torchaudio=={deps.torchaudio}+{cu}',
-        f'torchvision=={deps.torchvision}+{cu}',
-    ]
+    if torch_version.extra:
+        prefix = torch_index_url(torch_version, cuda_version)
+        py = f'cp{python_version.replace('.', '')}'
+        pkgs = [
+            f'torch @ {prefix}/torch-{torch_version}%2B{cu}-{py}-{py}-linux_x86_64.whl',
+            f'torchaudio @ {prefix}/torchaudio-{deps.torchaudio}%2B{cu}-{py}-{py}-linux_x86_64.whl',
+            f'torchvision @ {prefix}/torchvision-{deps.torchvision}%2B{cu}-{py}-{py}-linux_x86_64.whl',
+        ]
+    else:
+        pkgs = [
+            f'torch=={torch_version}+{cu}',
+            f'torchaudio=={deps.torchaudio}+{cu}',
+            f'torchvision=={deps.torchvision}+{cu}',
+        ]
     # Older Torch versions do not bundle CUDA or CuDNN
     nvidia_pkgs = []
     if torch_version < Version.parse('2.2.0'):
@@ -101,7 +110,7 @@ def update_venv(
     ]
     cmd = ['uv', 'pip', 'compile', '--python-platform', 'x86_64-unknown-linux-gnu']
     cmd = cmd + emit_args + index_args(t, cuda_version) + ['-']
-    pkgs = pip_packages(t, cuda_version, pip_pkgs)
+    pkgs = pip_packages(t, python_version, cuda_version, pip_pkgs)
     env = os.environ.copy()
     env['VIRTUAL_ENV'] = vdir
     try: