From a4e3f269d8810c0d68304825ba41a065a3e0470a Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 26 Aug 2026 19:24:16 +0000 Subject: [PATCH 1/3] Stop the openpi subprocess from taking the whole GPU 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 --- .github/workflows/unit-test.yaml | 20 +++++++++++++ docker/CONTEXTS.md | 6 ++++ positronic/vendors/openpi/README.md | 20 +++++++++++++ positronic/vendors/openpi/server.py | 6 +++- .../vendors/openpi/tests/test_server.py | 28 +++++++++++++++++++ 5 files changed, 79 insertions(+), 1 deletion(-) create mode 100644 positronic/vendors/openpi/tests/test_server.py diff --git a/.github/workflows/unit-test.yaml b/.github/workflows/unit-test.yaml index c5389b8af..d7c510732 100644 --- a/.github/workflows/unit-test.yaml +++ b/.github/workflows/unit-test.yaml @@ -92,6 +92,26 @@ jobs: PYTHONPATH: ${{ env.PYTHONPATH }}:$PWD run: uv run pytest positronic/vendors/lerobot/tests + openpi: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Enable caching + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + + - name: Set up Python + run: uv venv --python 3.13 + + - name: Install dependencies + run: uv sync --locked --extra openpi + + - name: Run openpi vendor tests + env: + PYTHONPATH: ${{ env.PYTHONPATH }}:$PWD + run: uv run pytest positronic/vendors/openpi/tests + lockfile-portability: strategy: fail-fast: false diff --git a/docker/CONTEXTS.md b/docker/CONTEXTS.md index 84e2d59ba..6b4ad2048 100644 --- a/docker/CONTEXTS.md +++ b/docker/CONTEXTS.md @@ -31,6 +31,12 @@ When running `docker --context compose run ...`, volume paths in `docke CACHE_ROOT=/home/ docker --context vm-train compose run -d --service-ports openpi-server ... ``` +## Restart policy + +Start a container with `--restart on-failure:2`, not `--restart unless-stopped`. A container that crashes and +restarts reports `Up 1 second` to every `docker ps`, which reads as a slow start. A restart policy that hides a +crash costs more than the crash. + ## VM management Start: `../internal/scripts/start.sh train` diff --git a/positronic/vendors/openpi/README.md b/positronic/vendors/openpi/README.md index a5e6ad363..ed2a50c65 100644 --- a/positronic/vendors/openpi/README.md +++ b/positronic/vendors/openpi/README.md @@ -142,6 +142,26 @@ emits absolute `JointPosition` chunks executed at RoboLab's leaderboard cadence - `--recording_dir`: (Optional) Directory for server-side `.rrd` recordings (local or S3) - `--idle_timeout_min`: (Optional) Shut down after this many minutes without activity +### Serving More Than One Policy On One GPU + +The server starts its OpenPI subprocess with `XLA_PYTHON_CLIENT_PREALLOCATE=false`, so JAX allocates on +demand. JAX otherwise takes ~75% of the device at its first use, and a second server on that GPU then fails +with `RESOURCE_EXHAUSTED` while `nvidia-smi` reports the device almost free. + +With no preallocation, `XLA_PYTHON_CLIENT_MEM_FRACTION` is a hard cap on what one server allocates. Set it +per container when you co-host N policies. Leave it unset for one policy: a cap that is too low makes a large +model fail with the same `RESOURCE_EXHAUSTED`. Three policies held 30.4 GB together on an 80 GB H100 with +`XLA_PYTHON_CLIENT_MEM_FRACTION=.25`. + +The `openpi-server-8001` service is a second server on the same machine, on host port 8001: + +```bash +docker compose run --rm --service-ports -e XLA_PYTHON_CLIENT_MEM_FRACTION=.25 \ + -v ~/checkpoints:/checkpoints openpi-server-8001 ee \ + --pipeline.source.checkpoints_dir=/checkpoints/openpi/pi05_positronic_lowmem/experiment_v2/ \ + --pipeline.ee_frame=None +``` + ### API Endpoints The server exposes the following endpoints: diff --git a/positronic/vendors/openpi/server.py b/positronic/vendors/openpi/server.py index 9b72b2193..084a610e7 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -73,8 +73,12 @@ def start(self, on_progress: Callable[[str], None] | None = None): """Start the subprocess and block until it accepts connections, reporting progress.""" command = self._build_command() logger.info(f'Starting OpenPI subprocess: {" ".join(command)}') + env = os.environ.copy() + # JAX takes ~75% of the GPU at its first use, so a second server on that GPU finds none free. + # With no preallocation ``XLA_PYTHON_CLIENT_MEM_FRACTION`` caps each server. Set it per container. + env.setdefault('XLA_PYTHON_CLIENT_PREALLOCATE', 'false') # Don't pipeline stdout/stderr so we can see the output - self.process = subprocess.Popen(command, env=os.environ.copy(), cwd=str(self.openpi_root)) + self.process = subprocess.Popen(command, env=env, cwd=str(self.openpi_root)) self._wait_for_ready(on_progress) def _check_ready(self) -> bool: diff --git a/positronic/vendors/openpi/tests/test_server.py b/positronic/vendors/openpi/tests/test_server.py new file mode 100644 index 000000000..adb005ca9 --- /dev/null +++ b/positronic/vendors/openpi/tests/test_server.py @@ -0,0 +1,28 @@ +from unittest import mock + +import pytest + +pytest.importorskip('openpi_client') + +from positronic.vendors.openpi.server import OpenpiSubprocess # noqa: E402 + + +def _subprocess_env(monkeypatch) -> dict[str, str]: + """The environment ``start`` gives the openpi subprocess.""" + monkeypatch.setattr(OpenpiSubprocess, '_wait_for_ready', lambda self, on_progress: None) + with mock.patch('subprocess.Popen') as popen: + OpenpiSubprocess(checkpoint_dir='/checkpoints/exp/1000', config_name='pi05_positronic_lowmem').start() + return popen.call_args.kwargs['env'] + + +def test_the_subprocess_does_not_preallocate_the_gpu(monkeypatch): + """JAX takes ~75% of the device at its first use, which leaves a second server on that GPU none.""" + monkeypatch.delenv('XLA_PYTHON_CLIENT_PREALLOCATE', raising=False) + + assert _subprocess_env(monkeypatch)['XLA_PYTHON_CLIENT_PREALLOCATE'] == 'false' + + +def test_the_preallocation_setting_of_the_environment_reaches_the_subprocess(monkeypatch): + monkeypatch.setenv('XLA_PYTHON_CLIENT_PREALLOCATE', 'true') + + assert _subprocess_env(monkeypatch)['XLA_PYTHON_CLIENT_PREALLOCATE'] == 'true' From c9cc4d1e2b0d14456f32dacdab0b451a24ddc423 Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 26 Aug 2026 19:30:08 +0000 Subject: [PATCH 2/3] Name the preallocation variable once `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 --- positronic/vendors/openpi/server.py | 5 +++-- positronic/vendors/openpi/tests/test_server.py | 10 +++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/positronic/vendors/openpi/server.py b/positronic/vendors/openpi/server.py index 084a610e7..0f3372f72 100644 --- a/positronic/vendors/openpi/server.py +++ b/positronic/vendors/openpi/server.py @@ -24,6 +24,8 @@ logger = logging.getLogger(__name__) +PREALLOCATE_ENV = 'XLA_PYTHON_CLIENT_PREALLOCATE' + ########################################################################################### # Subprocess manager for OpenPI WebSocket server @@ -75,8 +77,7 @@ def start(self, on_progress: Callable[[str], None] | None = None): logger.info(f'Starting OpenPI subprocess: {" ".join(command)}') env = os.environ.copy() # JAX takes ~75% of the GPU at its first use, so a second server on that GPU finds none free. - # With no preallocation ``XLA_PYTHON_CLIENT_MEM_FRACTION`` caps each server. Set it per container. - env.setdefault('XLA_PYTHON_CLIENT_PREALLOCATE', 'false') + env.setdefault(PREALLOCATE_ENV, 'false') # Don't pipeline stdout/stderr so we can see the output self.process = subprocess.Popen(command, env=env, cwd=str(self.openpi_root)) self._wait_for_ready(on_progress) diff --git a/positronic/vendors/openpi/tests/test_server.py b/positronic/vendors/openpi/tests/test_server.py index adb005ca9..d1f04ed4b 100644 --- a/positronic/vendors/openpi/tests/test_server.py +++ b/positronic/vendors/openpi/tests/test_server.py @@ -4,7 +4,7 @@ pytest.importorskip('openpi_client') -from positronic.vendors.openpi.server import OpenpiSubprocess # noqa: E402 +from positronic.vendors.openpi.server import PREALLOCATE_ENV, OpenpiSubprocess # noqa: E402 def _subprocess_env(monkeypatch) -> dict[str, str]: @@ -17,12 +17,12 @@ def _subprocess_env(monkeypatch) -> dict[str, str]: def test_the_subprocess_does_not_preallocate_the_gpu(monkeypatch): """JAX takes ~75% of the device at its first use, which leaves a second server on that GPU none.""" - monkeypatch.delenv('XLA_PYTHON_CLIENT_PREALLOCATE', raising=False) + monkeypatch.delenv(PREALLOCATE_ENV, raising=False) - assert _subprocess_env(monkeypatch)['XLA_PYTHON_CLIENT_PREALLOCATE'] == 'false' + assert _subprocess_env(monkeypatch)[PREALLOCATE_ENV] == 'false' def test_the_preallocation_setting_of_the_environment_reaches_the_subprocess(monkeypatch): - monkeypatch.setenv('XLA_PYTHON_CLIENT_PREALLOCATE', 'true') + monkeypatch.setenv(PREALLOCATE_ENV, 'true') - assert _subprocess_env(monkeypatch)['XLA_PYTHON_CLIENT_PREALLOCATE'] == 'true' + assert _subprocess_env(monkeypatch)[PREALLOCATE_ENV] == 'true' From 5955a0e8654a2c9d7a470af9a2db9a747ecd92ee Mon Sep 17 00:00:00 2001 From: Vladimir Yakunin Date: Wed, 26 Aug 2026 19:45:11 +0000 Subject: [PATCH 3/3] Say what `on-failure:2` does, in place of the aphorism 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 --- docker/CONTEXTS.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docker/CONTEXTS.md b/docker/CONTEXTS.md index 6b4ad2048..1de1c69d2 100644 --- a/docker/CONTEXTS.md +++ b/docker/CONTEXTS.md @@ -34,8 +34,8 @@ CACHE_ROOT=/home/ docker --context vm-train compose run -d --service-ports ## Restart policy Start a container with `--restart on-failure:2`, not `--restart unless-stopped`. A container that crashes and -restarts reports `Up 1 second` to every `docker ps`, which reads as a slow start. A restart policy that hides a -crash costs more than the crash. +restarts reports `Up 1 second` to every `docker ps`, which reads as a slow start. `on-failure:2` stops the +container after the second crash, so `docker ps` shows it dead and `docker logs` gives the fault. ## VM management