feat(api_server): Add OpenAI-compatible API server for MaxText models #2313
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
A9isha,
NuojCheng,
RissyRan,
SujeethJinesh,
SurbhiJainUSC,
aireenmei,
bvandermoon,
gagika,
gobbleturk,
hengtaoguo,
khatwanimohit,
mitalisi,
parambole,
richjames0,
shauryagup,
shralex,
vipannalla and
yangyuwei
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
Collaborator
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
Great work! I will take 2nd round of review for files maxtet_generator, maxtext_server, and server_models.
babyplutokurt
force-pushed
the
api_sever_v1
branch
5 times, most recently
from
b70e083 to
18d3055
Compare
RissyRan
reviewed
Sep 12, 2025
Collaborator
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
Thanks! LGTM in general. Could you leverage Gemini to build some unit tests, especailly for maxtext_generator? More unit tests are very welcome!
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
Lumosis,
gpolovets1,
jrplatin,
mailvijayasingh,
michelle-yooh and
patemotter
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
Collaborator
RissyRan
left a comment
RissyRan
left a comment
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.
babyplutokurt
force-pushed
the
api_sever_v1
branch
4 times, most recently
from
c703914 to
07a0e66
Compare
RissyRan
approved these changes
Sep 24, 2025
Collaborator
RissyRan
left a comment
RissyRan
left a comment
There was a problem hiding this comment.
Discussing with @bvandermoon and @hengtaoguo, we will follow it up if any issues.
babyplutokurt
force-pushed
the
api_sever_v1
branch
from
07a0e66 to
e8bdeef
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
e8bdeef to
f8031a0
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):