Disable multi-threaded load by default when prefetch is on (#30146)

test/registered/unit/model_loader/test_prefetch_checkpoints.py

@@ -9,17 +9,21 @@ import os
 import tempfile
 import unittest
 from concurrent.futures import Future
+from types import SimpleNamespace
 from unittest.mock import patch
 
 import safetensors.torch
 import torch
 
+from sglang.srt.configs.load_config import LoadConfig, LoadFormat
+from sglang.srt.model_loader.loader import DefaultModelLoader
 from sglang.srt.model_loader.weight_utils import (
     _prefetch_all_checkpoints,
     buffered_multi_thread_safetensors_weights_iterator,
     safetensors_weights_iterator,
 )
 from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
 
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
@@ -56,7 +60,7 @@ def _wait_all(fs, return_when):
     return set(fs), set()
 
 
-class TestPrefetchCheckpoints(unittest.TestCase):
+class TestPrefetchCheckpoints(CustomTestCase):
     """Verify coordinated checkpoint prefetch behavior."""
 
     def _create_safetensors_files(self, tmpdir, num_shards=3):
@@ -218,5 +222,192 @@ class TestPrefetchCheckpoints(unittest.TestCase):
             )
 
 
+class TestPrefetchDispatch(CustomTestCase):
+    """Verify _get_weights_iterator dispatches to the right safetensors
+    iterator based on prefetch / multi-thread config.
+
+    Prefetch + default (multi-thread on) must fall back to the
+    single-threaded iterator; an explicit enable_multithread_load=true is
+    honored as an opt-out; FASTSAFETENSORS and disable_mmap bypass the
+    override.
+    """
+
+    def _make_loader(self, extra_config, load_format=LoadFormat.SAFETENSORS):
+        load_config = LoadConfig(
+            load_format=load_format,
+            model_loader_extra_config=extra_config,
+        )
+        return DefaultModelLoader(load_config)
+
+    def _make_source(self):
+        # model_config=None skips maybe_add_mtp_safetensors.
+        return SimpleNamespace(
+            model_or_path="/dummy",
+            revision=None,
+            fall_back_to_pt=False,
+            model_config=None,
+            prefix="",
+        )
+
+    def _server_args(self, prefetch, disable_mmap=False):
+        return SimpleNamespace(
+            weight_loader_disable_mmap=disable_mmap,
+            weight_loader_prefetch_checkpoints=prefetch,
+            weight_loader_prefetch_num_threads=4,
+            weight_loader_drop_cache_after_load=False,
+        )
+
+    def _run(self, loader):
+        # _get_weights_iterator returns a generator wrapping the chosen
+        # iterator; consuming it forces the eager dispatch (the if/elif/else
+        # that calls the iterator factory) to execute.
+        list(loader._get_weights_iterator(self._make_source()))
+
+    def _patch_dispatch(self, prefetch, disable_mmap=False):
+        return (
+            patch.object(
+                DefaultModelLoader,
+                "_prepare_weights",
+                return_value=("/dummy", ["f.safetensors"], True),
+            ),
+            patch(
+                "sglang.srt.model_loader.loader.get_global_server_args",
+                return_value=self._server_args(prefetch, disable_mmap),
+            ),
+            patch(
+                "sglang.srt.model_loader.loader."
+                "buffered_multi_thread_safetensors_weights_iterator",
+                return_value=iter([]),
+            ),
+            patch(
+                "sglang.srt.model_loader.loader.safetensors_weights_iterator",
+                return_value=iter([]),
+            ),
+            patch("sglang.srt.model_loader.loader.logger.warning"),
+        )
+
+    def test_prefetch_uses_single_thread_for_default_config(self):
+        """Prefetch on + no explicit multithread config -> single-threaded,
+        and the opt-out warning fires once."""
+        loader = self._make_loader({})
+        p_prep, p_args, p_buffered, p_single, p_warn = self._patch_dispatch(
+            prefetch=True
+        )
+        with (
+            p_prep,
+            p_args,
+            p_buffered as mock_buffered,
+            p_single as mock_single,
+            p_warn as mock_warning,
+        ):
+            self._run(loader)
+        mock_single.assert_called_once()
+        mock_buffered.assert_not_called()
+        mock_warning.assert_called_once()
+
+    def test_explicit_enable_multithread_keeps_buffered_with_prefetch(self):
+        """Explicit enable_multithread_load=true is the escape hatch; the
+        override and its warning must not fire."""
+        loader = self._make_loader({"enable_multithread_load": True})
+        p_prep, p_args, p_buffered, p_single, p_warn = self._patch_dispatch(
+            prefetch=True
+        )
+        with (
+            p_prep,
+            p_args,
+            p_buffered as mock_buffered,
+            p_single as mock_single,
+            p_warn as mock_warning,
+        ):
+            self._run(loader)
+        mock_buffered.assert_called_once()
+        mock_single.assert_not_called()
+        mock_warning.assert_not_called()
+
+    def test_num_threads_only_keeps_buffered_with_prefetch(self):
+        """num_threads alone (relying on the enable_multithread_load=True
+        default) also signals multi-thread intent, so the override must not
+        fire and num_threads stays live."""
+        loader = self._make_loader({"num_threads": 64})
+        p_prep, p_args, p_buffered, p_single, p_warn = self._patch_dispatch(
+            prefetch=True
+        )
+        with (
+            p_prep,
+            p_args,
+            p_buffered as mock_buffered,
+            p_single as mock_single,
+            p_warn as mock_warning,
+        ):
+            self._run(loader)
+        mock_buffered.assert_called_once()
+        # num_threads is forwarded as max_workers to the buffered iterator.
+        self.assertEqual(mock_buffered.call_args.kwargs["max_workers"], 64)
+        mock_single.assert_not_called()
+        mock_warning.assert_not_called()
+
+    def test_no_prefetch_uses_multithread(self):
+        """Prefetch off -> multi-threaded iterator is used (default), no
+        override warning."""
+        loader = self._make_loader({})
+        p_prep, p_args, p_buffered, p_single, p_warn = self._patch_dispatch(
+            prefetch=False
+        )
+        with (
+            p_prep,
+            p_args,
+            p_buffered as mock_buffered,
+            p_single as mock_single,
+            p_warn as mock_warning,
+        ):
+            self._run(loader)
+        mock_buffered.assert_called_once()
+        mock_single.assert_not_called()
+        mock_warning.assert_not_called()
+
+    def test_prefetch_does_not_override_when_mmap_disabled(self):
+        """Prefetch is a no-op without mmap, so the override and its warning
+        must not fire."""
+        loader = self._make_loader({})
+        p_prep, p_args, p_buffered, p_single, p_warn = self._patch_dispatch(
+            prefetch=True, disable_mmap=True
+        )
+        with (
+            p_prep,
+            p_args,
+            p_buffered as mock_buffered,
+            p_single as mock_single,
+            p_warn as mock_warning,
+        ):
+            self._run(loader)
+        mock_buffered.assert_called_once()
+        mock_single.assert_not_called()
+        mock_warning.assert_not_called()
+
+    def test_prefetch_does_not_override_for_fastsafetensors(self):
+        """FASTSAFETENSORS ignores both flags; override + warning must not
+        fire."""
+        loader = self._make_loader({}, load_format=LoadFormat.FASTSAFETENSORS)
+        p_prep, p_args, p_buffered, p_single, p_warn = self._patch_dispatch(
+            prefetch=True
+        )
+        with (
+            patch(
+                "sglang.srt.model_loader.loader.fastsafetensors_weights_iterator",
+                return_value=iter([]),
+            ) as mock_fast,
+            p_prep,
+            p_args,
+            p_buffered as mock_buffered,
+            p_single as mock_single,
+            p_warn as mock_warning,
+        ):
+            self._run(loader)
+        mock_fast.assert_called_once()
+        mock_buffered.assert_not_called()
+        mock_single.assert_not_called()
+        mock_warning.assert_not_called()
+
+
 if __name__ == "__main__":
     unittest.main()