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:
Theodore Li
2026-07-10 16:21:45 -04:00
committed by GitHub
co-authored by Claude Opus 4.8
parent f0d85cb7ab
commit 97bb727eeb
+30 -6
View File
@@ -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