Fix bounded checkpoint prefetching and buffered drop-cache handling (#29156)

python/sglang/srt/model_loader/weight_utils.py

@@ -67,7 +67,7 @@ from sglang.utils import is_in_ci
 
 try:
     from fastsafetensors import SafeTensorsFileLoader, SingleGroup
-except ImportError as e:
+except ImportError:
     SafeTensorsFileLoader = SingleGroup = None
 
 logger = logging.getLogger(__name__)
@@ -814,10 +814,12 @@ def _prefetch_all_checkpoints(
     naturally adapts to any RAM size — even if the full checkpoint does
     not fit in page cache, the prefetch thread stays ahead of the loader.
     """
-    import asyncio
     import threading
     import time
 
+    if num_threads < 1:
+        raise ValueError("weight loader prefetch num_threads must be >= 1")
+
     # Use node-local rank so that each node independently prefetches the
     # full checkpoint into its own page cache. Global rank would split files
     # across nodes, but page cache is not shared across nodes.
@@ -842,40 +844,60 @@ def _prefetch_all_checkpoints(
         num_threads,
     )
 
-    async def _prefetch_all() -> None:
-        semaphore = asyncio.Semaphore(num_threads)
+    def _prefetch_all() -> None:
         completed = 0
         next_log_pct = 10
 
-        async def prefetch_one(path: str) -> None:
+        def record_complete() -> None:
             nonlocal completed, next_log_pct
-            try:
-                async with semaphore:
-                    await asyncio.to_thread(_prefetch_checkpoint_file, path)
-                completed += 1
-                if total_for_rank > 0 and next_log_pct <= 100:
-                    pct = 100 * completed / total_for_rank
-                    if pct >= next_log_pct:
-                        logger.info(
-                            "Rank %d: prefetching checkpoint files: %d%% (%d/%d)",
-                            local_rank,
-                            next_log_pct,
-                            completed,
-                            total_for_rank,
-                        )
-                        next_log_pct += 10
-            except Exception:
-                logger.warning(
-                    "Failed to prefetch checkpoint file %r.",
-                    path,
-                    exc_info=True,
-                )
 
-        await asyncio.gather(*(prefetch_one(p) for p in my_files))
+            completed += 1
+            if total_for_rank > 0 and next_log_pct <= 100:
+                pct = 100 * completed / total_for_rank
+                while pct >= next_log_pct and next_log_pct <= 100:
+                    logger.info(
+                        "Rank %d: prefetching checkpoint files: %d%% (%d/%d)",
+                        local_rank,
+                        next_log_pct,
+                        completed,
+                        total_for_rank,
+                    )
+                    next_log_pct += 10
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=num_threads) as executor:
+            file_iter = iter(my_files)
+            pending: Dict[concurrent.futures.Future, str] = {}
+
+            for path in itertools.islice(file_iter, num_threads):
+                pending[executor.submit(_prefetch_checkpoint_file, path)] = path
+
+            while pending:
+                done, _ = concurrent.futures.wait(
+                    pending,
+                    return_when=concurrent.futures.FIRST_COMPLETED,
+                )
+                for future in done:
+                    path = pending.pop(future)
+                    try:
+                        future.result()
+                    except Exception:
+                        logger.warning(
+                            "Failed to prefetch checkpoint file %r.",
+                            path,
+                            exc_info=True,
+                        )
+                    finally:
+                        record_complete()
+
+                    next_path = next(file_iter, None)
+                    if next_path is not None:
+                        pending[
+                            executor.submit(_prefetch_checkpoint_file, next_path)
+                        ] = next_path
 
     def _run_prefetch() -> None:
         start = time.perf_counter()
-        asyncio.run(_prefetch_all())
+        _prefetch_all()
         elapsed = time.perf_counter() - start
         logger.info(
             "Rank %d: prefetching checkpoint files into page cache "
@@ -1081,7 +1103,7 @@ def buffered_multi_thread_safetensors_weights_iterator(
 
         # Seed the buffer.
         for st_file in itertools.islice(file_iter, buffer_size):
-            pending.append(executor.submit(_load_file, st_file))
+            pending.append((st_file, executor.submit(_load_file, st_file)))
 
         with tqdm(
             total=len(hf_weights_files),
@@ -1091,18 +1113,22 @@ def buffered_multi_thread_safetensors_weights_iterator(
             position=tqdm._get_free_pos(),
         ) as pbar:
             while pending:
-                future = pending.popleft()
+                st_file, future = pending.popleft()
                 state_dict = future.result()
                 del future  # let GC reclaim the Future's internal result
 
                 # Replenish: submit the next file to keep the buffer full.
                 next_file = next(file_iter, None)
                 if next_file is not None:
-                    pending.append(executor.submit(_load_file, next_file))
+                    pending.append((next_file, executor.submit(_load_file, next_file)))
 
                 for name in sorted(state_dict.keys()):
                     yield name, state_dict[name]
                 del state_dict
+                if drop_cache_after_load:
+                    # DONTNEED reduces page-cache pressure after copying weights,
+                    # but later mmap-backed tensor access may fault pages again.
+                    _drop_file_cache_after_load(st_file)
                 pbar.update(1)
 
 

test/registered/unit/model_loader/test_prefetch_checkpoints.py

@@ -8,12 +8,15 @@ to weights loaded without prefetch.
 import os
 import tempfile
 import unittest
+from concurrent.futures import Future
 from unittest.mock import patch
 
 import safetensors.torch
 import torch
 
 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
@@ -21,8 +24,40 @@ from sglang.test.ci.ci_register import register_cpu_ci
 register_cpu_ci(est_time=10, suite="base-a-test-cpu")
 
 
-class TestPrefetchWeightsIdentical(unittest.TestCase):
-    """Verify that loading with prefetch yields identical weights to without."""
+class _InlineThread:
+    def __init__(self, target, daemon=None):
+        self.target = target
+        self.daemon = daemon
+
+    def start(self):
+        self.target()
+
+
+class _InlineExecutor:
+    def __init__(self, max_workers):
+        self.max_workers = max_workers
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc, tb):
+        return False
+
+    def submit(self, fn, *args, **kwargs):
+        future = Future()
+        try:
+            future.set_result(fn(*args, **kwargs))
+        except Exception as exc:
+            future.set_exception(exc)
+        return future
+
+
+def _wait_all(fs, return_when):
+    return set(fs), set()
+
+
+class TestPrefetchCheckpoints(unittest.TestCase):
+    """Verify coordinated checkpoint prefetch behavior."""
 
     def _create_safetensors_files(self, tmpdir, num_shards=3):
         """Create real safetensors files with known tensor content."""
@@ -44,12 +79,144 @@ class TestPrefetchWeightsIdentical(unittest.TestCase):
             paths = self._create_safetensors_files(tmpdir)
 
             without = dict(safetensors_weights_iterator(paths, prefetch=False))
-            with_pf = dict(safetensors_weights_iterator(paths, prefetch=True))
+            with (
+                patch("threading.Thread", _InlineThread),
+                patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor),
+                patch("concurrent.futures.wait", side_effect=_wait_all),
+            ):
+                with_pf = dict(safetensors_weights_iterator(paths, prefetch=True))
 
             self.assertEqual(set(without.keys()), set(with_pf.keys()))
             for name in without:
                 torch.testing.assert_close(without[name], with_pf[name])
 
+    def test_prefetch_rejects_invalid_thread_count(self):
+        with self.assertRaisesRegex(ValueError, "num_threads"):
+            _prefetch_all_checkpoints(["dummy.safetensors"], num_threads=0)
+
+    @patch("torch.distributed.is_initialized", return_value=False)
+    def test_prefetch_keeps_bounded_pending_window(self, _):
+        paths = [f"model-{i:05d}.safetensors" for i in range(20)]
+        pending_sizes = []
+        submitted_paths = []
+
+        class RecordingExecutor(_InlineExecutor):
+            def submit(self, fn, path):
+                submitted_paths.append(path)
+                return super().submit(fn, path)
+
+        def record_pending_size(fs, return_when):
+            pending_sizes.append(len(fs))
+            return _wait_all(fs, return_when)
+
+        with (
+            patch("threading.Thread", _InlineThread),
+            patch("concurrent.futures.ThreadPoolExecutor", RecordingExecutor),
+            patch("concurrent.futures.wait", side_effect=record_pending_size),
+            patch("sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file"),
+        ):
+            _prefetch_all_checkpoints(paths, num_threads=4)
+
+        self.assertEqual(submitted_paths, paths)
+        self.assertLessEqual(max(pending_sizes), 4)
+
+    @patch("torch.distributed.is_initialized", return_value=False)
+    def test_prefetch_logs_failed_futures(self, _):
+        paths = ["bad.safetensors"]
+
+        def fail_prefetch(path):
+            raise OSError(f"failed {path}")
+
+        with (
+            patch("threading.Thread", _InlineThread),
+            patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor),
+            patch("concurrent.futures.wait", side_effect=_wait_all),
+            patch(
+                "sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file",
+                side_effect=fail_prefetch,
+            ),
+            patch("sglang.srt.model_loader.weight_utils.logger.warning") as warning,
+        ):
+            _prefetch_all_checkpoints(paths, num_threads=1)
+
+        warning.assert_called_once()
+        self.assertEqual(
+            warning.call_args.args[0],
+            "Failed to prefetch checkpoint file %r.",
+        )
+        self.assertEqual(warning.call_args.args[1], paths[0])
+        self.assertTrue(warning.call_args.kwargs["exc_info"])
+
+    @patch("torch.distributed.is_initialized", return_value=False)
+    def test_prefetch_progress_logs_all_crossed_buckets(self, _):
+        paths = [f"model-{i:05d}.safetensors" for i in range(3)]
+
+        with (
+            patch("threading.Thread", _InlineThread),
+            patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor),
+            patch("concurrent.futures.wait", side_effect=_wait_all),
+            patch("sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file"),
+            patch("sglang.srt.model_loader.weight_utils.logger.info") as log_info,
+        ):
+            _prefetch_all_checkpoints(paths, num_threads=1)
+
+        progress_pcts = [
+            call.args[2]
+            for call in log_info.call_args_list
+            if call.args
+            and call.args[0] == "Rank %d: prefetching checkpoint files: %d%% (%d/%d)"
+        ]
+        self.assertEqual(progress_pcts, list(range(10, 101, 10)))
+
+    @patch("torch.distributed.is_initialized", return_value=True)
+    def test_prefetch_uses_node_local_rank_partitioning(self, _):
+        paths = [f"model-{i:05d}.safetensors" for i in range(10)]
+        loaded_paths = []
+
+        class FakeWorldGroup:
+            local_rank = 1
+            local_size = 3
+            world_size = 99
+
+        with (
+            patch("threading.Thread", _InlineThread),
+            patch("concurrent.futures.ThreadPoolExecutor", _InlineExecutor),
+            patch("concurrent.futures.wait", side_effect=_wait_all),
+            patch(
+                "sglang.srt.model_loader.weight_utils.get_world_group",
+                return_value=FakeWorldGroup(),
+            ),
+            patch(
+                "sglang.srt.model_loader.weight_utils._prefetch_checkpoint_file",
+                side_effect=loaded_paths.append,
+            ),
+        ):
+            _prefetch_all_checkpoints(paths, num_threads=2)
+
+        self.assertEqual(sorted(loaded_paths), sorted(paths[1::3]))
+
+    @patch("torch.distributed.is_initialized", return_value=False)
+    def test_buffered_loader_drops_cache_after_each_loaded_shard(self, _):
+        with tempfile.TemporaryDirectory() as tmpdir:
+            paths = self._create_safetensors_files(tmpdir, num_shards=3)
+
+            with patch(
+                "sglang.srt.model_loader.weight_utils._drop_file_cache_after_load"
+            ) as drop_cache:
+                loaded = list(
+                    buffered_multi_thread_safetensors_weights_iterator(
+                        paths,
+                        max_workers=2,
+                        drop_cache_after_load=True,
+                    )
+                )
+
+            self.assertEqual(len(loaded), 6)
+            self.assertEqual(
+                [call.args[0] for call in drop_cache.call_args_list],
+                paths,
+            )
+
 
 if __name__ == "__main__":
     unittest.main()