refactor: unify rollout onto inference engine KV-cache path
- RolloutGenerator now delegates prefill/decode to InferenceScheduler.run_batch (sync API, no background thread), sharing one KV-cache code path with the inference server and eliminating O(n^2) recompute in rollout - Add sample(return_logprobs=) and Executor.execute_decode(return_logprobs=) to expose behaviour-policy log-probs through the engine; Task gains output_logprobs - RolloutResult now subclasses RawRollout (adds rewards only), removing duplicated fields - RolloutRunner.__call__ returns (result, is_fresh) instead of relying on object identity, removing the fragile refresh-detection contract - Remove O(n^2) generate_responses helper and dead code (_tokenize_prompts, unused old_model arg) - train_context.py wires InferenceScheduler directly instead of hand-rolling SamplingPipeline - Tests: +11 covering return_logprobs, run_batch, and KV-cache-backed rollout semantics; 404 pass
This commit is contained in:
@@ -71,22 +71,31 @@ def _make_rollout_result(B=2, G=4, P=6, R=8, device="cpu"):
|
||||
|
||||
|
||||
class _RecordingRunner:
|
||||
"""Fake RolloutRunner that returns a fixed result and tracks calls."""
|
||||
"""Fake RolloutRunner returning a fixed result with freshness tracking.
|
||||
|
||||
Freshness is ``True`` on the first call after construction or after
|
||||
:meth:`swap_result`; ``False`` on subsequent cached calls — mirroring
|
||||
the real ``RolloutRunner`` contract without invoking generation.
|
||||
"""
|
||||
|
||||
def __init__(self, result):
|
||||
self.result = result
|
||||
self.calls = 0
|
||||
self.step_calls = 0
|
||||
self._fresh = True
|
||||
|
||||
def __call__(self, batch):
|
||||
self.calls += 1
|
||||
return self.result
|
||||
fresh = self._fresh
|
||||
self._fresh = False
|
||||
return self.result, fresh
|
||||
|
||||
def step(self):
|
||||
self.step_calls += 1
|
||||
|
||||
def swap_result(self, result):
|
||||
self.result = result
|
||||
self._fresh = True
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
||||
Reference in New Issue
Block a user