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'