feat(retrieval): preload Qwen models during setup

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:

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():

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'

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)