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..1de1c69d2 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. `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` 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..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 @@ -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: diff --git a/positronic/vendors/openpi/tests/test_server.py b/positronic/vendors/openpi/tests/test_server.py new file mode 100644 index 000000000..d1f04ed4b --- /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 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'