[Bugfix][Core] Restore logging of stats in the async engine

[ronensc]

Following the refactor of #3894, I noticed that the async engine has stopped logging stats to Prometheus.

The reason is that the call to stat_logger.log() was relocated from the _process_model_outputs() method (which is shared between async and non-async engines) to LLMEngine.step(), exclusively used in the non-async engine.

This PR fixes this by reintroducing the call to stat_logger.log() in the async engine.

cc @cadedaniel

[ronensc]

Would it be worth considering moving the call to stat_logger.log() into _process_model_outputs()?
Currently, it's being invoked in both LLMEngine.step() and AsyncLLMEngine.step_async(), duplicating the functionality.

vllm/vllm/engine/llm_engine.py

Lines 542 to 544 in fe3b5bb

	# Log stats.
	if self.log_stats:
	self.stat_logger.log(self._get_stats(scheduler_outputs))

What are your thoughts?

[rkooo567]

feel like having it here makes more sense for abstraction? it is technically not related to process model output

[rkooo567]

alternatively, we can create a method like post_process_step() and add logging logic here and just call it in both sync/async?

[sfc-gh-zhwang]

it is technically not related to process model output

it's a logging to me and can be fold into anywhere lol..

i feel it's simple to just fold it into _process_model_outputs, maybe make _process_model_outputs accept scheduler_outputs instead of scheduled_seq_groups and ignored_seq_groups?

[rkooo567]

IMO relevant logging should go to the relevant API. E.g., logging about e2e stats is a little bit weird to be in process model output API, but it makes sense to be in step API. It is fine now because of all assuptions we have (process_model_output is called only by step), so not very strong opinion or a blocker at the moment

[rkooo567]

Is it possible to add a simple regression test?

[rkooo567]

feel like having it here makes more sense for abstraction? it is technically not related to process model output

[sfc-gh-zhwang]

Nit: i think you can just do the logging before below instead of assigning it to request_outputs

        return self._process_model_outputs(
            output, scheduler_outputs.scheduled_seq_groups,
            scheduler_outputs.ignored_seq_groups)

[ronensc]

I'm not sure it will work, as stat_logger.log() should be called after seq_group.maybe_set_first_token_time(now).

vllm/vllm/engine/llm_engine.py

Line 467 in e8cc796

seq_group.maybe_set_first_token_time(now)

[sfc-gh-zhwang]

is first_token_time used in stat_logger? https://github.com/search?q=repo%3Avllm-project%2Fvllm%20first_token_time&type=code

[ronensc]

My bad, It's not. However, I still think that the logging should be after _process_model_outputs(). That's because during logging, we calculate the latency between tokens, which I believe should include the processing of the model outputs.

vllm/vllm/engine/llm_engine.py

Lines 601 to 603 in fe3b5bb

	# Time since last token.
	# (n.b. updates seq_group.metrics.last_token_time)
	time_last_iters.append(seq_group.get_last_latency(now))

[rkooo567]

Let's merge it after adding a simple regression test !

[mgoin]

It seems the openai api server is not reporting tokens/s metric correctly anymore (on main as of today):

INFO 04-19 13:38:10 metrics.py:224] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 0.0 tokens/s, Running: 49 reqs, Swapped: 0 reqs, Pending: 981 reqs, GPU KV cache usage: 4.5%, CPU KV cache usage: 0.0%
INFO 04-19 13:38:20 metrics.py:224] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 0.0 tokens/s, Running: 79 reqs, Swapped: 0 reqs, Pending: 951 reqs, GPU KV cache usage: 6.9%, CPU KV cache usage: 0.0%
INFO 04-19 13:38:30 metrics.py:224] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 0.0 tokens/s, Running: 99 reqs, Swapped: 0 reqs, Pending: 931 reqs, GPU KV cache usage: 8.6%, CPU KV cache usage: 0.0%
INFO 04-19 13:38:40 metrics.py:224] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 0.0 tokens/s, Running: 129 reqs, Swapped: 0 reqs, Pending: 901 reqs, GPU KV cache usage: 11.2%, CPU KV cache usage: 0.0%
INFO 04-19 13:38:50 metrics.py:224] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 0.0 tokens/s, Running: 149 reqs, Swapped: 0 reqs, Pending: 881 reqs, GPU KV cache usage: 12.9%, CPU KV cache usage: 0.0%

This PR won't affect this, but any ideas?

[simon-mo]

Thank you for fixing it. I'm merging this to unblock release. Please follow up with testing which is desperately needed!