HI all,
I want to run dataflow with PyTorch, Cuda and Debian bullseye.
Dataflow can run with docker image on Container-Optimized OS.
This docker container works fine and torch is able to use GPU:
FROM pytorch/pytorch:2.6.0-cuda12.6-cudnn9-runtime
WORKDIR /pipeline
COPY requirements.txt .
COPY *.py ./
RUN pip install --no-cache-dir --upgrade pip \
&& pip install --no-cache-dir -r requirements.txt \
&& pip check
COPY --from=apache/beam_python3.11_sdk:2.61.0 /opt/apache/beam /opt/apache/beam
ENTRYPOINT [ "/opt/apache/beam/boot" ]
When I try to install torch+cuda following the instruction here Start Locally | PyTorch, torch is installed but orch.cuda.is_available() return false. In both cases NVIDIA driver are installed separately.
FROM python:3.11-bullseye
WORKDIR /pipeline
COPY requirements.txt .
COPY *.py ./
RUN pip install --no-cache-dir --upgrade pip \
&& pip install --no-cache-dir -r requirements.txt \
&& pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu126 \
&& pip check
COPY --from=apache/beam_python3.11_sdk:2.61.0 /opt/apache/beam /opt/apache/beam
ENTRYPOINT [ "/opt/apache/beam/boot" ]
GPU looks ok:
ldd --version
ldd (Debian GLIBC 2.31-13+deb11u11) 2.31