art-dag/mono
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add JAX CPU to L1 worker image, CUDA JAX to GPU image
All checks were successful
Build and Deploy / build-and-deploy (push) Successful in 2m49s
Details

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
giles
2026-02-25 15:24:59 +00:00
parent 84e3ff3a91
commit 3bffb97ca1
2 changed files with 3 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
l1/Dockerfile.gpu
View File

@@ -63,8 +63,8 @@ RUN python3 -m pip install --upgrade pip
COPY l1/requirements.txt . COPY l1/requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt RUN pip install --no-cache-dir -r requirements.txt
# Install GPU-specific dependencies (CuPy for CUDA 12.x) # Install GPU-specific dependencies (CuPy for CUDA 12.x, JAX with CUDA)
RUN pip install --no-cache-dir cupy-cuda12x RUN pip install --no-cache-dir cupy-cuda12x jax[cuda12]
# Install PyNvVideoCodec for zero-copy GPU encoding # Install PyNvVideoCodec for zero-copy GPU encoding
RUN pip install --no-cache-dir PyNvVideoCodec RUN pip install --no-cache-dir PyNvVideoCodec

1
l1/requirements.txt
View File

@@ -13,6 +13,7 @@ markdown>=3.5.0
# Common effect dependencies (used by uploaded effects) # Common effect dependencies (used by uploaded effects)
numpy>=1.24.0 numpy>=1.24.0
opencv-python-headless>=4.8.0 opencv-python-headless>=4.8.0
jax[cpu]>=0.4.20
# core (artdag) and common (artdag_common) installed from local dirs in Dockerfile # core (artdag) and common (artdag_common) installed from local dirs in Dockerfile
psycopg2-binary psycopg2-binary
nest_asyncio nest_asyncio