Fix RunAI object-storage checkpoint index filtering (#38988)

python/sglang/srt/model_loader/loader.py

@@ -4118,7 +4118,11 @@ class RunaiModelStreamerLoader(BaseModelLoader):
         """Prepare weights for the model.
 
         If the model is not local, it will be downloaded."""
-        from sglang.srt.utils.runai_utils import is_runai_obj_uri, list_safetensors
+        from sglang.srt.utils.runai_utils import (
+            ObjectStorageModel,
+            is_runai_obj_uri,
+            list_safetensors,
+        )
 
         is_object_storage_path = is_runai_obj_uri(model_name_or_path)
         if self._is_distributed is None:
@@ -4160,6 +4164,10 @@ class RunaiModelStreamerLoader(BaseModelLoader):
                 self.load_config.download_dir,
                 revision,
             )
+        if is_object_storage_path:
+            index_file = os.path.abspath(
+                os.path.join(ObjectStorageModel.get_path(hf_folder), index_file)
+            )
         hf_weights_files = filter_duplicate_safetensors_files(
             hf_weights_files, hf_folder, index_file
         )

python/sglang/srt/model_loader/weight_utils.py

@@ -830,6 +830,9 @@ def filter_duplicate_safetensors_files(
             if any(fnmatch.fnmatch(rel_path, pattern) for pattern in allow_patterns):
                 files_to_validate.add(f)
 
+    if "://" in hf_folder:
+        missing_files = sorted(files_to_validate.difference(hf_weights_files))
+    else:
         missing_files = sorted(f for f in files_to_validate if not os.path.isfile(f))
     if missing_files:
         raise RuntimeError(
@@ -876,7 +879,16 @@ def maybe_add_mtp_safetensors(
 
     # Check if mtp.safetensors exists and is not already in the file list
     mtp_path = os.path.join(hf_folder, "mtp.safetensors")
-    if not os.path.isfile(mtp_path) or mtp_path in hf_weights_files:
+    if mtp_path in hf_weights_files:
+        return hf_weights_files
+
+    from sglang.srt.utils.runai_utils import is_runai_obj_uri, list_safetensors
+
+    if is_runai_obj_uri(hf_folder):
+        mtp_exists = mtp_path in list_safetensors(hf_folder)
+    else:
+        mtp_exists = os.path.isfile(mtp_path)
+    if not mtp_exists:
         return hf_weights_files
 
     # mtp.safetensors exists but not in index - this is a bug

test/registered/unit/model_loader/test_runai_model_streamer_loader.py

@@ -1,5 +1,8 @@
 import concurrent.futures
+import json
+import os
 import sys
+import tempfile
 import threading
 import unittest
 from types import SimpleNamespace
@@ -20,6 +23,7 @@ from sglang.srt.models.deepseek_v4 import (
     _dequant_fp8_wo_a_streaming,
 )
 from sglang.srt.models.deepseek_v4_dspark import DeepseekV4ForCausalLMDSpark
+from sglang.srt.utils import runai_utils
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -32,6 +36,80 @@ class _FakeModel:
 
 
 class TestRunaiModelStreamerLoader(CustomTestCase):
+    def test_remote_checkpoint_index(self):
+        for scheme in ("s3", "gs", "az"):
+            for relative_cache in (False, True):
+                with (
+                    self.subTest(scheme=scheme, relative_cache=relative_cache),
+                    tempfile.TemporaryDirectory() as cache_dir,
+                ):
+                    model_path = f"{scheme}://bucket/model"
+                    shard = "model-00001-of-00001.safetensors"
+                    files = [
+                        f"{model_path}/{shard}",
+                        f"{model_path}/mtp.safetensors",
+                        f"{model_path}/unused.safetensors",
+                    ]
+                    cache_root = (
+                        os.path.relpath(cache_dir) if relative_cache else cache_dir
+                    )
+                    with (
+                        patch.object(
+                            runai_utils.envs.SGLANG_CACHE_DIR,
+                            "get",
+                            return_value=cache_root,
+                        ),
+                        patch.object(loader_mod, "get_server_args", return_value=None),
+                        patch.object(
+                            runai_utils, "list_safetensors", return_value=files
+                        ) as list_safetensors,
+                        patch.object(
+                            weight_utils,
+                            "runai_safetensors_weights_iterator",
+                            return_value=iter(()),
+                        ) as streamer,
+                    ):
+                        metadata = runai_utils.ObjectStorageModel.get_path(model_path)
+                        os.makedirs(metadata, exist_ok=True)
+                        loader = loader_mod.RunaiModelStreamerLoader(
+                            LoadConfig(load_format=LoadFormat.RUNAI_STREAMER)
+                        )
+                        self.assertEqual(
+                            loader._prepare_weights(model_path, revision=None),
+                            (model_path, files),
+                        )
+                        list_safetensors.assert_called_with(path=model_path)
+
+                        with open(
+                            os.path.join(metadata, "model.safetensors.index.json"), "w"
+                        ) as f:
+                            json.dump({"weight_map": {"weight": shard}}, f)
+                        self.assertEqual(
+                            loader._prepare_weights(model_path, revision=None),
+                            (model_path, [files[0]]),
+                        )
+
+                        loader.target_device_str = "cpu"
+                        source = loader_mod.RunaiModelStreamerLoader.Source(
+                            model_or_path=model_path,
+                            revision=None,
+                            model_config=cast(
+                                ModelConfig,
+                                SimpleNamespace(
+                                    hf_config=SimpleNamespace(
+                                        architectures=["Glm4MoeForCausalLMNextN"],
+                                        num_nextn_predict_layers=1,
+                                    )
+                                ),
+                            ),
+                        )
+                        list(loader._get_weights_iterator(source))
+                        streamer.assert_called_once_with(files[:2], True, "cpu")
+
+                        list_safetensors.return_value = [files[-1]]
+                        with self.assertRaisesRegex(RuntimeError, shard):
+                            loader._prepare_weights(model_path, revision=None)
+
     def test_passes_quant_config_to_model_init(self):
         quant_config = object()
         fake_model = _FakeModel()

test/registered/unit/model_loader/test_weight_utils.py

@@ -4,8 +4,14 @@ import json
 import os
 import tempfile
 import unittest
+from types import SimpleNamespace
+from unittest.mock import patch
 
-from sglang.srt.model_loader.weight_utils import filter_duplicate_safetensors_files
+from sglang.srt.model_loader.weight_utils import (
+    filter_duplicate_safetensors_files,
+    maybe_add_mtp_safetensors,
+)
+from sglang.srt.utils import runai_utils
 from sglang.test.ci.ci_register import register_cpu_ci
 from sglang.test.test_utils import CustomTestCase
 
@@ -85,14 +91,16 @@ class TestFilterDuplicateSafetensorsFiles(CustomTestCase):
                 "vit": "vision_encoder/model.safetensors",
             },
         )
-        transformer = _touch(
+        shard = "transformer/diffusion_pytorch_model.safetensors"
-            self.folder, "transformer/diffusion_pytorch_model.safetensors"
+        _touch(self.folder, shard)
-        )
 
+        for folder in (self.folder, "s3://bucket/model"):
+            with self.subTest(folder=folder):
+                transformer = os.path.join(folder, shard)
                 result = filter_duplicate_safetensors_files(
                     hf_weights_files=[transformer],
-            hf_folder=self.folder,
+                    hf_folder=folder,
-            index_file=INDEX_NAME,
+                    index_file=os.path.join(self.folder, INDEX_NAME),
                     allow_patterns=["transformer/*.safetensors"],
                 )
                 self.assertEqual(result, [transformer])
@@ -106,18 +114,34 @@ class TestFilterDuplicateSafetensorsFiles(CustomTestCase):
                 "vit": "vision_encoder/model.safetensors",
             },
         )
