[Doc] RayServe Single-Host TPU v6e Example with vLLM

[ryanaoleary]

Why are these changes needed?

This PR adds a new guide to the Ray docs that details how to serve an LLM with vLLM and single-host TPUs on GKE. This PR has been tested by running through the steps in the proposed guide and verifying correct output. This PR uses sample code from https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/tree/main/ai-ml/gke-ray/rayserve/llm/tpu.

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[ryanaoleary]

cc: @andrewsykim

[stale]

This pull request has been automatically marked as stale because it has not had recent activity. It will be closed in 14 days if no further activity occurs. Thank you for your contributions.

• If you'd like to keep this open, just leave any comment, and the stale label will be removed.

[ryanaoleary]

I went through and updated this guide for the newer TPU Trillium so that it'll be more useful, it no longer requires multi-host since v6e can fit Llama 70B on a single node. I'll create a PR with a separate guide showcasing serving with Llama 405B and multi-host TPUs. cc: @andrewsykim @kevin85421

[andrewsykim]

@angelinalg can you help us review this please?

[kouroshHakha]

Hi @ryanaoleary , Thanks for the contribution. We just ossed our native ray serve LLM apis. Do you mind converting this tutorial to use that instead?

https://docs.ray.io/en/latest/serve/llm/overview.html

Basically there is no need to reference a custom application code that is hosted in the kuberay code. Please use the latest apis from master.

Also another question is if this tutorial is part of the ci/cd. For example what will happen if vllm / some api from ray changes?

[andrewsykim]

@kouroshHakha thanks for letting us know!

There's actually some existing docs that use KubeRay + vLLM right now, for example see https://docs.ray.io/en/latest/cluster/kubernetes/examples/vllm-rayservice.html. Would it be possible to merge Ryan's PR as-is and follow-up with a sweeping change that updates all the docs using vLLM for consistency?

[kouroshHakha]

Yes we can merge this as it's a self-contained, and the code that most likely won't break in the future so this is a plus for this approach. But can we add a note on the very top to both new files in this PR and the existing kubernetes example (at the same level and visibility as where prerequisites are listed) in this tutorial that mentions there is a native serve LLM API with some pointers so that the reader does go under the impression that this is the only application code that works? In fact the new API is more recommended path for running LLMs through ray serve. The kubernetes deployment configs remain the same and only application code changes to use those APIs. If we do this then there is no need for a sweeping change now as it would be more clear to the reader.

[andrewsykim]

@kouroshHakha does the new LLM API support TPUs?

[ryanaoleary]

Yes we can merge this as it's a self-contained, and the code that most likely won't break in the future so this is a plus for this approach. But can we add a note on the very top to both new files in this PR and the existing kubernetes example (at the same level and visibility as where prerequisites are listed) in this tutorial that mentions there is a native serve LLM API with some pointers so that the reader does go under the impression that this is the only application code that works? In fact the new API is more recommended path for running LLMs through ray serve. The kubernetes deployment configs remain the same and only application code changes to use those APIs. If we do this then there is no need for a sweeping change now as it would be more clear to the reader.

Added a note in the vLLM examples in 466ad03 about the new Ray Serve LLM API.

[kouroshHakha]

does the new LLM API support TPUs?

@andrewsykim

Right now users can call config.get_serve_options() to get the deployment options (including the placement group) That function currently does not support TPUs. But you can bypass that by providing your own placement group similar to how it's done currently in https://github.com/ray-project/kuberay/blob/master/ray-operator/config/samples/vllm/serve_tpu.py#L102 and calling .options on the deployment object. Everything else remains the same similar to your example.

It would be awesome if you / someone from your team could contribute and extend the current placement group creation within ray.serve.llm to support TPUs as well.

[kouroshHakha]

LGTM. Thanks. 🚀 The doc build is failing tho. It might be the reference that is messed up. Once fixed we can merge.

[ryanaoleary]

LGTM. Thanks. 🚀 The doc build is failing tho. It might be the reference that is messed up. Once fixed we can merge.

Oh oops, I fixed the links. The build looks to be passing now.

[kouroshHakha]

cc @angelinalg can you stamp this PR? It needs docs approval

[angelinalg]

stamp

[kevin85421]

Hi @ryanaoleary, I will revert this PR because my colleague merges this PR accidentally. We can discuss the next step of this PR in our 1:1 next time.