[Frontend] Add progress reporting to run_batch.py

[alugowski]

Add a tqdm progress indicator to run_batch.py.

run_batch.py runs requests from a OpenAI Batch protocol-compatible .jsonl file where each line is one request. Batches can have request counts in the hundreds of thousands and can take hours to complete.

Currently there is no way to tell how many requests are to be executed (short of externally inspecting the input file), what the current progress is, or when to expect it to finish.

A nice tqdm progress indicator is a great quality of life improvement, and enables external progress monitoring in automated workflows.

The progress indicator here is modeled after the shard loader tqdm in weight_utils.py.

Example output:

INFO 09-01 12:14:46 model_runner.py:913] Starting to load model /var/data/models/meta-llama/Meta-Llama-3.1-8B-Instruct...
Loading safetensors checkpoint shards:   0% Completed | 0/4 [00:00<?, ?it/s]
Loading safetensors checkpoint shards:  25% Completed | 1/4 [00:00<00:02,  1.45it/s]
Loading safetensors checkpoint shards:  50% Completed | 2/4 [00:01<00:01,  1.99it/s]
Loading safetensors checkpoint shards:  75% Completed | 3/4 [00:01<00:00,  1.65it/s]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:02<00:00,  1.40it/s]
Loading safetensors checkpoint shards: 100% Completed | 4/4 [00:02<00:00,  1.50it/s]

INFO 09-01 12:14:49 model_runner.py:924] Loading model weights took 14.9888 GB
INFO 09-01 12:14:49 gpu_executor.py:122] # GPU blocks: 9778, # CPU blocks: 2048
INFO 09-01 12:14:52 model_runner.py:1215] Capturing the model for CUDA graphs. This may lead to unexpected consequences if the model is not static. To run the model in eager mode, set 'enforce_eager=True' or use '--enforce-eager' in the CLI.
INFO 09-01 12:14:52 model_runner.py:1219] CUDA graphs can take additional 1~3 GiB memory per GPU. If you are running out of memory, consider decreasing `gpu_memory_utilization` or enforcing eager mode. You can also reduce the `max_num_seqs` as needed to decrease memory usage.
INFO 09-01 12:15:03 model_runner.py:1333] Graph capturing finished in 11 secs.
WARNING 09-01 12:15:03 serving_embedding.py:190] embedding_mode is False. Embedding API will not work.
INFO 09-01 12:15:03 run_batch.py:204] Reading batch from many.jsonl...
Running batch:   0% Completed | 0/938 [00:00<?, ?req/s]
INFO 09-01 12:15:08 metrics.py:351] Avg prompt throughput: 3911.7 tokens/s, Avg generation throughput: 1110.1 tokens/s, Running: 187 reqs, Swapped: 0 reqs, Pending: 534 reqs, GPU KV cache usage: 9.6%, CPU KV cache usage: 0.0%.
Running batch:  23% Completed | 215/938 [00:05<00:16, 42.71req/s]
INFO 09-01 12:15:13 metrics.py:351] Avg prompt throughput: 2794.5 tokens/s, Avg generation throughput: 4242.6 tokens/s, Running: 256 reqs, Swapped: 0 reqs, Pending: 257 reqs, GPU KV cache usage: 24.6%, CPU KV cache usage: 0.0%.
Running batch:  46% Completed | 429/938 [00:10<00:12, 40.90req/s]
INFO 09-01 12:15:19 metrics.py:351] Avg prompt throughput: 1199.3 tokens/s, Avg generation throughput: 4569.3 tokens/s, Running: 252 reqs, Swapped: 0 reqs, Pending: 129 reqs, GPU KV cache usage: 32.7%, CPU KV cache usage: 0.0%.
Running batch:  68% Completed | 634/938 [00:16<00:08, 37.18req/s]
INFO 09-01 12:15:24 metrics.py:351] Avg prompt throughput: 1312.6 tokens/s, Avg generation throughput: 4513.6 tokens/s, Running: 179 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 27.8%, CPU KV cache usage: 0.0%.
Running batch:  88% Completed | 822/938 [00:24<00:03, 31.17req/s]
INFO 09-01 12:15:29 metrics.py:351] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 3889.4 tokens/s, Running: 106 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 26.3%, CPU KV cache usage: 0.0%.
INFO 09-01 12:15:34 metrics.py:351] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 2867.9 tokens/s, Running: 70 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 27.4%, CPU KV cache usage: 0.0%.
INFO 09-01 12:15:39 metrics.py:351] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 2237.7 tokens/s, Running: 49 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 24.0%, CPU KV cache usage: 0.0%.
Running batch:  99% Completed | 928/938 [00:40<00:00, 31.17req/s]
Running batch:  99% Completed | 929/938 [00:40<00:00, 17.01req/s]
INFO 09-01 12:15:44 metrics.py:351] Avg prompt throughput: 0.0 tokens/s, Avg generation throughput: 1336.6 tokens/s, Running: 5 reqs, Swapped: 0 reqs, Pending: 0 reqs, GPU KV cache usage: 3.0%, CPU KV cache usage: 0.0%.
Running batch: 100% Completed | 938/938 [00:41<00:00, 22.35req/s]

INFO 09-01 12:15:45 async_llm_engine.py:60] Engine is gracefully shutting down.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[alugowski]

/ready