From 40f15463210f868b7858847f483816426c78ce29 Mon Sep 17 00:00:00 2001 From: Orbax Authors Date: Sat, 4 Jul 2026 15:45:11 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 942603774 --- .../_src/testing/benchmarks/xpk/Dockerfile.orbax_benchmark | 2 +- .../checkpoint/_src/testing/benchmarks/xpk/Dockerfile.sidecar | 3 +-- 2 files changed, 2 insertions(+), 3 deletions(-) diff --git a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/Dockerfile.orbax_benchmark b/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/Dockerfile.orbax_benchmark index 2c583b0ac..461878506 100644 --- a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/Dockerfile.orbax_benchmark +++ b/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/Dockerfile.orbax_benchmark @@ -108,7 +108,7 @@ RUN if [ "$JAX_VERSION" = "newest" ]; then \ # 4. Install Orbax from Source WORKDIR /app/orbax_repo/checkpoint RUN pip install --no-cache-dir . -RUN pip install pathwaysutils tensorboard cloudpickle +RUN pip install "pathwaysutils>=0.1.10" tensorboard cloudpickle # 5. Environment Setup # Set PYTHONPATH so 'import orbax' works from the correctly mapped directory diff --git a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/Dockerfile.sidecar b/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/Dockerfile.sidecar index 053deb7b1..cce9bdf61 100644 --- a/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/Dockerfile.sidecar +++ b/checkpoint/orbax/checkpoint/_src/testing/benchmarks/xpk/Dockerfile.sidecar @@ -55,11 +55,10 @@ RUN if [ -f "requirements.txt" ]; then pip install --no-cache-dir -r requirement # 4. Install Orbax from Source WORKDIR /app/orbax_repo/checkpoint RUN pip install --no-cache-dir . -RUN pip install pathwaysutils tensorboard cloudpickle gcsfs +RUN pip install "pathwaysutils>=0.1.10" tensorboard cloudpickle gcsfs # Ensure Orbax are in PYTHONPATH ENV PYTHONPATH=/app/orbax/checkpoint # Verify installation RUN python3 -c "import orbax.checkpoint; print('Orbax installed:', orbax.checkpoint.__file__)" -