config: keep runtime hicache and weight-version updates off ServerArgs (#33336)

The scheduler's runtime HiCache attach/detach wrote its own ServerArgs so the
internal-state readback would show the change; that readback already reports
the resolved config, so the writes become get_context().override(...) and the
namespace readers see them too.

The tokenizer side is per-engine — several Engines can share one process — so
its control-plane updates (weight version, model path + load format, HiCache
attach/detach) stay with the manager instead of moving to the process-global
bags. TokenizerManager gains record_config_updates / config_value /
resolved_config_dict, and the readbacks that used to observe the instance write
(/server_info, /model_info, the HiCache status endpoint, the gRPC bridge) now
overlay those updates onto the startup config.

test_server_info's stub grew the real manager instead of a SimpleNamespace, so
the overlay it now exercises cannot drift from production.

Writer ratchet 26 -> 19.
This commit is contained in:
Cheng Wan
2026-08-02 21:23:38 -07:00
committed by GitHub
parent 9bc8848fcf
commit 0b3e8bedd1
16 changed files with 432 additions and 42 deletions
+3 -1
View File
@@ -1271,7 +1271,9 @@ class Engine(EngineScoreMixin, EngineBase):
)
return msgspec_to_builtins(
{
**dataclasses.asdict(self.tokenizer_manager.server_args),
**self.tokenizer_manager.resolved_config_dict(
dataclasses.asdict(self.tokenizer_manager.server_args)
),
**self._scheduler_init_result.scheduler_infos[0],
"internal_states": internal_states,
"version": __version__,
+4 -2
View File
@@ -379,7 +379,7 @@ class RuntimeHandle:
"model_path": self.tokenizer_manager.model_path,
"tokenizer_path": self.tokenizer_manager.server_args.tokenizer_path,
"is_generation": self.tokenizer_manager.is_generation,
"weight_version": self.tokenizer_manager.server_args.weight_version,
"weight_version": self.tokenizer_manager.config_value("weight_version"),
"model_type": getattr(model_config.hf_config, "model_type", None),
"architectures": getattr(model_config.hf_config, "architectures", None),
}
@@ -393,7 +393,9 @@ class RuntimeHandle:
return json.dumps(result, default=str)
def get_server_info(self) -> str:
result: Dict[str, Any] = dataclasses.asdict(self.server_args)
result: Dict[str, Any] = self.tokenizer_manager.resolved_config_dict(
dataclasses.asdict(self.tokenizer_manager.server_args)
)
result.update(self.scheduler_info)
return json.dumps(msgspec_to_builtins(result), default=str)
+14 -9
View File
@@ -710,12 +710,13 @@ async def model_info():
"tokenizer_path": _global_state.tokenizer_manager.server_args.tokenizer_path,
"is_generation": _global_state.tokenizer_manager.is_generation,
"preferred_sampling_params": _global_state.tokenizer_manager.server_args.preferred_sampling_params,
"weight_version": _global_state.tokenizer_manager.server_args.weight_version,
"weight_version": _global_state.tokenizer_manager.config_value(
"weight_version"
),
"has_image_understanding": model_config.is_image_understandable_model,
"has_audio_understanding": model_config.is_audio_understandable_model,
"model_type": getattr(model_config.hf_config, "model_type", None),
"architectures": getattr(model_config.hf_config, "architectures", None),
"weight_version": _global_state.tokenizer_manager.server_args.weight_version,
# "hf_config": model_config.hf_config.to_dict(),
}
embedding_model_spec = getattr(model_config, "embedding_model_spec", None)
@@ -761,7 +762,9 @@ async def server_info():
# server_args.model_config is not serializable but should be excluded by asdict.
return msgspec_to_builtins(
{
**dataclasses.asdict(server_args),
**_global_state.tokenizer_manager.resolved_config_dict(
dataclasses.asdict(server_args)
),
**_global_state.scheduler_info,
"internal_states": internal_states,
"version": __version__,
@@ -1091,10 +1094,13 @@ async def hicache_storage_backend_status():
return _admin_api_key_missing_response()
return {
"hicache_storage_backend": _global_state.tokenizer_manager.server_args.hicache_storage_backend,
"hicache_storage_backend_extra_config": _global_state.tokenizer_manager.server_args.hicache_storage_backend_extra_config,
"hicache_storage_prefetch_policy": _global_state.tokenizer_manager.server_args.hicache_storage_prefetch_policy,
"hicache_write_policy": _global_state.tokenizer_manager.server_args.hicache_write_policy,
name: _global_state.tokenizer_manager.config_value(name)
for name in (
"hicache_storage_backend",
"hicache_storage_backend_extra_config",
"hicache_storage_prefetch_policy",
"hicache_write_policy",
)
}
@@ -1385,8 +1391,7 @@ async def update_weight_version(
# Use a simple approach without the complex lock mechanism for now
# since weight_version update is a simple operation that doesn't affect model weights
try:
# Update the weight version in server args (the single source of truth)
_global_state.tokenizer_manager.server_args.override(
_global_state.tokenizer_manager.record_config_updates(
"http.update_weight_version", weight_version=obj.new_version
)
@@ -338,12 +338,12 @@ class RealtimeConnection:
if (
transcription is not None
and transcription.model
and transcription.model != self.server_args.served_model_name
and transcription.model != self.tokenizer_manager.served_model_name
):
await self._send_error(
"not_supported",
f"Model {transcription.model!r} is not served by this endpoint "
f"(serving {self.server_args.served_model_name!r}); set "
f"(serving {self.tokenizer_manager.served_model_name!r}); set "
f"transcription.model to null or to the server's model name.",
param="session.audio.input.transcription.model",
)
@@ -1279,7 +1279,7 @@ class OpenAIServingChat(OpenAIServingBase):
logger.warning(
"Model '%s' supports only 'low' reasoning effort; "
"requested '%s' treated as default thinking",
self.tokenizer_manager.server_args.served_model_name,
self.tokenizer_manager.served_model_name,
request.reasoning_effort,
)
@@ -39,7 +39,7 @@ class OpenAIServingClassify(OpenAIServingBase):
self.model_name = (
self.tokenizer_manager.served_model_name
if self.tokenizer_manager.served_model_name
else self.tokenizer_manager.server_args.model_path
else self.tokenizer_manager.model_path
)
if not self.id2label:
raise ValueError("id2label mapping is missing")
+2 -2
View File
@@ -4009,7 +4009,7 @@ class Scheduler(
)
if recv_req.hicache_write_policy is not None:
hicache_fields["hicache_write_policy"] = recv_req.hicache_write_policy
self.server_args.override("scheduler.attach_hicache", **hicache_fields)
get_context().override("scheduler.attach_hicache", **hicache_fields)
logger.info(
f"Attached HiCache storage backend: {recv_req.hicache_storage_backend}"
)
@@ -4050,7 +4050,7 @@ class Scheduler(
if ok or (not self.enable_hicache_storage):
# Treat "already disabled / nothing to do" as success for idempotence.
self.enable_hicache_storage = False
self.server_args.override(
get_context().override(
"scheduler.detach_hicache",
hicache_storage_backend=None,
hicache_storage_backend_extra_config=None,
@@ -343,7 +343,7 @@ class TokenizerControlMixin:
)
if hicache_write_policy is not None:
hicache_fields["hicache_write_policy"] = hicache_write_policy
self.server_args.override("tokenizer.attach_hicache", **hicache_fields)
self.record_config_updates("tokenizer.attach_hicache", **hicache_fields)
return out
async def detach_hicache_storage(
@@ -359,7 +359,7 @@ class TokenizerControlMixin:
out = DetachHiCacheStorageReqOutput(success=all_success, message=all_message)
# TODO: partial rollback if failed
if all_success:
self.server_args.override(
self.record_config_updates(
"tokenizer.detach_hicache",
hicache_storage_backend=None,
hicache_storage_backend_extra_config=None,
@@ -920,6 +920,6 @@ class TokenizerControlMixin:
) -> None:
"""Update weight version if provided."""
if weight_version is not None:
self.server_args.override(
self.record_config_updates(
"tokenizer.weight_version", weight_version=weight_version
)
@@ -364,6 +364,11 @@ class InputFormat(Enum):
CROSS_ENCODER_PAIRS = 3 # Cross-encoder pairs like [["query", "document"]]
_SERVER_ARGS_FIELDS = frozenset(f.name for f in dataclasses.fields(ServerArgs))
_MANAGER_OWNED_FIELDS = ("model_path", "served_model_name")
class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
"""TokenizerManager is a process that tokenizes the text."""
@@ -386,6 +391,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
):
# Parse args
self.server_args = server_args
self._config_updates: List[Tuple[str, Dict[str, Any]]] = []
self.elastic_worker_count = server_args.dp_size
self.elastic_pending_ep_size = None
self.elastic_scale_phase = "idle"
@@ -1888,9 +1894,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
) -> Tuple[bool, str]:
self.auto_create_handle_loop()
# default the load format to the server_args
if obj.load_format is None:
obj.load_format = self.server_args.load_format
obj.load_format = self.config_value("load_format")
logger.info("Start update_weights. Load format=%s", obj.load_format)
if obj.abort_all_requests:
@@ -1914,11 +1919,57 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
return success, message, num_paused_requests
def record_config_updates(self, source: str, **fields) -> None:
"""Record a control-plane config change for this engine.
Per-engine state: several ``Engine``s can share one tokenizer process.
The readback endpoints overlay these onto the startup config. The
process-global sibling is ``RuntimeContext.override`` /
``resolved_server_args_dict``, which writes the config bags every
process shares.
"""
unknown = sorted(f for f in fields if f not in _SERVER_ARGS_FIELDS)
if unknown:
raise ValueError(
f"{unknown} are not ServerArgs fields; the readback endpoints "
"overlay these onto a serialized ServerArgs, so an unknown key "
"would surface as a phantom config entry."
)
self._config_updates.append((source, dict(fields)))
def config_value(self, name: str):
"""The value in effect for one config field, control-plane updates first."""
if name in _MANAGER_OWNED_FIELDS:
return getattr(self, name)
for _source, fields in reversed(self._config_updates):
if name in fields:
return fields[name]
return getattr(self.server_args, name)
def _dump_config_snapshot(self) -> Optional[Dict[str, Any]]:
"""The config in effect, or None when it cannot be serialized.
A dump is worth having even when the config is not: request data is the
part that cannot be reconstructed afterwards.
"""
try:
return self.resolved_config_dict(dataclasses.asdict(self.server_args))
except Exception as e:
logger.error(f"Failed to snapshot the resolved config for the dump: {e!r}")
return None
def resolved_config_dict(self, base: Dict[str, Any]) -> Dict[str, Any]:
"""``base`` (a serialized ``ServerArgs``) with the control-plane updates on top."""
resolved = dict(base)
for _source, fields in self._config_updates:
resolved.update(fields)
for name in _MANAGER_OWNED_FIELDS:
resolved[name] = getattr(self, name)
return resolved
def _update_model_path_info(self, model_path: str, load_format: str):
self.served_model_name = model_path
self.server_args.override(
"tokenizer.update_weights", model_path=model_path, load_format=load_format
)
self.record_config_updates("tokenizer.update_weights", load_format=load_format)
self.model_path = model_path
async def _wait_for_model_update_from_disk(
@@ -2060,7 +2111,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
"id": rid,
"finish_reason": recv_obj.finished_reasons[i],
"prompt_tokens": recv_obj.prompt_tokens[i],
"weight_version": self.server_args.weight_version,
"weight_version": self.config_value("weight_version"),
"num_retractions": recv_obj.retraction_counts[i],
}
@@ -2789,6 +2840,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
logger.info(log_message)
to_dump_with_server_args = {
"server_args": self.server_args,
"config_updates": list(self._config_updates),
"resolved_config": self._dump_config_snapshot(),
"requests": data_list.copy(),
}
@@ -2808,6 +2861,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
f.seek(0)
f.truncate()
to_dump_with_server_args["server_args"] = None
# The snapshot copies the same object field by field.
to_dump_with_server_args["resolved_config"] = None
pickle.dump(to_dump_with_server_args, f)
asyncio.create_task(asyncio.to_thread(background_task))
@@ -2870,6 +2925,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
# Write the data to the file
data_to_dump_with_server_args = {
"server_args": self.server_args,
"config_updates": list(self._config_updates),
"resolved_config": self._dump_config_snapshot(),
"requests": data_to_dump,
"launch_command": " ".join(sys.argv),
}
@@ -2887,6 +2944,8 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
f.seek(0)
f.truncate()
data_to_dump_with_server_args["server_args"] = None
# The snapshot copies the same object field by field.
data_to_dump_with_server_args["resolved_config"] = None
pickle.dump(data_to_dump_with_server_args, f)
logger.error(
f"Dumped {len(self.crash_dump_request_list)} finished and {len(unfinished_requests)} unfinished requests before crash to {filename}"
@@ -2999,7 +3058,7 @@ class TokenizerManager(TokenizerControlMixin, TokenizerManagerScoreMixin):
meta_info = {
"id": recv_obj.rid,
"finish_reason": finish_reason,
"weight_version": self.server_args.weight_version,
"weight_version": self.config_value("weight_version"),
"e2e_latency": state.time_stats.get_e2e_latency(),
}
is_stream = getattr(state.obj, "stream", False)
+14 -8
View File
@@ -912,17 +912,23 @@ class RuntimeContext:
"""Serialize the *resolved* config: the pristine ``server_args`` fields
with every post-publish ``override`` overlaid.
Reporting endpoints (``/server_info``, ``get_internal_state``) surface
the config the process is *currently* running, not the startup record,
so they read this rather than serializing ``server_args`` directly
otherwise runtime updates (weight version, model path, tunables set via
``/set_internal_state``) never show up in the readback.
``get_internal_state`` reports this, and ``/server_info`` carries it in
the ``internal_states`` block, so scheduler-side runtime changes show up
in a readback: HiCache attach/detach, the generated forward-pass-metrics
endpoint, tunables set via ``/set_internal_state``.
``base`` defaults to ``dict(vars(server_args))`` (matching the legacy
``vars`` dump); pass ``dataclasses.asdict(server_args)`` when nested
dataclass fields must be expanded first (``/server_info``). Override
leaves are flat ``ServerArgs`` field names, so overlaying them onto the
top level of either base is exact.
dataclass fields must be expanded first. Override leaves are flat
``ServerArgs`` field names, so overlaying them onto the top level of
either base is exact.
This covers the process-global bags only. Per-engine control-plane
changes (weight version, model path, the tokenizer's HiCache mirror)
live on the tokenizer manager several ``Engine``s can share one
process and ``TokenizerManager.resolved_config_dict`` overlays those
for the top-level ``/server_info`` body. The two are separate logs, not
one merged dict.
"""
d = dict(vars(self.server_args)) if base is None else dict(base)
for _source, fields in self._overrides_log: