minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Aurick Qiao
2026-09-12 12:25:57 +08:00
committed by GitHub
parent a984c78330
commit c389dd0863
4 changed files with 191 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_loader/loader.py
+9 -1
View File
@@ -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
+14 -2
View File
@@ -830,7 +830,10 @@ def filter_duplicate_safetensors_files(
if any(fnmatch.fnmatch(rel_path, pattern) for pattern in allow_patterns):
files_to_validate.add(f)
missing_files = sorted(f for f in files_to_validate if not os.path.isfile(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(
f"{index_file} references {len(missing_files)} shard file(s) missing "
@@ -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
+78
View File
@@ -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
+90 -22
View File
@@ -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,17 +91,19 @@ class TestFilterDuplicateSafetensorsFiles(CustomTestCase):
"vit": "vision_encoder/model.safetensors",
},
)
transformer = _touch(
self.folder, "transformer/diffusion_pytorch_model.safetensors"
)
shard = "transformer/diffusion_pytorch_model.safetensors"
_touch(self.folder, shard)
result = filter_duplicate_safetensors_files(
hf_weights_files=[transformer],
hf_folder=self.folder,
index_file=INDEX_NAME,
allow_patterns=["transformer/*.safetensors"],
)
self.assertEqual(result, [transformer])
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=folder,
index_file=os.path.join(self.folder, INDEX_NAME),
allow_patterns=["transformer/*.safetensors"],
)
self.assertEqual(result, [transformer])
def test_missing_shard_inside_allow_patterns_raises(self):
_write_index(
@@ -106,18 +114,34 @@ class TestFilterDuplicateSafetensorsFiles(CustomTestCase):
"vit": "vision_encoder/model.safetensors",
},
)
transformer = _touch(
self.folder, "transformer/model-00001-of-00002.safetensors"
)
shard = "transformer/model-00001-of-00002.safetensors"
_touch(self.folder, shard)
with self.assertRaises(RuntimeError) as cm:
filter_duplicate_safetensors_files(
hf_weights_files=[transformer],
hf_folder=self.folder,
index_file=INDEX_NAME,
allow_patterns=["transformer/*.safetensors"],
)
self.assertIn("model-00002-of-00002.safetensors", str(cm.exception))
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=[os.path.join(folder, shard)],
hf_folder=folder,
index_file=os.path.join(self.folder, INDEX_NAME),
allow_patterns=["transformer/*.safetensors"],
)
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()