feat: add coordinated checkpoint prefetch for network filesystem loading (#20843)

test/registered/model_loading/test_prefetch_checkpoints_multi_gpu.py

@@ -0,0 +1,49 @@
+import unittest
+
+import sglang as sgl
+from sglang.test.ci.ci_register import register_cuda_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cuda_ci(est_time=300, suite="nightly-4-gpu")
+
+PROMPTS = [
+    "Hello, my name is",
+    "The president of the United States is",
+    "The capital of France is",
+    "The future of AI is",
+]
+
+
+class TestPrefetchCheckpointsMultiGPU(CustomTestCase):
+    """Verify that --weight-loader-prefetch-checkpoints works with DP attention."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.engine = sgl.Engine(
+            model_path="Qwen/Qwen1.5-MoE-A2.7B-Chat",
+            tp_size=4,
+            dp_size=4,
+            enable_dp_attention=True,
+            disable_radix_cache=True,
+            weight_loader_prefetch_checkpoints=True,
+            cuda_graph_max_bs=1,
+            max_total_tokens=256,
+        )
+
+    @classmethod
+    def tearDownClass(cls):
+        if hasattr(cls, "engine") and cls.engine:
+            cls.engine.shutdown()
+
+    def test_generate_with_prefetch(self):
+        """Server launched with prefetch must produce valid output."""
+        outputs = self.engine.generate(PROMPTS)
+        self.assertEqual(len(outputs), len(PROMPTS))
+        for i, output in enumerate(outputs):
+            text = output["text"]
+            self.assertIsInstance(text, str)
+            self.assertGreater(len(text), 0, f"Prompt {i} produced empty output")
+
+
+if __name__ == "__main__":
+    unittest.main()

test/registered/unit/model_loader/test_prefetch_checkpoints.py

@@ -0,0 +1,55 @@
+"""
+Unit tests for coordinated checkpoint prefetch.
+
+Verifies that weights loaded with prefetch enabled are bit-identical
+to weights loaded without prefetch.
+"""
+
+import os
+import tempfile
+import unittest
+from unittest.mock import patch
+
+import safetensors.torch
+import torch
+
+from sglang.srt.model_loader.weight_utils import (
+    safetensors_weights_iterator,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
+
+
+class TestPrefetchWeightsIdentical(unittest.TestCase):
+    """Verify that loading with prefetch yields identical weights to without."""
+
+    def _create_safetensors_files(self, tmpdir, num_shards=3):
+        """Create real safetensors files with known tensor content."""
+        paths = []
+        for i in range(num_shards):
+            tensors = {
+                f"layer{i}.weight": torch.randn(32, 32),
+                f"layer{i}.bias": torch.randn(32),
+            }
+            path = os.path.join(tmpdir, f"model-{i:05d}.safetensors")
+            safetensors.torch.save_file(tensors, path)
+            paths.append(path)
+        return paths
+
+    @patch("torch.distributed.is_initialized", return_value=False)
+    def test_weights_match_with_and_without_prefetch(self, _):
+        """Tensors yielded must be bit-identical regardless of prefetch flag."""
+        with tempfile.TemporaryDirectory() as tmpdir:
+            paths = self._create_safetensors_files(tmpdir)
+
+            without = dict(safetensors_weights_iterator(paths, prefetch=False))
+            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])
+
+
+if __name__ == "__main__":
+    unittest.main()