Emit vLLM engine stats from the RL rollout sampler

[wenxindongwork]

What

Starts a periodic do_log_stats() task in RLVllmSampler.start() and cancels it in stop(), so the rollout logs vLLM's standard engine stats line:

Avg prompt throughput: ... tokens/s, Avg generation throughput: ... tokens/s,
Running: N reqs, Waiting: N reqs, GPU KV cache usage: ...%, Prefix cache hit rate: ...%

Why

That line never appears in rollout logs — not once. It isn't emitted by the engine on a timer; it's emitted by whoever calls do_log_stats(), and in this configuration nobody does.

LoggingStatLogger.log() is reachable from exactly one place in AsyncLLM:

# vllm/v1/engine/async_llm.py:1021
async def do_log_stats(self) -> None:
    if self.logger_manager:
        self.logger_manager.log()

do_log_stats has no internal caller. The ticker lives in the API server's lifespan:

# vllm/entrypoints/launchers/utils/server_utils.py:70-76
async def _force_log():
    while True:
        await asyncio.sleep(envs.VLLM_LOG_STATS_INTERVAL)
        await engine_client.do_log_stats()

AsyncLLM starts only the input-stream and output handlers; neither logs stats. (The synchronous LLMEngine is different — it drives do_log_stats_with_interval() from its own step loop.) RLVllmSampler binds AsyncLLMEngine directly with no server, so the stats are recorded and never flushed.

This is not a config problem: disable_log_stats defaults to False, so log_stats=True and the StatLoggerManager is already constructed and recording.

There is already a precedent in this repo — VLLMInProcessDriver spawns a thread for exactly this (_log_loop, tunix/generate/vllm_async_driver.py:271-282). The v2 sampler never got the equivalent. Since its engine is async, this uses an asyncio task mirroring the server's _force_log rather than a thread.

Testing

Drove _log_stats_loop against a stub engine with VLLM_LOG_STATS_INTERVAL=0.1:

do_log_stats calls: 5     # over 0.55s
task cleared: True
engine cleared: True

Confirms the cadence, and that stop() cancels the task cleanly (no CancelledError escaping, no pending-task warning).

Not yet exercised against a real engine on TPU, hence draft. Note log() downgrades itself to logger.debug when engine_is_idle (vllm/v1/metrics/loggers.py:260), so expect gaps between rollout batches rather than a strict every-10s line.