feat(retrieval): prefer CUDA devices

This commit is contained in:
Julian lechner committed 2026-09-17 17:50:33 +02:00
1 parent 5ec6d0b6ac
commit e5b63965b8
2 files changed
+15 -2

No files matched your search

+6
View File
@@ -27,6 +27,12 @@ def embedding():
class RetrievalServerTests(unittest.TestCase):
def test_preferred_device_uses_cuda_with_cpu_fallback(self):
with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: True))):
self.assertEqual(MODULE.preferred_device(), 'cuda')
with patch.dict('sys.modules', torch=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False))):
self.assertEqual(MODULE.preferred_device(), 'cpu')
@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):
os.environ['HF_HOME'] = os.environ['QWEN_TEST_MODEL_CACHE']