Move create_scheduler_watchdog from runtime_checker mixin to scheduler.py (#25622)
This commit is contained in:
@@ -187,7 +187,6 @@ from sglang.srt.managers.scheduler_pp_mixin import SchedulerPPMixin
|
|||||||
from sglang.srt.managers.scheduler_recv_skipper import SchedulerRecvSkipper
|
from sglang.srt.managers.scheduler_recv_skipper import SchedulerRecvSkipper
|
||||||
from sglang.srt.managers.scheduler_runtime_checker_mixin import (
|
from sglang.srt.managers.scheduler_runtime_checker_mixin import (
|
||||||
SchedulerRuntimeCheckerMixin,
|
SchedulerRuntimeCheckerMixin,
|
||||||
create_scheduler_watchdog,
|
|
||||||
)
|
)
|
||||||
from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length
|
from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length
|
||||||
from sglang.srt.mem_cache import kv_cache_builder
|
from sglang.srt.mem_cache import kv_cache_builder
|
||||||
@@ -238,6 +237,7 @@ from sglang.srt.utils.network import get_zmq_socket
|
|||||||
from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node
|
from sglang.srt.utils.numa_utils import get_numa_node_if_available, numa_bind_to_node
|
||||||
from sglang.srt.utils.tensor_bridge import use_mlx
|
from sglang.srt.utils.tensor_bridge import use_mlx
|
||||||
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
|
from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
|
||||||
|
from sglang.srt.utils.watchdog import WatchdogRaw
|
||||||
from sglang.utils import TypeBasedDispatcher, get_exception_traceback
|
from sglang.utils import TypeBasedDispatcher, get_exception_traceback
|
||||||
|
|
||||||
if is_mps():
|
if is_mps():
|
||||||
@@ -322,6 +322,30 @@ def validate_dflash_request(req: Req) -> Optional[str]:
|
|||||||
return None
|
return None
|
||||||
|
|
||||||
|
|
||||||
|
def create_scheduler_watchdog(
|
||||||
|
scheduler: "Scheduler", watchdog_timeout: float, soft: bool = False
|
||||||
|
) -> WatchdogRaw:
|
||||||
|
def dump_info() -> str:
|
||||||
|
if scheduler.is_initializing:
|
||||||
|
return ""
|
||||||
|
_, messages = scheduler._check_all_pools(
|
||||||
|
scheduler.pool_stats_observer.get_pool_stats()
|
||||||
|
)
|
||||||
|
return (
|
||||||
|
f"{scheduler.cur_batch.batch_size()=}\n"
|
||||||
|
f"{scheduler.cur_batch.reqs=}\n" + "\n".join(messages)
|
||||||
|
)
|
||||||
|
|
||||||
|
return WatchdogRaw(
|
||||||
|
debug_name="Scheduler",
|
||||||
|
get_counter=lambda: scheduler.forward_ct,
|
||||||
|
is_active=lambda: scheduler.is_initializing or scheduler.cur_batch is not None,
|
||||||
|
watchdog_timeout=watchdog_timeout,
|
||||||
|
soft=soft,
|
||||||
|
dump_info=dump_info,
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
class Scheduler(
|
class Scheduler(
|
||||||
SchedulerOutputProcessorMixin,
|
SchedulerOutputProcessorMixin,
|
||||||
SchedulerMetricsMixin,
|
SchedulerMetricsMixin,
|
||||||
|
|||||||
@@ -9,7 +9,6 @@ from sglang.srt.disaggregation.utils import DisaggregationMode
|
|||||||
from sglang.srt.environ import envs
|
from sglang.srt.environ import envs
|
||||||
from sglang.srt.observability.metrics_collector import QueueCount
|
from sglang.srt.observability.metrics_collector import QueueCount
|
||||||
from sglang.srt.utils.common import ceil_align, raise_error_or_warn
|
from sglang.srt.utils.common import ceil_align, raise_error_or_warn
|
||||||
from sglang.srt.utils.watchdog import WatchdogRaw
|
|
||||||
|
|
||||||
if TYPE_CHECKING:
|
if TYPE_CHECKING:
|
||||||
from sglang.srt.managers.scheduler import Scheduler
|
from sglang.srt.managers.scheduler import Scheduler
|
||||||
@@ -281,27 +280,3 @@ class SchedulerRuntimeCheckerMixin:
|
|||||||
or (self.is_hybrid_ssm and self.tree_cache.supports_mamba())
|
or (self.is_hybrid_ssm and self.tree_cache.supports_mamba())
|
||||||
):
|
):
|
||||||
self.tree_cache.sanity_check()
|
self.tree_cache.sanity_check()
|
||||||
|
|
||||||
|
|
||||||
def create_scheduler_watchdog(
|
|
||||||
scheduler: Scheduler, watchdog_timeout: float, soft: bool = False
|
|
||||||
) -> WatchdogRaw:
|
|
||||||
def dump_info() -> str:
|
|
||||||
if scheduler.is_initializing:
|
|
||||||
return ""
|
|
||||||
_, messages = scheduler._check_all_pools(
|
|
||||||
scheduler.pool_stats_observer.get_pool_stats()
|
|
||||||
)
|
|
||||||
return (
|
|
||||||
f"{scheduler.cur_batch.batch_size()=}\n"
|
|
||||||
f"{scheduler.cur_batch.reqs=}\n" + "\n".join(messages)
|
|
||||||
)
|
|
||||||
|
|
||||||
return WatchdogRaw(
|
|
||||||
debug_name="Scheduler",
|
|
||||||
get_counter=lambda: scheduler.forward_ct,
|
|
||||||
is_active=lambda: scheduler.is_initializing or scheduler.cur_batch is not None,
|
|
||||||
watchdog_timeout=watchdog_timeout,
|
|
||||||
soft=soft,
|
|
||||||
dump_info=dump_info,
|
|
||||||
)
|
|
||||||
|
|||||||
Reference in New Issue
Block a user