mirror of
https://github.com/simstudioai/sim.git
synced 2026-09-21 13:00:04 +08:00
fix(pii): install CUDA torch on amd64 so GLiNER can run on GPU (#5552)
* fix(pii): install CUDA torch on amd64 so GLiNER can run on GPU The published pii image installed a CPU-only torch build, so GLiNER on the ECS GPU fleet died at model load with "Attempting to deserialize object on a CUDA device but torch.cuda.is_available() is False". The Dockerfile already had a TORCH_INDEX_URL arg, but no CI job ever passed --build-arg, so every image silently took the cpu default. Select the wheel index from TARGETARCH instead: amd64 gets cu128, arm64 keeps the cpu index (cu128 publishes no aarch64 wheel at 2.11.0, and no arm64 target has a GPU). CUDA torch falls back to CPU when no GPU is present, so one image still serves both the Fargate CPU tasks and the EC2 GPU tasks off the same tag — no CI or CDK changes needed. cu128 keeps sm_75, the compute capability of the fleet's T4s, and its CUDA 12.8 runtime needs driver >=525 via minor-version compatibility, which the ECS GPU AMI satisfies. cu121 was not an option: that index stops at torch 2.5.1. Verified in an amd64 build of the changed block: 2.11.0+cu128 cuda=12.8 arch=sm_75 sm_80 sm_86 sm_90 sm_100 sm_120 arm64 still resolves to 2.11.0+cpu. A build-time assert now fails the image if amd64 ever silently regresses to a cpu wheel. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QHNEWVrh7k89m8Wtqzhs18 * fix(pii): assert torch CUDA state after every pip install The check sat directly after the torch install, but requirements-gliner.txt and requirements-dev.txt are installed afterwards and resolve against PyPI with no torch pin, so a future gliner bump could swap the wheel that torch_index selected without tripping the assert. Neither file changes torch today (verified: torch is 2.11.0+cu128 both before and after the gliner install), so this guards the invariant rather than fixing a live regression. Moving it below the last pip install makes it certify the torch that actually ships. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01QHNEWVrh7k89m8Wtqzhs18 --------- Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
co-authored by
Claude Opus 4.8
parent
f0d85cb7ab
commit
97bb727eeb
+30
-6
@@ -6,10 +6,10 @@
|
||||
# gliner package, and the baked GLiNER weights all ship in it, so flipping
|
||||
# engines never requires an image swap.
|
||||
#
|
||||
# GPU variant (EC2-GPU fleet follow-up): same Dockerfile, CUDA torch wheels —
|
||||
# docker build --build-arg TORCH_INDEX_URL=https://download.pytorch.org/whl/cu128 ...
|
||||
# (torch CUDA wheels bundle their own CUDA libs; the host only needs the
|
||||
# nvidia container runtime.)
|
||||
# ONE image also serves both fleets: the amd64 build ships CUDA torch, which
|
||||
# falls back to CPU when no GPU is present, so the Fargate CPU tasks and the
|
||||
# EC2-GPU tasks pull the same tag. (torch CUDA wheels bundle their own CUDA
|
||||
# libs; the host only needs the nvidia driver + container runtime.)
|
||||
#
|
||||
# Source files are COPY'd last so code edits never re-download deps or models.
|
||||
# ========================================
|
||||
@@ -47,10 +47,25 @@ RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
# torch is pinned here (not requirements-gliner.txt) because the CPU and CUDA
|
||||
# builds install the same version from different wheel indexes. 2.11.0 is the
|
||||
# newest release published on both the cpu and cu128 indexes for py312.
|
||||
#
|
||||
# cu128's arch list keeps sm_75, the compute capability of the GPU fleet's T4s.
|
||||
# cu121 could not serve this pin anyway — that index stops at torch 2.5.1.
|
||||
# CUDA 12.8 needs an NVIDIA driver >=525 via minor-version compatibility, which
|
||||
# the ECS GPU AMI's nvidia-driver-latest-dkms satisfies.
|
||||
#
|
||||
# arm64 takes the cpu index: cu128 publishes no aarch64 wheel at 2.11.0, and no
|
||||
# arm64 target has a GPU.
|
||||
ARG TORCH_VERSION=2.11.0
|
||||
ARG TORCH_INDEX_URL=https://download.pytorch.org/whl/cpu
|
||||
ARG TORCH_CUDA_INDEX_URL=https://download.pytorch.org/whl/cu128
|
||||
ARG TORCH_CPU_INDEX_URL=https://download.pytorch.org/whl/cpu
|
||||
ARG TARGETARCH
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
pip install torch==${TORCH_VERSION} --index-url ${TORCH_INDEX_URL}
|
||||
case "${TARGETARCH}" in \
|
||||
amd64) torch_index="${TORCH_CUDA_INDEX_URL}" ;; \
|
||||
arm64) torch_index="${TORCH_CPU_INDEX_URL}" ;; \
|
||||
*) echo "unsupported TARGETARCH: ${TARGETARCH}" >&2; exit 1 ;; \
|
||||
esac && \
|
||||
pip install torch==${TORCH_VERSION} --index-url "${torch_index}"
|
||||
|
||||
COPY apps/pii/requirements-gliner.txt ./requirements-gliner.txt
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
@@ -84,6 +99,15 @@ COPY apps/pii/requirements-dev.txt ./requirements-dev.txt
|
||||
RUN --mount=type=cache,target=/root/.cache/pip \
|
||||
pip install -r requirements-dev.txt
|
||||
|
||||
# Runs after every pip install, because the requirements above resolve against
|
||||
# PyPI and could swap the wheel torch_index chose. A cpu-only torch on amd64
|
||||
# otherwise surfaces only as "torch.cuda.is_available() is False" once GLiNER
|
||||
# loads on a GPU host.
|
||||
RUN python -c "import torch; \
|
||||
have = torch.version.cuda is not None; \
|
||||
want = '${TARGETARCH}' == 'amd64'; \
|
||||
assert have == want, f'{torch.__version__}: cuda build={have}, expected={want}'"
|
||||
|
||||
RUN groupadd -g 1001 pii && \
|
||||
useradd -u 1001 -g pii pii && \
|
||||
chown -R pii:pii /app
|
||||
|
||||
Reference in New Issue
Block a user