[mem] Introduce PoolStats dataclass; unify pool metrics and token_usage (#22554)
This commit is contained in:
@@ -9,6 +9,7 @@ maybe_stub_sgl_kernel()
|
||||
|
||||
from sglang.srt.managers.io_struct import PauseGenerationReqInput
|
||||
from sglang.srt.managers.scheduler import Scheduler
|
||||
from sglang.srt.managers.scheduler_runtime_checker_mixin import PoolStats
|
||||
|
||||
register_cpu_ci(est_time=10, suite="stage-a-test-cpu")
|
||||
|
||||
@@ -33,7 +34,14 @@ class TestSchedulerPauseGeneration(unittest.TestCase):
|
||||
scheduler.token_to_kv_pool_allocator = MagicMock()
|
||||
scheduler.token_to_kv_pool_allocator.available_size.return_value = 1000
|
||||
scheduler.max_total_num_tokens = 1000
|
||||
scheduler._get_token_info = MagicMock(return_value=(0, 0, 1000, 0))
|
||||
scheduler._get_token_info = MagicMock(
|
||||
return_value=PoolStats(
|
||||
full_num_used=0,
|
||||
full_token_usage=0,
|
||||
full_available_size=1000,
|
||||
full_evictable_size=0,
|
||||
)
|
||||
)
|
||||
return scheduler
|
||||
|
||||
def test_inplace_only_sets_flag(self):
|
||||
|
||||
Reference in New Issue
Block a user