diff --git a/scripts/lib/semantic-retrieval.py b/scripts/lib/semantic-retrieval.py index 886668a..1828cb0 100644 --- a/scripts/lib/semantic-retrieval.py +++ b/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 diff --git a/scripts/test-semantic-retrieval.py b/scripts/test-semantic-retrieval.py index a0bca46..e4462bf 100644 --- a/scripts/test-semantic-retrieval.py +++ b/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'