Stop the openpi subprocess from taking the whole GPU - #681
Conversation
JAX preallocates ~75% of the device at its first use. The `OpenpiSubprocess` launched the openpi server with the inherited environment, so the first server on a GPU claimed it all and a second one failed with `RESOURCE_EXHAUSTED` while `nvidia-smi` reported 84% of the device free. The shipped `openpi-server-8001` service exists to run that second server, so the configuration could not work. `start` now sets `XLA_PYTHON_CLIENT_PREALLOCATE=false` on the subprocess environment, and leaves an operator-set value alone. It sets no default for `XLA_PYTHON_CLIENT_MEM_FRACTION`: with no preallocation that value is a hard cap, and a default would make a large single-tenant model fail. The README says how to set it per container, with the figures that held three policies in 30.4 GB on an 80 GB H100. The vendor tests now also run under the `openpi` extra, which is what reaches the new test. Ticket: Positronic-Robotics/internal#774 #open
`hardcoded-keys`: the server and its test spelled `XLA_PYTHON_CLIENT_PREALLOCATE` five times between them. `PREALLOCATE_ENV` holds it, and the test imports it. `diff-comments`: the second comment line told the reader to set `XLA_PYTHON_CLIENT_MEM_FRACTION`, which this code neither reads nor writes. The README says it, where an operator reads it. Ticket: Positronic-Robotics/internal#774 #open
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: c9cc4d1e2b
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| pytest.importorskip('openpi_client') | ||
|
|
||
| from positronic.vendors.openpi.server import PREALLOCATE_ENV, OpenpiSubprocess # noqa: E402 |
There was a problem hiding this comment.
Rule grandfathered-violation violated:
The newly added # noqa: E402 suppresses a diagnostic in a new file without explaining the optional-dependency import constraint; restructure the import gate to avoid the suppression, or retain the narrow suppression with a reason at this site.
AGENTS.md reference: AGENTS.md:L14-L22
Useful? React with 👍 / 👎.
The restart-policy note ended on a closing aphorism, which the writing rules ban. It now says what the flag does for the reader who follows it. Ticket: Positronic-Robotics/internal#774 #open
1 participant
JAX preallocates ~75% of the GPU at its first use.
OpenpiSubprocess.startlaunched theopenpi server with the inherited environment, so the first server on a GPU claimed the device
and a second one died with
RESOURCE_EXHAUSTED: Out of memory while trying to allocate 2415919104 byteswhilenvidia-smiread 12.7 GB used of 81.5 GB — 84% of the device free.docker/docker-compose.ymlships anopenpi-server-8001service whose only purpose is asecond openpi server on one box, so this was a shipped configuration that could not work.
Ticket: internal#774 — openpi server preallocates the whole GPU
What this ships
positronic/vendors/openpi/server.py—startsetsXLA_PYTHON_CLIENT_PREALLOCATE=falseon the subprocess environment with
setdefault, so an operator-set value survives.PREALLOCATE_ENVholds the name, and the test imports it.positronic/vendors/openpi/tests/test_server.py— the subprocess getsfalseby default,and an environment that already names the variable reaches it unchanged. The first fails
without the fix; the second fails if
setdefaultbecomes an assignment.positronic/vendors/openpi/README.md— a section under 4. Serve Inference on servingmore than one policy on one GPU, with the
openpi-server-8001invocation.docker/CONTEXTS.md— a restart policy note:--restart unless-stoppedreports acrash-looping container as
Up 1 secondto everydocker ps..github/workflows/unit-test.yaml— anopenpijob, mirroringlerobot. The new testneeds
openpi_client, which no job installed, sopytest.importorskipwould otherwisehave skipped it everywhere.
No default for
XLA_PYTHON_CLIENT_MEM_FRACTIONWith no preallocation that variable is a hard cap on total allocation, so a default would
fail a large model in the ordinary single-tenant case. It is a deployment knob: the README
gives the per-container form and the measured figures —
.25held three policies in 30.4 GBon an 80 GB H100.
positronic/vendors/openpi/train.pykeeps its ownXLA_PYTHON_CLIENT_MEM_FRACTION=0.995.That is single-tenant training and wants the device.