140 lines
4.3 KiB
Python
140 lines
4.3 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Dict, Iterator, List, Optional
|
|
|
|
if TYPE_CHECKING:
|
|
from sglang.srt.managers.schedule_batch import Req
|
|
from sglang.test.scripted_runtime.context.api import ScriptedContext
|
|
|
|
|
|
def _get_all_reqs(ctx: ScriptedContext) -> Iterator[Req]:
|
|
s = ctx.scheduler
|
|
if s.chunked_req is not None:
|
|
yield s.chunked_req
|
|
yield from s.waiting_queue
|
|
if s.ps.pp_size > 1:
|
|
for mb in (*s.mbs, *s.last_mbs, *s.running_mbs):
|
|
if mb is not None:
|
|
yield from mb.reqs
|
|
else:
|
|
if s.running_batch is not None:
|
|
yield from s.running_batch.reqs
|
|
if s.last_batch is not None:
|
|
yield from s.last_batch.reqs
|
|
|
|
|
|
def list_active_reqs(ctx: ScriptedContext) -> List[Req]:
|
|
return list(set(_get_all_reqs(ctx)))
|
|
|
|
|
|
def batch_composition(ctx: ScriptedContext) -> Dict[str, List[str]]:
|
|
s = ctx.scheduler
|
|
chunked_rid = s.chunked_req.rid if s.chunked_req is not None else None
|
|
chunked = [chunked_rid] if chunked_rid is not None else []
|
|
running = (
|
|
[r.rid for r in s.running_batch.reqs] if s.running_batch is not None else []
|
|
)
|
|
|
|
prefill: List[str] = []
|
|
decode: List[str] = []
|
|
batch = s.last_batch
|
|
if batch is not None and not batch.is_empty() and batch.forward_mode is not None:
|
|
bucket = prefill if batch.forward_mode.is_extend() else decode
|
|
bucket.extend(r.rid for r in batch.reqs if r.rid != chunked_rid)
|
|
|
|
return {
|
|
"prefill": prefill,
|
|
"decode": decode,
|
|
"chunked": chunked,
|
|
"running": running,
|
|
}
|
|
|
|
|
|
def is_idle(ctx: ScriptedContext) -> bool:
|
|
s = ctx.scheduler
|
|
return (
|
|
s.chunked_req is None
|
|
and len(s.waiting_queue) == 0
|
|
and (s.running_batch is None or s.running_batch.is_empty())
|
|
)
|
|
|
|
|
|
def is_fully_idle(ctx: ScriptedContext) -> bool:
|
|
s = ctx.scheduler
|
|
return is_idle(ctx) and (s.last_batch is None or s.last_batch.is_empty())
|
|
|
|
|
|
def last_batch_forward_mode(ctx: ScriptedContext) -> Optional[str]:
|
|
s = ctx.scheduler
|
|
if s.last_batch is not None and s.last_batch.forward_mode is not None:
|
|
return s.last_batch.forward_mode.name
|
|
return None
|
|
|
|
|
|
def find_req_by_rid(ctx: ScriptedContext, rid: str) -> Optional[Req]:
|
|
req = next((r for r in _get_all_reqs(ctx) if r.rid == rid), None)
|
|
if req is not None:
|
|
ctx._seen_rids.add(rid)
|
|
return req
|
|
|
|
|
|
def is_finished(ctx: ScriptedContext, rid: str) -> bool:
|
|
req = find_req_by_rid(ctx, rid)
|
|
if req is not None:
|
|
return req.finished()
|
|
if rid in ctx._seen_rids:
|
|
return True
|
|
# Fallback: if the req ran in a forward batch (recorded in _batch_log) but
|
|
# is now absent from all active scheduler sets, it must have finished.
|
|
# This catches requests that completed without ever being observed via
|
|
# find_req_by_rid (e.g. when Python short-circuit evaluation prevents the
|
|
# query while another request is still running).
|
|
log = ctx._scheduler_hook._batch_log
|
|
if any(rid in record.rids for record in log):
|
|
ctx._seen_rids.add(rid)
|
|
return True
|
|
return False
|
|
|
|
|
|
def is_chunking(ctx: ScriptedContext, rid: str) -> bool:
|
|
s = ctx.scheduler
|
|
return s.chunked_req is not None and s.chunked_req.rid == rid
|
|
|
|
|
|
def status(ctx: ScriptedContext, rid: str) -> str:
|
|
s = ctx.scheduler
|
|
if rid in {r.rid for r in s.waiting_queue}:
|
|
return "waiting"
|
|
req = find_req_by_rid(ctx, rid)
|
|
if req is not None:
|
|
return "finished" if req.finished() else "running"
|
|
if rid in ctx._seen_rids:
|
|
return "finished"
|
|
return "unknown"
|
|
|
|
|
|
def remaining_prompt_tokens(ctx: ScriptedContext, rid: str) -> int:
|
|
req = find_req_by_rid(ctx, rid)
|
|
if req is None:
|
|
return 0
|
|
return max(0, len(req.origin_input_ids) - req.kv_committed_len)
|
|
|
|
|
|
def chunks_done(ctx: ScriptedContext, rid: str) -> int:
|
|
log = ctx._scheduler_hook._batch_log
|
|
held = sum(1 for record in log if record.chunked_rid == rid and rid in record.rids)
|
|
if held == 0:
|
|
return 0
|
|
completed = any(
|
|
rid in record.extend_rids and record.chunked_rid != rid for record in log
|
|
)
|
|
return held + (1 if completed else 0)
|
|
|
|
|
|
def chunked_parks(ctx: ScriptedContext, rid: str) -> int:
|
|
return sum(
|
|
1
|
|
for record in ctx._scheduler_hook._batch_log
|
|
if record.chunked_rid == rid and rid not in record.rids
|
|
)
|