174 lines
6.8 KiB
Python
174 lines
6.8 KiB
Python
"""`model_config_of()` caches, and the key is the path the record carried.
|
|
|
|
Two movements of a `model_path` reach this cache, and only the first one means
|
|
the cached configuration describes the wrong checkpoint:
|
|
|
|
- the record's own path moves, because the GGUF and ModelScope handlers declare
|
|
a local path for a Hub reference;
|
|
- `ModelConfig` moves its own path to the local pull directory when the weights
|
|
sit behind an object-store URI, while still describing the checkpoint the
|
|
record asked for.
|
|
|
|
So the key is the path the record carried when the cache was filled, and these
|
|
cases pin both movements against it, on a raw record and on a resolved one.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import shutil
|
|
import tempfile
|
|
import unittest
|
|
from types import SimpleNamespace
|
|
|
|
from sglang.srt.arg_groups.overrides import declare_resolution, model_config_of
|
|
from sglang.srt.configs.model_config import ModelConfig
|
|
from sglang.srt.environ import EnvField, envs
|
|
from sglang.srt.server_args import ServerArgs
|
|
from sglang.srt.utils import runai_utils
|
|
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")
|
|
|
|
_MINI_CONFIG = {
|
|
"architectures": ["LlamaForCausalLM"],
|
|
"model_type": "llama",
|
|
"hidden_size": 16,
|
|
"intermediate_size": 32,
|
|
"num_attention_heads": 2,
|
|
"num_key_value_heads": 2,
|
|
"num_hidden_layers": 2,
|
|
"vocab_size": 128,
|
|
"max_position_embeddings": 2048,
|
|
}
|
|
|
|
_OBJECT_STORE_URI = "gs://sglang-test-bucket/mini-llama/"
|
|
|
|
|
|
class TestTheModelConfigCache(CustomTestCase):
|
|
def setUp(self):
|
|
# Resolution writes the environment on the way through
|
|
# (`SGLANG_USE_CUDA_IPC_TRANSPORT` among others) and flips the
|
|
# descriptor-level flag `EnvField.set()` keeps, which `os.environ` does
|
|
# not carry. The cases that resolve put both back.
|
|
fields = {}
|
|
for klass in reversed(type(envs).__mro__):
|
|
for name, field in vars(klass).items():
|
|
if isinstance(field, EnvField):
|
|
fields[name] = field
|
|
state = (
|
|
dict(os.environ),
|
|
{name: field._set_to_none for name, field in fields.items()},
|
|
)
|
|
self.addCleanup(self._restore, state)
|
|
|
|
@staticmethod
|
|
def _restore(state):
|
|
saved_environ, saved_none_flags = state
|
|
os.environ.clear()
|
|
os.environ.update(saved_environ)
|
|
for name, was_none in saved_none_flags.items():
|
|
getattr(type(envs), name)._set_to_none = was_none
|
|
|
|
def _checkpoint(self) -> str:
|
|
directory = tempfile.mkdtemp(prefix="model_config_cache_")
|
|
self.addCleanup(shutil.rmtree, directory, ignore_errors=True)
|
|
with open(os.path.join(directory, "config.json"), "w") as handle:
|
|
json.dump(_MINI_CONFIG, handle)
|
|
return directory
|
|
|
|
def _pulled_to(self, directory: str) -> None:
|
|
"""The launcher has already pulled the metadata for the URI.
|
|
|
|
Both entry points answer with the local directory, which is what a
|
|
second process on the same host sees: the download runs once, in the
|
|
launcher, and everything after it resolves the path.
|
|
"""
|
|
for name in ("download_and_get_path", "get_path"):
|
|
original = getattr(runai_utils.ObjectStorageModel, name)
|
|
self.addCleanup(setattr, runai_utils.ObjectStorageModel, name, original)
|
|
setattr(
|
|
runai_utils.ObjectStorageModel,
|
|
name,
|
|
classmethod(lambda cls, model_path, _dir=directory: _dir),
|
|
)
|
|
|
|
def _resolved(self, **kwargs) -> ServerArgs:
|
|
# device="cuda" keeps the golden path host-independent: an
|
|
# accelerator-less runner resolves only the base platform, where
|
|
# get_device() raises.
|
|
kwargs.setdefault("device", "cuda")
|
|
kwargs.setdefault("random_seed", 42)
|
|
server_args = ServerArgs(**kwargs)
|
|
server_args.resolve_once()
|
|
return server_args
|
|
|
|
def test_a_configuration_that_repoints_itself_stays_cached(self):
|
|
"""The object-store shape: the configuration's own path moved."""
|
|
pulled = self._checkpoint()
|
|
self._pulled_to(pulled)
|
|
|
|
server_args = self._resolved(
|
|
model_path=_OBJECT_STORE_URI, load_format="runai_streamer"
|
|
)
|
|
cached = server_args.__dict__["_model_config"]
|
|
self.assertIsInstance(cached, ModelConfig)
|
|
# The record still carries the URI the operator typed, and the
|
|
# configuration carries the directory it read the metadata from.
|
|
self.assertEqual(server_args.model_path, _OBJECT_STORE_URI)
|
|
self.assertEqual(cached.model_path, pulled)
|
|
|
|
self.assertIs(model_config_of(server_args), cached)
|
|
|
|
def test_a_declared_model_path_rebuilds_the_configuration(self):
|
|
"""The GGUF and ModelScope shape: the record's own path moved."""
|
|
first_checkpoint = self._checkpoint()
|
|
second_checkpoint = self._checkpoint()
|
|
|
|
server_args = ServerArgs(model_path=first_checkpoint, device="cuda")
|
|
first = model_config_of(server_args)
|
|
self.assertEqual(first.model_path, first_checkpoint)
|
|
|
|
declare_resolution(
|
|
server_args,
|
|
"test_a_declared_model_path_rebuilds_the_configuration",
|
|
model_path=second_checkpoint,
|
|
)
|
|
second = model_config_of(server_args)
|
|
self.assertIsNot(second, first)
|
|
self.assertEqual(second.model_path, second_checkpoint)
|
|
|
|
def test_the_cache_refills_on_a_resolved_record(self):
|
|
"""A rebuild has to be storable wherever the key can invalidate.
|
|
|
|
The record is read-only once resolution has finished, and the cache is
|
|
the record's own bookkeeping, so the guard lets the refill through.
|
|
"""
|
|
first_checkpoint = self._checkpoint()
|
|
second_checkpoint = self._checkpoint()
|
|
|
|
server_args = self._resolved(model_path=first_checkpoint)
|
|
copy_ = server_args.replace_resolved(
|
|
"test_the_cache_refills_on_a_resolved_record",
|
|
model_path=second_checkpoint,
|
|
)
|
|
|
|
rebuilt = model_config_of(copy_)
|
|
self.assertEqual(rebuilt.model_path, second_checkpoint)
|
|
self.assertIs(model_config_of(copy_), rebuilt)
|
|
# The parent keeps the configuration it resolved with.
|
|
self.assertEqual(model_config_of(server_args).model_path, first_checkpoint)
|
|
|
|
def test_a_supplied_configuration_is_handed_back(self):
|
|
"""A configuration nothing in here built carries no key, so nothing
|
|
invalidates it."""
|
|
server_args = ServerArgs(model_path=self._checkpoint(), device="cuda")
|
|
stand_in = SimpleNamespace(model_path="somewhere/else")
|
|
server_args._model_config = stand_in
|
|
|
|
self.assertIs(model_config_of(server_args), stand_in)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|