diff --git a/scripts/lib/semantic-retrieval.py b/scripts/lib/semantic-retrieval.py index f0ca25e..189a68a 100644 --- a/scripts/lib/semantic-retrieval.py +++ b/scripts/lib/semantic-retrieval.py @@ -10,6 +10,7 @@ import sys FILES = ('server.py', 'benchmark.py', 'requirements.txt') CUDA_INDEX = 'https://download.pytorch.org/whl/cu126' +MODEL_PRELOAD_TIMEOUT = 900 def digest(path): @@ -41,6 +42,11 @@ def ensure_runtime(target, dry_run): return python +def ensure_models(python, target): + subprocess.run([str(python), str(target / 'server.py'), 'preload', '--model-cache', str(target / 'model-cache')], + check=True, timeout=MODEL_PRELOAD_TIMEOUT) + + def load_state(path): if not path.exists(): return {'version': 1, 'clients': [], 'files': {}} @@ -96,7 +102,9 @@ def sync(root, home, clients, enabled, dry_run, update, summary=False): for name in FILES: shutil.copy2(source / name, target / name) state['files'][name] = digest(target / name) + save_state(state_path, state, False) python = ensure_runtime(target, False) + ensure_models(python, target) state['clients'] = list(dict.fromkeys(state['clients'])) save_state(state_path, state, False) for client in clients: diff --git a/scripts/tests/test-retrieval-server.py b/scripts/tests/test-retrieval-server.py index 82cab72..aea6ce7 100644 --- a/scripts/tests/test-retrieval-server.py +++ b/scripts/tests/test-retrieval-server.py @@ -82,6 +82,25 @@ class RetrievalServerTests(unittest.TestCase): with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False))): self.assertEqual(MODULE.preferred_device(), 'cpu') + def test_preload_models_warms_embedding_and_reranking_from_model_cache(self): + with TemporaryDirectory() as temporary: + cache = Path(temporary) / 'model-cache' + calls = [] + encoder = SimpleNamespace(encode=lambda values, **kwargs: calls.append(('encode', values, kwargs)) or [embedding()]) + reranker = SimpleNamespace(predict=lambda pairs, **kwargs: calls.append(('predict', pairs, kwargs)) or [0.0]) + + with patch.object(MODULE, 'shared_encoder', return_value=encoder), \ + patch.object(MODULE, 'shared_reranker', return_value=reranker), \ + patch.dict(os.environ, {'HF_HOME': 'other-cache'}): + getattr(MODULE, 'preload_models', lambda _: None)(cache) + self.assertEqual(str(cache.resolve()), os.environ['HF_HOME']) + + self.assertTrue(cache.is_dir()) + self.assertEqual([ + ('encode', ['Verify semantic retrieval models.'], {'prompt_name': 'query'}), + ('predict', [('Verify semantic retrieval models.', 'Verify semantic retrieval models.')], {'show_progress_bar': False}), + ], calls) + @unittest.skipUnless(os.environ.get('QWEN_TEST_MODEL_CACHE'), 'Set QWEN_TEST_MODEL_CACHE to run the real-model check.') def test_real_models(self): if not huggingface_available(): diff --git a/scripts/tests/test-semantic-retrieval.py b/scripts/tests/test-semantic-retrieval.py index 3ef3475..ce78718 100644 --- a/scripts/tests/test-semantic-retrieval.py +++ b/scripts/tests/test-semantic-retrieval.py @@ -48,6 +48,19 @@ class SemanticRetrievalTests(unittest.TestCase): self.assertEqual(run.call_count, 2) self.assertNotIn('--extra-index-url', run.call_args.args[0]) + def test_sync_preloads_models_after_installing_runtime(self): + with tempfile.TemporaryDirectory() as directory: + home = Path(directory) / 'home' + target = home / '.my-ai-configuration/semantic-retrieval' + python = MODULE.python_path(target / '.venv') + with patch.object(MODULE, 'ensure_runtime', return_value=python), \ + patch.object(MODULE.subprocess, 'run') as run: + MODULE.sync(ROOT, home, ['codex'], True, False, False) + + self.assertEqual([[ + str(python), str(target / 'server.py'), 'preload', '--model-cache', str(target / 'model-cache') + ]], [call.args[0] for call in run.call_args_list]) + def test_summary_dry_run_omits_details_without_writes(self): with tempfile.TemporaryDirectory() as directory: home = Path(directory) / 'home' diff --git a/shared/retrieval/server.py b/shared/retrieval/server.py index 161d1b0..3e6dd7c 100644 --- a/shared/retrieval/server.py +++ b/shared/retrieval/server.py @@ -201,6 +201,18 @@ def shared_reranker(): return _shared_reranker +def preload_models(model_cache: Path): + model_cache.mkdir(parents=True, exist_ok=True) + os.environ["HF_HOME"] = str(model_cache.resolve()) + query = "Verify semantic retrieval models." + embeddings = shared_encoder().encode([query], prompt_name="query") + if len(embeddings) != 1 or len(embeddings[0]) != EMBEDDING_DIMENSIONS: + raise ValueError("Embedding model returned unexpected dimensions.") + scores = shared_reranker().predict([(query, query)], show_progress_bar=False) + if len(scores) != 1: + raise ValueError("Reranking model returned unexpected scores.") + + def reciprocal_rank_fusion(*rankings): scores = {} for ranking in rankings: @@ -489,15 +501,20 @@ def run_direct(action: str, root: Path, data_dir: Path, model_cache: Path, query def main(): parser = argparse.ArgumentParser() - parser.add_argument("action", choices=("search", "rebuild", "daemon")) + parser.add_argument("action", choices=("search", "rebuild", "daemon", "preload")) parser.add_argument("--query") parser.add_argument("--top-k", type=int, default=FINAL_RESULTS) parser.add_argument("--root", default=os.getcwd(), type=Path) - parser.add_argument("--data-dir", required=True, type=Path) + parser.add_argument("--data-dir", type=Path) parser.add_argument("--model-cache", required=True, type=Path) args = parser.parse_args() - data_dir = args.data_dir.resolve() model_cache = args.model_cache.resolve() + if args.action == "preload": + preload_models(model_cache) + return + if args.data_dir is None: + parser.error(f"{args.action} requires --data-dir") + data_dir = args.data_dir.resolve() if args.action == "daemon": data_dir.mkdir(parents=True, exist_ok=True) run_daemon(data_dir, model_cache)