-        transformer = _touch(
+        shard = "transformer/model-00001-of-00002.safetensors"
-            self.folder, "transformer/model-00001-of-00002.safetensors"
+        _touch(self.folder, shard)
-        )
 
-        with self.assertRaises(RuntimeError) as cm:
+        for folder in (self.folder, "s3://bucket/model"):
+            with (
+                self.subTest(folder=folder),
+                self.assertRaisesRegex(
+                    RuntimeError, r"model-00002-of-00002\.safetensors"
+                ),
+            ):
                 filter_duplicate_safetensors_files(
-                hf_weights_files=[transformer],
+                    hf_weights_files=[os.path.join(folder, shard)],
-                hf_folder=self.folder,
+                    hf_folder=folder,
-                index_file=INDEX_NAME,
+                    index_file=os.path.join(self.folder, INDEX_NAME),
                     allow_patterns=["transformer/*.safetensors"],
                 )
-        self.assertIn("model-00002-of-00002.safetensors", str(cm.exception))
+
+    def test_local_index_allows_subset_of_existing_shards(self):
+        _write_index(
+            self.folder,
+            {"w1": "first.safetensors", "w2": "second.safetensors"},
+        )
+        first = _touch(self.folder, "first.safetensors")
+        _touch(self.folder, "second.safetensors")
+
+        result = filter_duplicate_safetensors_files([first], self.folder, INDEX_NAME)
+
+        self.assertEqual(result, [first])
 
     def test_single_file_model_no_index_returns_unchanged(self):
         # No index on disk (single-file / dummy / object-storage): early return.
@@ -131,5 +155,49 @@ class TestFilterDuplicateSafetensorsFiles(CustomTestCase):
         self.assertEqual(result, [single])
 
 
+class TestMaybeAddMtpSafetensors(CustomTestCase):
+    def test_remote_mtp_guards(self):
+        folder = "s3://bucket/model"
+        model = f"{folder}/model.safetensors"
+        mtp = f"{folder}/mtp.safetensors"
+        cases = (
+            ("Glm4MoeForCausalLM", 1, [model, mtp], [model, mtp], False),
+            ("Glm4MoeForCausalLM", 1, [model], [model], True),
+            ("Glm4MoeForCausalLM", 0, [model], [model, mtp], False),
+            ("LlamaForCausalLM", 1, [model], [model, mtp], False),
+        )
+        for arch, nextn, selected, available, may_list in cases:
+            with (
+                self.subTest(arch=arch, nextn=nextn, selected=selected),
+                patch.object(
+                    runai_utils, "list_safetensors", return_value=available
+                ) as listing,
+            ):
+                config = SimpleNamespace(
+                    architectures=[arch], num_nextn_predict_layers=nextn
+                )
+                self.assertEqual(
+                    maybe_add_mtp_safetensors(selected, folder, INDEX_NAME, config),
+                    selected,
+                )
+                if not may_list:
+                    listing.assert_not_called()
+
+    def test_local_unindexed_mtp(self):
+        with tempfile.TemporaryDirectory() as folder:
+            model = _touch(folder, "model.safetensors")
+            mtp = _touch(folder, "mtp.safetensors")
+            config = SimpleNamespace(
+                architectures=["Glm4MoeLiteForCausalLMNextN"],
+                num_nextn_predict_layers=1,
+            )
+            with patch.object(runai_utils, "list_safetensors") as listing:
+                self.assertEqual(
+                    maybe_add_mtp_safetensors([model], folder, INDEX_NAME, config),
+                    [model, mtp],
+                )
+                listing.assert_not_called()
+
+
 if __name__ == "__main__":
     unittest.main()