diff --git a/packages/sie_server/README.md b/packages/sie_server/README.md index f3d5c76ee..eb75f379c 100644 --- a/packages/sie_server/README.md +++ b/packages/sie_server/README.md @@ -1220,6 +1220,16 @@ attention. Calibrated profiles for these configurations retain their overrides. Automatic sizing is an estimate; unusually large runtime workspaces can still invoke the existing OOM recovery path. +A profile whose engine needs more memory outside its static pool than these +allowances can declare the measured excess as +`adapter_options.loadtime.runtime_headroom_gib`. An example is a vision encoder +processing one full-size page. SIE adds the value to the reservation in both +modes. Load staging then evicts idle eligible models until that much memory is +free, and an automatically sized engine does not launch without it. The +headroom stays outside the engine's static pool. A model loaded later is still +admitted on actual device usage, so it can take headroom that the engine has not +allocated yet. + | Env var | Default | Effect | |--|--|--| | `SIE_MEMORY_PRESSURE_THRESHOLD_PERCENT` | `95` | VRAM utilisation that triggers reactive LRU eviction by the pressure monitor. | diff --git a/packages/sie_server/src/sie_server/adapters/sglang/_memory.py b/packages/sie_server/src/sie_server/adapters/sglang/_memory.py index 08ad6d6ca..f670dffb4 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/_memory.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/_memory.py @@ -116,6 +116,7 @@ class SGLangMemoryBudget: mem_fraction_static: float | None = None tensor_parallel_size: int = 1 has_auxiliary_weights: bool = False + runtime_headroom_gib: float = 0 reserved_bytes: int = 0 def __post_init__(self) -> None: @@ -126,6 +127,9 @@ def __post_init__(self) -> None: isinstance(fraction, bool) or not isinstance(fraction, int | float) or not 0 < fraction <= 1 ): raise ValueError("mem_fraction_static must be a finite number in (0, 1]") + headroom = self.runtime_headroom_gib + if isinstance(headroom, bool) or not isinstance(headroom, int | float) or not 0 <= headroom < math.inf: + raise ValueError("runtime_headroom_gib must be a finite non-negative number") @property def token_capacity(self) -> int: @@ -150,15 +154,18 @@ def static_bytes(self) -> int: def required_bytes(self, *, device_type: str, device_total_bytes: int) -> int | None: if device_type != "cuda" or device_total_bytes <= 0: return None + headroom = math.ceil(self.runtime_headroom_gib * GIB) if self.mem_fraction_static is not None: required = _server.estimate_load_required_memory_bytes( device_type=device_type, device_total_bytes=device_total_bytes, mem_fraction_static=self.mem_fraction_static, ) + if required is not None: + required = min(required + headroom, device_total_bytes) self.reserved_bytes = (required or 0) * self.tensor_parallel_size return required - required = self.static_bytes + RUNTIME_HEADROOM_BYTES + required = self.static_bytes + RUNTIME_HEADROOM_BYTES + headroom self.reserved_bytes = required return required diff --git a/packages/sie_server/src/sie_server/adapters/sglang/generation.py b/packages/sie_server/src/sie_server/adapters/sglang/generation.py index 2cf2f9d34..a97721a1e 100644 --- a/packages/sie_server/src/sie_server/adapters/sglang/generation.py +++ b/packages/sie_server/src/sie_server/adapters/sglang/generation.py @@ -552,6 +552,10 @@ def __init__( # it, then follows as usual. None leaves reasoning unbounded and the # launch unchanged. Needs a reasoning parser with known control tokens. thinking_budget: int | None = None, + # GiB the engine allocates outside its static pool beyond the default + # allowance at peak, such as a vision encoder's activations for one + # full-size image. Load staging requires it free before launch. + runtime_headroom_gib: float = 0, **kwargs: Any, # accept extra args from loader for compatibility ) -> None: _ = kwargs @@ -571,6 +575,7 @@ def __init__( mem_fraction_static=mem_fraction_static, tensor_parallel_size=self._tensor_parallel_size, has_auxiliary_weights=bool(lora_paths or (speculative and speculative.get("enabled"))), + runtime_headroom_gib=runtime_headroom_gib, ) self._watchdog_timeout_s = _server.validate_optional_positive(watchdog_timeout_s, field="watchdog_timeout_s") self._declared_nccl_port = _server.validate_optional_port(nccl_port) diff --git a/packages/sie_server/tests/adapters/test_sglang_memory.py b/packages/sie_server/tests/adapters/test_sglang_memory.py index b0f6022ad..07b7747a0 100644 --- a/packages/sie_server/tests/adapters/test_sglang_memory.py +++ b/packages/sie_server/tests/adapters/test_sglang_memory.py @@ -93,6 +93,50 @@ def test_explicit_fraction_keeps_admission_and_unbounded_kv_behavior(monkeypatch assert budget.launch_args(0) == ["--mem-fraction-static", "0.85"] +def test_runtime_headroom_is_reserved_in_both_admission_paths(model_dir): + automatic = _budget(model_dir).required_bytes(device_type="cuda", device_total_bytes=80 * GIB) + budget = _budget(model_dir, runtime_headroom_gib=4) + assert budget.required_bytes(device_type="cuda", device_total_bytes=80 * GIB) == automatic + 4 * GIB + assert budget.reserved_bytes == automatic + 4 * GIB + + explicit = _budget(Path("missing-model"), mem_fraction_static=0.85, runtime_headroom_gib=4) + assert explicit.required_bytes(device_type="cuda", device_total_bytes=80 * GIB) == 73 * GIB + assert explicit.reserved_bytes == 73 * GIB + oversized = _budget(Path("missing-model"), mem_fraction_static=0.85, runtime_headroom_gib=20) + assert oversized.required_bytes(device_type="cuda", device_total_bytes=80 * GIB) == 80 * GIB + + +def test_runtime_headroom_stays_outside_the_static_pool(model_dir, monkeypatch): + budget = _budget(model_dir, runtime_headroom_gib=4) + without_headroom = _budget(model_dir).required_bytes(device_type="cuda", device_total_bytes=80 * GIB) + monkeypatch.setattr(_memory, "device_memory", lambda _: (without_headroom, 80 * GIB)) + with pytest.raises(RuntimeError, match="CUDA out of memory"): + budget.launch_args(0) + monkeypatch.setattr(_memory, "device_memory", lambda _: (60 * GIB, 80 * GIB)) + assert float(budget.launch_args(0)[1]) * 60 * GIB == pytest.approx(budget.static_bytes) + + +@pytest.mark.parametrize("headroom", [-1, True, float("nan"), float("inf"), "4"]) +def test_invalid_runtime_headroom_rejected(headroom): + with pytest.raises(ValueError, match="runtime_headroom_gib"): + _budget(Path("missing"), runtime_headroom_gib=headroom) + + +def test_load_staging_evicts_until_the_runtime_headroom_is_free(model_dir): + plain = SGLangVisionExtractAdapter(str(model_dir), max_seq_length=32768, kv_budget_tokens=32768) + vision = SGLangVisionExtractAdapter( + str(model_dir), max_seq_length=32768, kv_budget_tokens=32768, runtime_headroom_gib=4 + ) + plain_required = plain.load_required_memory_bytes(device_type="cuda", device_total_bytes=80 * GIB) + assert plain_required is not None + manager = MemoryManager(device="cuda:0") + free = plain_required + GIB + with patch.object(manager, "get_stats", return_value=MemoryStats(80 * GIB - free, 80 * GIB, "cuda")): + assert not manager.should_evict_for_load(_adapter_load_required_bytes(plain, manager)) + assert _adapter_load_required_bytes(vision, manager) == plain_required + 4 * GIB + assert manager.should_evict_for_load(_adapter_load_required_bytes(vision, manager)) + + def test_sharded_weights_use_index_and_serving_precision(model_dir): (model_dir / "model.safetensors").unlink() _write_weights(model_dir / "model-1.safetensors", 1000, "F32", 4) diff --git a/packages/sie_server/tests/config/test_sglang_memory_profiles.py b/packages/sie_server/tests/config/test_sglang_memory_profiles.py index c3175e34a..c5a445b5d 100644 --- a/packages/sie_server/tests/config/test_sglang_memory_profiles.py +++ b/packages/sie_server/tests/config/test_sglang_memory_profiles.py @@ -107,6 +107,19 @@ def test_glm_ocr_launches_with_its_explicit_fraction(): assert budget.launch_args(0) == ["--mem-fraction-static", "0.85"] +def test_profile_runtime_headroom_reaches_the_load_reservation(): + config = load_model_config(MODELS / "zai-org__GLM-OCR.yaml") + baseline = load_adapter(config, MODELS, device="cuda:0") + assert baseline._memory_budget.runtime_headroom_gib == 0 + values = config.model_dump() + values["profiles"]["default"]["adapter_options"]["loadtime"]["runtime_headroom_gib"] = 4 + adapter = load_adapter(type(config).model_validate(values), MODELS, device="cuda:0") + assert adapter._memory_budget.runtime_headroom_gib == 4 + required = adapter.load_required_memory_bytes(device_type="cuda", device_total_bytes=80 * GIB) + assert required == baseline.load_required_memory_bytes(device_type="cuda", device_total_bytes=80 * GIB) + 4 * GIB + assert required == 73 * GIB + + def test_guard_calibrated_profile_retains_explicit_override(): config = load_model_config(MODELS / "Qwen__Qwen3Guard-Gen-4B.yaml") assert config.resolve_profile("default").kv_budget_tokens == 32768