Move invariant checks to SchedulerInvariantChecker and retire runtime_checker mixin (#25624)
This commit is contained in:
fzyzcjy
@@ -211,7 +211,7 @@ class SchedulerMlxOverlapMixin:
|
||||
self.cur_batch = pending_curr.batch_copy
|
||||
self.last_batch = pending_curr.batch_copy
|
||||
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
|
||||
self.self_check_during_busy()
|
||||
self.invariant_checker.self_check_during_busy()
|
||||
continue
|
||||
|
||||
# 4. Chain is broken. Finalise pending_next (if any), then
|
||||
@@ -230,4 +230,4 @@ class SchedulerMlxOverlapMixin:
|
||||
|
||||
self.last_batch = next_batch
|
||||
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
|
||||
self.self_check_during_busy()
|
||||
self.invariant_checker.self_check_during_busy()
|
||||
|
||||
@@ -331,8 +331,7 @@ def create_scheduler_watchdog(
|
||||
def dump_info() -> str:
|
||||
if scheduler.is_initializing:
|
||||
return ""
|
||||
_, messages = scheduler._check_all_pools(
|
||||
scheduler.invariant_checker,
|
||||
_, messages = scheduler.invariant_checker._check_all_pools(
|
||||
scheduler.pool_stats_observer.get_pool_stats(),
|
||||
)
|
||||
return (
|
||||
@@ -1523,7 +1522,7 @@ class Scheduler(
|
||||
# Update last_batch
|
||||
self.last_batch = batch
|
||||
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
|
||||
self.self_check_during_busy(self.invariant_checker)
|
||||
self.invariant_checker.self_check_during_busy()
|
||||
|
||||
@DynamicGradMode()
|
||||
def event_loop_overlap(self):
|
||||
@@ -1578,7 +1577,7 @@ class Scheduler(
|
||||
self.last_batch = batch
|
||||
|
||||
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get():
|
||||
self.self_check_during_busy(self.invariant_checker)
|
||||
self.invariant_checker.self_check_during_busy()
|
||||
|
||||
def is_disable_overlap_for_batch(self, batch: ScheduleBatch) -> bool:
|
||||
# For two consecutive prefill batches, we disable overlap to improve the TTFT of the first batch.
|
||||
@@ -3105,16 +3104,15 @@ class Scheduler(
|
||||
# memory leak check (skipped for hisparse — pool counters intentionally
|
||||
# diverge during host-backup, see _get_swa_token_info clamp).
|
||||
if not self.enable_hisparse:
|
||||
has_leak, messages = self._check_all_pools(
|
||||
self.invariant_checker,
|
||||
has_leak, messages = self.invariant_checker._check_all_pools(
|
||||
self.pool_stats_observer.get_pool_stats(),
|
||||
)
|
||||
if has_leak:
|
||||
self._report_leak(self.invariant_checker, "pool", "\n".join(messages))
|
||||
self._check_req_pool(self.invariant_checker)
|
||||
self.invariant_checker._report_leak("pool", "\n".join(messages))
|
||||
self.invariant_checker._check_req_pool()
|
||||
|
||||
# tree cache sanity check
|
||||
self._check_tree_cache(self.invariant_checker)
|
||||
self.invariant_checker._check_tree_cache()
|
||||
|
||||
# metrics every 30s
|
||||
self._maybe_log_idle_metrics()
|
||||
|
||||
@@ -1,17 +1,29 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import warnings
|
||||
from dataclasses import dataclass
|
||||
from typing import Callable, Optional
|
||||
from typing import (
|
||||
Callable,
|
||||
List,
|
||||
Optional,
|
||||
Tuple,
|
||||
)
|
||||
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.managers.scheduler_components.pool_stats_observer import (
|
||||
PoolStats,
|
||||
SchedulerPoolStatsObserver,
|
||||
)
|
||||
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
|
||||
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
|
||||
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
from sglang.srt.utils.common import (
|
||||
ceil_align,
|
||||
raise_error_or_warn,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@@ -34,3 +46,224 @@ class SchedulerInvariantChecker:
|
||||
get_running_batch: Callable
|
||||
count_req_pool_leak_warnings: int = 0
|
||||
count_memory_leak_warnings: int = 0
|
||||
|
||||
@staticmethod
|
||||
def _check_pool_invariant(
|
||||
pool_name: str,
|
||||
available: int,
|
||||
evictable: int,
|
||||
protected: int,
|
||||
session_held: int,
|
||||
total: int,
|
||||
uncached: int = 0,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Check: available + evictable + protected + session_held + uncached == total."""
|
||||
total_accounted = available + evictable + protected + session_held + uncached
|
||||
leak = total_accounted != total
|
||||
msg = (
|
||||
f"[{pool_name}] {total=}, {available=}, {evictable=}, "
|
||||
f"{protected=}, {session_held=}, {uncached=}"
|
||||
)
|
||||
return leak, msg
|
||||
|
||||
def _check_full_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]:
|
||||
if self.is_hybrid_swa:
|
||||
protected = self.tree_cache.full_protected_size()
|
||||
session_held = self.pool_stats_observer.session_held_full_tokens()
|
||||
total = self.full_tokens_per_layer
|
||||
elif self.is_hybrid_ssm and self.tree_cache.supports_mamba():
|
||||
protected = self.tree_cache.full_protected_size()
|
||||
session_held = self.pool_stats_observer.session_held_tokens()
|
||||
total = self.token_to_kv_pool_allocator.size
|
||||
else:
|
||||
protected = self.tree_cache.protected_size()
|
||||
session_held = self.pool_stats_observer.session_held_tokens()
|
||||
total = self.max_total_num_tokens
|
||||
return self._check_pool_invariant(
|
||||
"full",
|
||||
ps.full_available_size,
|
||||
ps.full_evictable_size,
|
||||
protected,
|
||||
session_held,
|
||||
total,
|
||||
uncached,
|
||||
)
|
||||
|
||||
def _check_swa_pool(self, ps: PoolStats, uncached: int = 0) -> Tuple[bool, str]:
|
||||
return self._check_pool_invariant(
|
||||
"swa",
|
||||
ps.swa_available_size,
|
||||
ps.swa_evictable_size,
|
||||
self.tree_cache.swa_protected_size(),
|
||||
self.pool_stats_observer.session_held_swa_tokens(),
|
||||
self.swa_tokens_per_layer,
|
||||
uncached,
|
||||
)
|
||||
|
||||
def _check_mamba_pool(self, ps: PoolStats) -> Tuple[bool, str]:
|
||||
leak, msg = self._check_pool_invariant(
|
||||
"mamba",
|
||||
ps.mamba_available_size,
|
||||
ps.mamba_evictable_size,
|
||||
self.tree_cache.mamba_protected_size(),
|
||||
self.pool_stats_observer.session_held_mamba_slots(),
|
||||
self.req_to_token_pool.mamba_pool.size,
|
||||
)
|
||||
if leak:
|
||||
# Page-level leak diagnosis for mamba
|
||||
free_full_pages = set(
|
||||
self.token_to_kv_pool_allocator.free_pages.tolist()
|
||||
+ self.token_to_kv_pool_allocator.release_pages.tolist()
|
||||
)
|
||||
cached_full_pages = set(self.tree_cache.all_values_flatten().tolist())
|
||||
expected_full_pages = set(
|
||||
range(1, self.token_to_kv_pool_allocator.size + 1)
|
||||
)
|
||||
leaked_full_pages = (
|
||||
expected_full_pages - free_full_pages - cached_full_pages
|
||||
)
|
||||
free_mamba_pages = set(
|
||||
self.req_to_token_pool.mamba_pool.free_slots.tolist()
|
||||
)
|
||||
cached_mamba_pages = set(
|
||||
self.tree_cache.all_mamba_values_flatten().tolist()
|
||||
)
|
||||
expected_mamba_pages = set(range(self.req_to_token_pool.mamba_pool.size))
|
||||
leaked_mamba_pages = (
|
||||
expected_mamba_pages - free_mamba_pages - cached_mamba_pages
|
||||
)
|
||||
msg += (
|
||||
f", leaked_full_pages={leaked_full_pages or None}"
|
||||
f", leaked_mamba_pages={leaked_mamba_pages or None}"
|
||||
)
|
||||
return leak, msg
|
||||
|
||||
def _get_total_uncached_sizes(
|
||||
self,
|
||||
) -> Tuple[int, int]:
|
||||
"""Sum uncached tokens for full and SWA pools across all active batches.
|
||||
|
||||
Returns (full_uncached, swa_uncached). For non-SWA models, swa_uncached is 0.
|
||||
|
||||
For full pool: uncached = allocated - cache_protected_len
|
||||
For SWA pool: uncached = allocated - max(cache_protected_len, swa_evicted_seqlen)
|
||||
"""
|
||||
# After decode: running_batch IS last_batch (same object), count once.
|
||||
# After prefill: they differ, both hold uncached tokens.
|
||||
batches = [self.get_last_batch()]
|
||||
if (
|
||||
self.get_running_batch() not in (None, self.get_last_batch())
|
||||
and not self.get_running_batch().is_empty()
|
||||
):
|
||||
batches.append(self.get_running_batch())
|
||||
|
||||
full_uncached = 0
|
||||
swa_uncached = 0
|
||||
for batch in batches:
|
||||
for req in batch.reqs:
|
||||
assert req.kv_committed_freed == req.kv_overallocated_freed
|
||||
if req.kv_committed_freed or req.req_pool_idx is None:
|
||||
continue
|
||||
|
||||
allocated_len = req.kv_allocated_len
|
||||
if self.page_size > 1:
|
||||
allocated_len = ceil_align(allocated_len, self.page_size)
|
||||
assert req.cache_protected_len % self.page_size == 0
|
||||
|
||||
full_uncached += allocated_len - req.cache_protected_len
|
||||
if self.is_hybrid_swa:
|
||||
swa_uncached += allocated_len - max(
|
||||
req.cache_protected_len, req.swa_evicted_seqlen
|
||||
)
|
||||
|
||||
return full_uncached, swa_uncached
|
||||
|
||||
def self_check_during_busy(self):
|
||||
if self.get_last_batch() is None:
|
||||
return
|
||||
|
||||
spec_topk = self.server_args.speculative_eagle_topk or 1
|
||||
if spec_topk > 1:
|
||||
warnings.warn(
|
||||
"Runtime memory check (busy) is not supported when speculation topk > 1."
|
||||
)
|
||||
return
|
||||
|
||||
ps = self.pool_stats_observer.get_pool_stats()
|
||||
full_uncached, swa_uncached = self._get_total_uncached_sizes()
|
||||
|
||||
full_leak, full_msg = self._check_full_pool(ps, uncached=full_uncached)
|
||||
|
||||
swa_leak, swa_msg = False, ""
|
||||
if self.is_hybrid_swa:
|
||||
swa_leak, swa_msg = self._check_swa_pool(ps, uncached=swa_uncached)
|
||||
|
||||
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get() > 1:
|
||||
logger.info(f"[Mem Check (BUSY)] {full_msg}")
|
||||
if swa_msg:
|
||||
logger.info(f"[Mem Check (BUSY)] {swa_msg}")
|
||||
assert not full_leak, f"Full Pool Mem Leak Detected! {full_msg}"
|
||||
assert not swa_leak, f"SWA Pool Mem Leak Detected! {swa_msg}"
|
||||
|
||||
def _check_req_pool(self):
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE:
|
||||
req_total_size = (
|
||||
self.req_to_token_pool.size + self.req_to_token_pool.pre_alloc_size
|
||||
)
|
||||
else:
|
||||
req_total_size = self.req_to_token_pool.size
|
||||
|
||||
session_req_count = self.pool_stats_observer.session_held_req_count()
|
||||
if len(self.req_to_token_pool.free_slots) + session_req_count != req_total_size:
|
||||
msg = (
|
||||
"req_to_token_pool memory leak detected!"
|
||||
f"available_size={len(self.req_to_token_pool.free_slots)}, "
|
||||
f"session_held={session_req_count}, "
|
||||
f"total_size={self.req_to_token_pool.size}\n"
|
||||
)
|
||||
raise_error_or_warn(
|
||||
self,
|
||||
envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
|
||||
"count_req_pool_leak_warnings",
|
||||
msg,
|
||||
)
|
||||
|
||||
def _report_leak(self, pool_name: str, token_msg: str):
|
||||
msg = f"{pool_name} memory leak detected! {token_msg}"
|
||||
raise_error_or_warn(
|
||||
self,
|
||||
envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
|
||||
"count_memory_leak_warnings",
|
||||
msg,
|
||||
)
|
||||
|
||||
def _check_all_pools(
|
||||
self, ps: PoolStats, uncached: int = 0
|
||||
) -> Tuple[bool, List[str]]:
|
||||
"""Check memory invariant across all pools. Returns (has_leak, messages)."""
|
||||
has_leak = False
|
||||
messages = []
|
||||
|
||||
full_leak, full_msg = self._check_full_pool(ps, uncached=uncached)
|
||||
has_leak |= full_leak
|
||||
messages.append(full_msg)
|
||||
|
||||
if self.is_hybrid_swa:
|
||||
swa_leak, swa_msg = self._check_swa_pool(ps)
|
||||
has_leak |= swa_leak
|
||||
messages.append(swa_msg)
|
||||
|
||||
if self.is_hybrid_ssm and self.tree_cache.supports_mamba():
|
||||
mamba_leak, mamba_msg = self._check_mamba_pool(ps)
|
||||
has_leak |= mamba_leak
|
||||
messages.append(mamba_msg)
|
||||
|
||||
return has_leak, messages
|
||||
|
||||
def _check_tree_cache(self):
|
||||
if (
|
||||
self.tree_cache.is_tree_cache()
|
||||
and (self.is_hybrid_swa and self.tree_cache.supports_swa())
|
||||
or (self.is_hybrid_ssm and self.tree_cache.supports_mamba())
|
||||
):
|
||||
self.tree_cache.sanity_check()
|
||||
|
||||
@@ -2,261 +2,18 @@ from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import time
|
||||
import warnings
|
||||
from typing import TYPE_CHECKING, List, Tuple
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sglang.srt.disaggregation.utils import DisaggregationMode
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.observability.metrics_collector import QueueCount
|
||||
from sglang.srt.utils.common import ceil_align, raise_error_or_warn
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
from sglang.srt.managers.scheduler_components.invariant_checker import (
|
||||
SchedulerInvariantChecker,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SchedulerRuntimeCheckerMixin:
|
||||
@staticmethod
|
||||
def _check_pool_invariant(
|
||||
pool_name: str,
|
||||
available: int,
|
||||
evictable: int,
|
||||
protected: int,
|
||||
session_held: int,
|
||||
total: int,
|
||||
uncached: int = 0,
|
||||
) -> Tuple[bool, str]:
|
||||
"""Check: available + evictable + protected + session_held + uncached == total."""
|
||||
total_accounted = available + evictable + protected + session_held + uncached
|
||||
leak = total_accounted != total
|
||||
msg = (
|
||||
f"[{pool_name}] {total=}, {available=}, {evictable=}, "
|
||||
f"{protected=}, {session_held=}, {uncached=}"
|
||||
)
|
||||
return leak, msg
|
||||
|
||||
@staticmethod
|
||||
def _check_full_pool(
|
||||
self: "SchedulerInvariantChecker", ps: PoolStats, uncached: int = 0
|
||||
) -> Tuple[bool, str]:
|
||||
if self.is_hybrid_swa:
|
||||
protected = self.tree_cache.full_protected_size()
|
||||
session_held = self.pool_stats_observer.session_held_full_tokens()
|
||||
total = self.full_tokens_per_layer
|
||||
elif self.is_hybrid_ssm and self.tree_cache.supports_mamba():
|
||||
protected = self.tree_cache.full_protected_size()
|
||||
session_held = self.pool_stats_observer.session_held_tokens()
|
||||
total = self.token_to_kv_pool_allocator.size
|
||||
else:
|
||||
protected = self.tree_cache.protected_size()
|
||||
session_held = self.pool_stats_observer.session_held_tokens()
|
||||
total = self.max_total_num_tokens
|
||||
return SchedulerRuntimeCheckerMixin._check_pool_invariant(
|
||||
"full",
|
||||
ps.full_available_size,
|
||||
ps.full_evictable_size,
|
||||
protected,
|
||||
session_held,
|
||||
total,
|
||||
uncached,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _check_swa_pool(
|
||||
self: "SchedulerInvariantChecker", ps: PoolStats, uncached: int = 0
|
||||
) -> Tuple[bool, str]:
|
||||
return SchedulerRuntimeCheckerMixin._check_pool_invariant(
|
||||
"swa",
|
||||
ps.swa_available_size,
|
||||
ps.swa_evictable_size,
|
||||
self.tree_cache.swa_protected_size(),
|
||||
self.pool_stats_observer.session_held_swa_tokens(),
|
||||
self.swa_tokens_per_layer,
|
||||
uncached,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _check_mamba_pool(
|
||||
self: "SchedulerInvariantChecker", ps: PoolStats
|
||||
) -> Tuple[bool, str]:
|
||||
leak, msg = SchedulerRuntimeCheckerMixin._check_pool_invariant(
|
||||
"mamba",
|
||||
ps.mamba_available_size,
|
||||
ps.mamba_evictable_size,
|
||||
self.tree_cache.mamba_protected_size(),
|
||||
self.pool_stats_observer.session_held_mamba_slots(),
|
||||
self.req_to_token_pool.mamba_pool.size,
|
||||
)
|
||||
if leak:
|
||||
# Page-level leak diagnosis for mamba
|
||||
free_full_pages = set(
|
||||
self.token_to_kv_pool_allocator.free_pages.tolist()
|
||||
+ self.token_to_kv_pool_allocator.release_pages.tolist()
|
||||
)
|
||||
cached_full_pages = set(self.tree_cache.all_values_flatten().tolist())
|
||||
expected_full_pages = set(
|
||||
range(1, self.token_to_kv_pool_allocator.size + 1)
|
||||
)
|
||||
leaked_full_pages = (
|
||||
expected_full_pages - free_full_pages - cached_full_pages
|
||||
)
|
||||
free_mamba_pages = set(
|
||||
self.req_to_token_pool.mamba_pool.free_slots.tolist()
|
||||
)
|
||||
cached_mamba_pages = set(
|
||||
self.tree_cache.all_mamba_values_flatten().tolist()
|
||||
)
|
||||
expected_mamba_pages = set(range(self.req_to_token_pool.mamba_pool.size))
|
||||
leaked_mamba_pages = (
|
||||
expected_mamba_pages - free_mamba_pages - cached_mamba_pages
|
||||
)
|
||||
msg += (
|
||||
f", leaked_full_pages={leaked_full_pages or None}"
|
||||
f", leaked_mamba_pages={leaked_mamba_pages or None}"
|
||||
)
|
||||
return leak, msg
|
||||
|
||||
@staticmethod
|
||||
def _get_total_uncached_sizes(
|
||||
self: "SchedulerInvariantChecker",
|
||||
) -> Tuple[int, int]:
|
||||
"""Sum uncached tokens for full and SWA pools across all active batches.
|
||||
|
||||
Returns (full_uncached, swa_uncached). For non-SWA models, swa_uncached is 0.
|
||||
|
||||
For full pool: uncached = allocated - cache_protected_len
|
||||
For SWA pool: uncached = allocated - max(cache_protected_len, swa_evicted_seqlen)
|
||||
"""
|
||||
# After decode: running_batch IS last_batch (same object), count once.
|
||||
# After prefill: they differ, both hold uncached tokens.
|
||||
batches = [self.get_last_batch()]
|
||||
if (
|
||||
self.get_running_batch() not in (None, self.get_last_batch())
|
||||
and not self.get_running_batch().is_empty()
|
||||
):
|
||||
batches.append(self.get_running_batch())
|
||||
|
||||
full_uncached = 0
|
||||
swa_uncached = 0
|
||||
for batch in batches:
|
||||
for req in batch.reqs:
|
||||
assert req.kv_committed_freed == req.kv_overallocated_freed
|
||||
if req.kv_committed_freed or req.req_pool_idx is None:
|
||||
continue
|
||||
|
||||
allocated_len = req.kv_allocated_len
|
||||
if self.page_size > 1:
|
||||
allocated_len = ceil_align(allocated_len, self.page_size)
|
||||
assert req.cache_protected_len % self.page_size == 0
|
||||
|
||||
full_uncached += allocated_len - req.cache_protected_len
|
||||
if self.is_hybrid_swa:
|
||||
swa_uncached += allocated_len - max(
|
||||
req.cache_protected_len, req.swa_evicted_seqlen
|
||||
)
|
||||
|
||||
return full_uncached, swa_uncached
|
||||
|
||||
@staticmethod
|
||||
def self_check_during_busy(self: "SchedulerInvariantChecker"):
|
||||
if self.get_last_batch() is None:
|
||||
return
|
||||
|
||||
spec_topk = self.server_args.speculative_eagle_topk or 1
|
||||
if spec_topk > 1:
|
||||
warnings.warn(
|
||||
"Runtime memory check (busy) is not supported when speculation topk > 1."
|
||||
)
|
||||
return
|
||||
|
||||
ps = self.pool_stats_observer.get_pool_stats()
|
||||
full_uncached, swa_uncached = (
|
||||
SchedulerRuntimeCheckerMixin._get_total_uncached_sizes(self)
|
||||
)
|
||||
|
||||
full_leak, full_msg = SchedulerRuntimeCheckerMixin._check_full_pool(
|
||||
self, ps, uncached=full_uncached
|
||||
)
|
||||
|
||||
swa_leak, swa_msg = False, ""
|
||||
if self.is_hybrid_swa:
|
||||
swa_leak, swa_msg = SchedulerRuntimeCheckerMixin._check_swa_pool(
|
||||
self, ps, uncached=swa_uncached
|
||||
)
|
||||
|
||||
if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get() > 1:
|
||||
logger.info(f"[Mem Check (BUSY)] {full_msg}")
|
||||
if swa_msg:
|
||||
logger.info(f"[Mem Check (BUSY)] {swa_msg}")
|
||||
assert not full_leak, f"Full Pool Mem Leak Detected! {full_msg}"
|
||||
assert not swa_leak, f"SWA Pool Mem Leak Detected! {swa_msg}"
|
||||
|
||||
@staticmethod
|
||||
def _check_req_pool(self: "SchedulerInvariantChecker"):
|
||||
if self.disaggregation_mode == DisaggregationMode.DECODE:
|
||||
req_total_size = (
|
||||
self.req_to_token_pool.size + self.req_to_token_pool.pre_alloc_size
|
||||
)
|
||||
else:
|
||||
req_total_size = self.req_to_token_pool.size
|
||||
|
||||
session_req_count = self.pool_stats_observer.session_held_req_count()
|
||||
if len(self.req_to_token_pool.free_slots) + session_req_count != req_total_size:
|
||||
msg = (
|
||||
"req_to_token_pool memory leak detected!"
|
||||
f"available_size={len(self.req_to_token_pool.free_slots)}, "
|
||||
f"session_held={session_req_count}, "
|
||||
f"total_size={self.req_to_token_pool.size}\n"
|
||||
)
|
||||
raise_error_or_warn(
|
||||
self,
|
||||
envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
|
||||
"count_req_pool_leak_warnings",
|
||||
msg,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _report_leak(self: "SchedulerInvariantChecker", pool_name: str, token_msg: str):
|
||||
msg = f"{pool_name} memory leak detected! {token_msg}"
|
||||
raise_error_or_warn(
|
||||
self,
|
||||
envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_IDLE.get(),
|
||||
"count_memory_leak_warnings",
|
||||
msg,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _check_all_pools(
|
||||
self: "SchedulerInvariantChecker", ps: PoolStats, uncached: int = 0
|
||||
) -> Tuple[bool, List[str]]:
|
||||
"""Check memory invariant across all pools. Returns (has_leak, messages)."""
|
||||
has_leak = False
|
||||
messages = []
|
||||
|
||||
full_leak, full_msg = SchedulerRuntimeCheckerMixin._check_full_pool(
|
||||
self, ps, uncached=uncached
|
||||
)
|
||||
has_leak |= full_leak
|
||||
messages.append(full_msg)
|
||||
|
||||
if self.is_hybrid_swa:
|
||||
swa_leak, swa_msg = SchedulerRuntimeCheckerMixin._check_swa_pool(self, ps)
|
||||
has_leak |= swa_leak
|
||||
messages.append(swa_msg)
|
||||
|
||||
if self.is_hybrid_ssm and self.tree_cache.supports_mamba():
|
||||
mamba_leak, mamba_msg = SchedulerRuntimeCheckerMixin._check_mamba_pool(
|
||||
self, ps
|
||||
)
|
||||
has_leak |= mamba_leak
|
||||
messages.append(mamba_msg)
|
||||
|
||||
return has_leak, messages
|
||||
|
||||
def _maybe_log_idle_metrics(self: Scheduler):
|
||||
"""Collect and log metrics every 30 seconds during idle."""
|
||||
if (
|
||||
@@ -297,12 +54,3 @@ class SchedulerRuntimeCheckerMixin:
|
||||
self.disagg_decode_transfer_queue.queue, priority_enabled
|
||||
)
|
||||
self.metrics_collector.log_stats(self.stats)
|
||||
|
||||
@staticmethod
|
||||
def _check_tree_cache(self: "SchedulerInvariantChecker"):
|
||||
if (
|
||||
self.tree_cache.is_tree_cache()
|
||||
and (self.is_hybrid_swa and self.tree_cache.supports_swa())
|
||||
or (self.is_hybrid_ssm and self.tree_cache.supports_mamba())
|
||||
):
|
||||
self.tree_cache.sanity_check()
|
||||
|
||||
Reference in New Issue
Block a user