feat(retrieval): install CUDA runtime

scripts/lib/semantic-retrieval.py

@@ -10,6 +10,7 @@ import sys
 
 NAME = 'my-ai-qwen3-retrieval'
 FILES = ('server.py', 'requirements.txt')
+CUDA_INDEX = 'https://download.pytorch.org/whl/cu126'
 
 
 def digest(path):
@@ -43,8 +44,15 @@ def ensure_runtime(target, dry_run):
     python = python_path(venv)
     if not python.is_file():
         raise ValueError('Semantic retrieval virtual environment is incomplete.')
-    subprocess.run([str(python), '-m', 'pip', 'install', '--disable-pip-version-check', '-r', str(target / 'requirements.txt')],
-                   check=True, timeout=900)
+    arguments = [str(python), '-m', 'pip', 'install', '--disable-pip-version-check', '-r', str(target / 'requirements.txt')]
+    if shutil.which('nvidia-smi'):
+        try:
+            subprocess.run([*arguments, '--extra-index-url', CUDA_INDEX], check=True, timeout=900)
+        except subprocess.CalledProcessError:
+            print('Semantic retrieval: CUDA runtime installation failed; falling back to CPU.', file=sys.stderr)
+            subprocess.run(arguments, check=True, timeout=900)
+    else:
+        subprocess.run(arguments, check=True, timeout=900)
     return python
 
 

scripts/test-semantic-retrieval.py

@@ -6,6 +6,7 @@ from contextlib import redirect_stdout
 from pathlib import Path
 import tempfile
 import unittest
+from unittest.mock import patch
 
 
 ROOT = Path(__file__).resolve().parent.parent
@@ -15,6 +16,30 @@ SPEC.loader.exec_module(MODULE)
 
 
 class SemanticRetrievalTests(unittest.TestCase):
+    def test_runtime_prefers_cuda_packages_for_nvidia_gpu(self):
+        with tempfile.TemporaryDirectory() as directory:
+            target = Path(directory)
+            python = MODULE.python_path(target / '.venv')
+            python.parent.mkdir(parents=True)
+            python.touch()
+            with patch.object(MODULE.shutil, 'which', return_value='nvidia-smi'), patch.object(MODULE.subprocess, 'run') as run:
+                MODULE.ensure_runtime(target, False)
+
+            self.assertEqual(run.call_args.args[0][-2:], ['--extra-index-url', MODULE.CUDA_INDEX])
+
+    def test_runtime_falls_back_to_cpu_packages_when_cuda_install_fails(self):
+        with tempfile.TemporaryDirectory() as directory:
+            target = Path(directory)
+            python = MODULE.python_path(target / '.venv')
+            python.parent.mkdir(parents=True)
+            python.touch()
+            failure = MODULE.subprocess.CalledProcessError(1, 'pip')
+            with patch.object(MODULE.shutil, 'which', return_value='nvidia-smi'), patch.object(MODULE.subprocess, 'run', side_effect=[failure, None]) as run:
+                MODULE.ensure_runtime(target, False)
+
+            self.assertEqual(run.call_count, 2)
+            self.assertNotIn('--extra-index-url', run.call_args.args[0])
+
     def test_summary_dry_run_omits_details_without_writes(self):
         with tempfile.TemporaryDirectory() as directory:
             home = Path(directory) / 'home'