Stop the openpi subprocess from taking the whole GPU

.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

docker/CONTEXTS.md

@@ -31,6 +31,12 @@ When running `docker --context <remote> compose run ...`, volume paths in `docke
 CACHE_ROOT=/home/<user> 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. `on-failure:2` stops the
+container after the second crash, so `docker ps` shows it dead and `docker logs` gives the fault.
+
 ## VM management
 
 Start: `../internal/scripts/start.sh train`

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:

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
@@ -73,8 +75,11 @@ 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.
+        env.setdefault(PREALLOCATE_ENV, '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:

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 PREALLOCATE_ENV, 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(PREALLOCATE_ENV, raising=False)
+
+    assert _subprocess_env(monkeypatch)[PREALLOCATE_ENV] == 'false'
+
+
+def test_the_preallocation_setting_of_the_environment_reaches_the_subprocess(monkeypatch):
+    monkeypatch.setenv(PREALLOCATE_ENV, 'true')
+
+    assert _subprocess_env(monkeypatch)[PREALLOCATE_ENV] == 'true'