Improve the warm-up models device logic

Update the `default_devices` logic to load models onto a single device (`cuda` if available, otherwise `cpu`), instead of both devices, to streamline resource usage.

Author: wanliAlex

src/marqo/inference/native_inference/remote/server/on_start_script.py

@@ -104,7 +104,7 @@ def __init__(self, config: Config):
         # TBD to include cross-encoder/ms-marco-TinyBERT-L-2-v2
 
         # TODO [Refactoring device logic] use device info gathered from device manager
-        self.default_devices = ['cpu'] if not torch.cuda.is_available() else ['cuda', 'cpu']
+        self.default_devices = ['cpu'] if not torch.cuda.is_available() else ['cuda']
 
         self.logger.info(f"pre-loading {self.models} onto devices={self.default_devices}")
 

src/marqo/version.py

@@ -1,4 +1,4 @@
-__version__ = "2.19.0"
+__version__ = "2.19.1"
 
 def get_version() -> str:
     return f"{__version__}"

tests/unit_tests/marqo/inference/native_inference/remote/server/test_on_start_script.py

@@ -216,4 +216,33 @@ def test_missing_punkt_downloaded(self):
                 checker = on_start_script.CheckNLTKTokenizers()
                 with self.assertRaises(StartupSanityCheckError):
                     checker.run()
-                mock_nltk_download.assert_any_call("punkt_tab")
+                mock_nltk_download.assert_any_call("punkt_tab")
+
+    def test_models_only_load_to_one_device(self):
+        """
+        Ensure models are only loaded to one device (cuda if available, else cpu) when warming up,
+        not to all devices.
+        """
+        with mock.patch("marqo.inference.native_inference.remote.server.on_start_script.torch.cuda.is_available") as mock_cuda_available, \
+             mock.patch("os.environ", {
+                 enums.EnvVars.MARQO_MODELS_TO_PRELOAD: json.dumps(["LanguageBind/Video_V1.5_FT_Audio_FT_Image"])
+             }):
+
+            for cuda_available in [True, False]:
+                expected_device = "cuda" if cuda_available else "cpu"
+                mock_cuda_available.return_value = cuda_available
+
+                cache_model_module = on_start_script.CacheModels(self.mock_config)
+                self.assertEqual(cache_model_module.default_devices, [expected_device])
+
+                with mock.patch.object(cache_model_module, "_preload_model") as mock_preload_model:
+                    cache_model_module.run()
+                    mock_preload_model.assert_called_with(
+                        model="LanguageBind/Video_V1.5_FT_Audio_FT_Image",
+                        content="this is a test string",
+                        device=expected_device
+                    )
+
+                    # Ensure the other device is not used
+                    other_device = "cpu" if expected_device == "cuda" else "cuda"
+                    self.assertNotIn(other_device, mock_preload_model.call_args[1]["device"])