From 3bffb97ca18fb716875026bea14e4b7271b85af1 Mon Sep 17 00:00:00 2001 From: giles Date: Wed, 25 Feb 2026 15:24:59 +0000 Subject: [PATCH] Add JAX CPU to L1 worker image, CUDA JAX to GPU image CPU workers can now run GPU-queue rendering tasks via JAX on CPU. GPU image overrides with jax[cuda12] for full CUDA acceleration. Co-Authored-By: Claude Opus 4.6 --- l1/Dockerfile.gpu | 4 ++-- l1/requirements.txt | 1 + 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/l1/Dockerfile.gpu b/l1/Dockerfile.gpu index 8f764da..ad33eb8 100644 --- a/l1/Dockerfile.gpu +++ b/l1/Dockerfile.gpu @@ -63,8 +63,8 @@ RUN python3 -m pip install --upgrade pip COPY l1/requirements.txt . RUN pip install --no-cache-dir -r requirements.txt -# Install GPU-specific dependencies (CuPy for CUDA 12.x) -RUN pip install --no-cache-dir cupy-cuda12x +# Install GPU-specific dependencies (CuPy for CUDA 12.x, JAX with CUDA) +RUN pip install --no-cache-dir cupy-cuda12x jax[cuda12] # Install PyNvVideoCodec for zero-copy GPU encoding RUN pip install --no-cache-dir PyNvVideoCodec diff --git a/l1/requirements.txt b/l1/requirements.txt index be6950f..8dd85aa 100644 --- a/l1/requirements.txt +++ b/l1/requirements.txt @@ -13,6 +13,7 @@ markdown>=3.5.0 # Common effect dependencies (used by uploaded effects) numpy>=1.24.0 opencv-python-headless>=4.8.0 +jax[cpu]>=0.4.20 # core (artdag) and common (artdag_common) installed from local dirs in Dockerfile psycopg2-binary nest_asyncio