Vllm model handler #32410
Vllm model handler #32410
Conversation
|
Assigning reviewers. If you would like to opt out of this review, comment R: @shunping for label python. Available commands:
The PR bot will only process comments in the main thread (not review comments). |
| formatted = [] | ||
| for message in messages: | ||
| formatted.append({"role": message.role, "content": message.content}) | ||
| completion = client.chat.completions.create( |
There was a problem hiding this comment.
What happens if an server exception occurs during the query?
There was a problem hiding this comment.
It will bubble up as an exception and either be retried or sent to a DLQ depending on user configuration. With that said, previously if the server died it would have just stayed dead. I added some logic to handle/avoid that problem by restarting the server if we can't connect.
| client = getVLLMClient(model.get_server_port()) | ||
| inference_args = inference_args or {} | ||
| predictions = [] | ||
| # TODO(https://github.com/apache/beam/issues/32528): We should add support |
There was a problem hiding this comment.
Do they support batch mode in the query?
There was a problem hiding this comment.
They do not as far as I can tell, unfortunately. I plan on addressing this in a follow up pr though - with vLLMs dynamic batching it still almost certainly makes sense to do something here.
This PR adds a model handler for running inference using vLLM. To leverage vLLM's dynamic batching, this spins up a central vLLM serving process, coordinated by a single global model wrapper, which individual worker threads can send RPCs to.
Testing this is tricky since it requires a gpu; I followed the same pattern we use for TensorRT - just launch some examples directly on Dataflow and validate that they run to completion. I didn't do any kind of result validation since results produced by a LLM are non-deterministic.
Part of #32528
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.