coop/rose-ash
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

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 .
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