feat(worker): KVCacheManager stats dual-path gate (Phase 5.1b, refs #175) - #194
Merged
Merged
Conversation
…se 5.1b, refs #175) Phase 2 / Slice 5 / sub-slice 1b — dual-path gate. Wires the read-only stats helpers landed in PR #193 into the worker behind BATCHGEN_WORKER_KV_STATS_NATIVE / ..._COMPARE env vars. Surface added: - Import KVCacheManager, KVStats, KVUtilizationRequest + dataclasses - `_kv_stats_native` + `_kv_stats_compare` env reads at __init__ - `_kv_cache_manager` lazy field (KVCacheManager constructed on first use) - `_make_kv_cache_manager()` — wires a TorchKVStatsBackend that reads from `host_paged_kv_worker_view` and `gpu_paged_kv_cache_manager` - `_make_kv_utilization_request()` — frozen snapshot per call - Gate wrappers around 3 methods (preserve original return types) - 3 `_legacy_<name>` methods (preserve bodies verbatim) Net: +87 LOC on batchgen_worker.py. Design note: the legacy `_get_host_kv_utilization` returns a dict; the new `KVCacheManager.get_host_utilization` returns a `HostKVUtilization` dataclass. The gate wrapper converts the dataclass via `dataclasses.asdict()` so callers see the same dict shape; the field names already match exactly. Phase 5.1c will keep the dict return type to avoid touching the 4 caller sites — only the legacy body gets deleted. Refs #175. Part of milestone Worker decouple.
10 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Phase 2 / Slice 5 / sub-slice 1b — dual-path gate step. Issue #175. Stacked on the Slice 5.1 port (PR #193, merged).
Wires the read-only stats helpers into the worker behind
BATCHGEN_WORKER_KV_STATS_NATIVE/..._COMPAREenv vars.Surface added (+87 LOC)
KVCacheManager,KVStats,KVUtilizationRequest,dataclasses)__init__: env reads + lazy_kv_cache_managerfield + comment_make_kv_cache_manager()(wiresTorchKVStatsBackendreading fromhost_paged_kv_worker_view+gpu_paged_kv_cache_manager)_make_kv_utilization_request()_get_host_kv_free_pages,_get_gpu_kv_free_pages,_get_host_kv_utilization)_legacy_<name>methods (preserve bodies verbatim)Design note: dict ↔ dataclass adapter
Legacy
_get_host_kv_utilizationreturns aDict[str, int]; the newKVCacheManager.get_host_utilizationreturns aHostKVUtilizationdataclass. The gate wrapper converts viadataclasses.asdict()so the 4 caller sites see the same dict shape (field names already match exactly). Phase 5.1c will keep the dict return type to avoid touching callers — only the legacy body gets deleted.Motivation
Refs #175. Part of milestone Worker decouple.
Type of Change
(Migration-window infrastructure. Production behavior unchanged when both env vars are unset.)
Checklist
tests/worker/test_kv_manager.pycoversKVCacheManagerdirectly. Gate validated via H20 dual-mode.Pre-merge verification
python -m py_compile batchgen/batchgen_worker.pysucceeds.pytest tests/worker/ -x→ 67/67 passing.pr-gpu-smoke.yml— runs on PR push.COMPARE=1 NATIVE=1— parity green; accuracy in sampling band.