feat(api_server): Add OpenAI-compatible API server for MaxText models #2313
+2,134
−0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
babyplutokurt
requested review from
SujeethJinesh,
bvandermoon,
richjames0,
shralex,
vipannalla,
mitalisi,
RissyRan,
shauryagup,
yangyuwei,
parambole,
gobbleturk,
khatwanimohit,
gagika,
SurbhiJainUSC,
hengtaoguo,
A9isha,
aireenmei and
NuojCheng
as code owners
babyplutokurt
force-pushed
the
api_sever_v1
branch
5 times, most recently
from
400ee8c to
ce5fab1
Compare
RissyRan
reviewed
Sep 11, 2025
babyplutokurt
force-pushed
the
api_sever_v1
branch
5 times, most recently
from
b70e083 to
18d3055
Compare
RissyRan
reviewed
Sep 12, 2025
hengtaoguo
approved these changes
Sep 12, 2025
babyplutokurt
force-pushed
the
api_sever_v1
branch
5 times, most recently
from
d1b5ad2 to
a29735a
Compare
babyplutokurt
requested review from
gpolovets1,
mailvijayasingh,
jrplatin,
patemotter,
Lumosis and
michelle-yooh
as code owners
babyplutokurt
force-pushed
the
api_sever_v1
branch
3 times, most recently
from
000872a to
750ffc2
Compare
RissyRan
reviewed
Sep 16, 2025
There was a problem hiding this comment.
I am fine to merge those at this moment (in a separate file, and no breakage for existing codebase), but could @hengtaoguo or @bvandermoon help test and verify end-to-end? Currently no tests for those scripts and functionality yet.
| # engine use its default configured `decode_sampling_strategy`. | ||
| return None | ||
|
|
||
| def _run_prefill_step(self, streams, decode_state, rng, logprobs, echo, temperature, top_k, top_p): |
There was a problem hiding this comment.
I agree that when involving with server, it may be difficult to mock. I guess my point is for those helper functions, we should test at least.
Since MaxEnginer has been in the codebase and verified many times
I think you make assumption that you verified many times, but not us? How do you ensure it's working end-to-end. I think currently only you tested those server and script. We have to "trust" you if no tests.
babyplutokurt
force-pushed
the
api_sever_v1
branch
3 times, most recently
from
92956a1 to
c703914
Compare
This commit introduces a fully-featured, OpenAI-compatible RESTful API server for serving MaxText models. The server is built with FastAPI, supports multi-host inference on TPUs, and is designed for both interactive use and large-scale benchmarking.
Key features and additions:
1. **Core Server Implementation:**
- Adds `maxtext_server.py`, a FastAPI application that serves `/v1/completions` and `/v1/chat/completions` endpoints.
- Implements dynamic request batching to efficiently utilize underlying hardware.
- Uses `maxtext_generator.py` to encapsulate the MaxText inference engine, handling model loading, tokenization, and the generation loop.
- Includes Pydantic models in `server_models.py` for robust, OpenAI-compliant request and response validation.
2. **Deployment and Utilities:**
- Provides `start_server.sh` to simplify launching the server from the project root.
- Adds `port_forward_xpk.sh`, a utility script to automatically find and connect to a server running on a GKE cluster via `xpk`, supporting custom namespaces.
- Isolates server-specific dependencies in `benchmarks/api_server/requirements.txt` (`uvicorn`, `fastapi`, `openai-harmony`).
3. **Comprehensive Documentation:**
- A new `README.md` in the `api_server` directory offers a complete guide covering:
- Installation and environment setup.
- Launching the server in both single-pod and multi-pod GKE environments.
- Detailed examples for interacting with the API using `curl` and the `openai` Python client.
- Step-by-step instructions for running benchmarks with `lm-evaluation-harness` and `evalchemy` for both log-likelihood and generative tasks.
babyplutokurt
force-pushed
the
api_sever_v1
branch
from
c703914 to
07a0e66
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This commit introduces a fully-featured, OpenAI-compatible RESTful API server for serving MaxText models. The server is built with FastAPI, supports multi-host inference on TPUs, and is designed for both interactive use and large-scale benchmarking.
Key features and additions:
Core Server Implementation:
maxtext_server.py, a FastAPI application that serves/v1/completionsand/v1/chat/completionsendpoints.maxtext_generator.pyto encapsulate the MaxText inference engine, handling model loading, tokenization, and the generation loop.server_models.pyfor robust, OpenAI-compliant request and response validation.Deployment and Utilities:
start_server.shto simplify launching the server from the project root.port_forward_xpk.sh, a utility script to automatically find and connect to a server running on a GKE cluster viaxpk, supporting custom namespaces.benchmarks/api_server/requirements.txt(uvicorn,fastapi,openai-harmony).Comprehensive Documentation:
README.mdin theapi_serverdirectory offers a complete guide covering: - Installation and environment setup. - Launching the server in both single-pod and multi-pod GKE environments. - Detailed examples for interacting with the API usingcurland theopenaiPython client. - Step-by-step instructions for running benchmarks withlm-evaluation-harnessandevalchemyfor both log-likelihood and generative tasks.Checklist
Before submitting this PR, please make sure (put X in square brackets):