Change installation of pytorch from CUDA-based torch to CPU-based torch in docker-build

[anates]

I have a poetry-based project I'd like to pack into a docker image. For that, I use the following docker script:

FROM python:3.12 AS builder

ARG CUR_GIT_COMMIT
ENV CUR_GIT_COMMIT=$CUR_GIT_COMMIT

RUN pip install poetry==1.8.4
RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 libegl-dev -y

ENV POETRY_NO_INTERACTION=1 \
    POETRY_VIRTUALENVS_IN_PROJECT=1 \
    POETRY_VIRTUALENVS_CREATE=1 \
    POETRY_CACHE_DIR=/tmp/poetry_cache

WORKDIR /test_project

COPY pyproject.toml poetry.lock ./
RUN touch README.md

# RUN poetry install --without dev --no-root
# RUN poetry install --no-root
RUN poetry install --without dev --no-root && rm -rf ${POETRY_CACHE_DIR}

FROM python:3.12-slim AS runtime 

ARG CUR_GIT_COMMIT
ENV CUR_GIT_COMMIT=$CUR_GIT_COMMIT

ENV VIRTUAL_ENV=/test_project/.venv \
    PATH="/test_project/.venv/bin:${PATH}"

WORKDIR /test_project

RUN apt-get update && apt-get install ffmpeg libsm6 libxext6 libegl-dev -y

COPY --from=builder ${VIRTUAL_ENV} ${VIRTUAL_ENV}

COPY test_project ./test_project
COPY run_test_project.py ./


ENTRYPOINT ["python", "-m", "run_test_project"]

However, as my project contains pytorch with CUDA, the resulting image has a size of > 8 GByte, much larger than what is supported by my repository. Therefore, I want to switch out my CUDA-based torch with a CPU-based torch build, but only in the docker-image. As my current implementation of torch in pyproject.toml is defined as

[tool.poetry.dependencies]
python = ">=3.10,<3.13"
torch = { version = "^2.5.1+cu124", source = "pytorch" }
torchaudio = { version = "^2.1.1", source = "pytorch" }
torchvision = { version = "^0.20.0+cu124", source = "pytorch" }

[[tool.poetry.source]]
name = "pytorch"
url = "https://download.pytorch.org/whl/cu124"
priority = "explicit"

the idea was to replace the package names and -url with the cpu-based url before creating the docker image, using

RUN sed -i -e "s%whl/cu124%whl/cpu%g" ./pyproject.toml
RUN sed -i -e "s%+cu124%%g" ./pyproject.toml
RUN poetry lock --no-update

However, this then results in a stuck poetry-loop while building the image. Is this the right way of switching, or is there another way I could solve my issue of the image being too large?

[radoering]

I do not see immediately see an issue. Locking may just take a while.

You may find some useful information in #6409.

In Poetry 2.0, you will have the following option: https://python-poetry.org/docs/main/dependency-specification/#extra-environment-marker

[anates]

@radoering : After waiting for a longer time, it locked, but the additional nvidia-packages necessary for the gpu-version are still installed (and therefore blow up the image significantly). The additional environmental markers might be an option, but have to be tested after poetry 2.0 has been released.

[radoering]

This is because https://download.pytorch.org/whl/cpu also provides the CUDA wheels from https://download.pytorch.org/whl/cu124. Actually, it looks like both point to the same list of wheels...

Further, 2.5.1+cu124 is a perfectly valid solution for ^2.5.1 so that running poetry lock --update after stripping cu124 from the constraints, keeps the cu124 versions and there is nothing wrong about that.

Unfortunately, I do not know a way to force 2.5.1 if the index also provides a local version like 2.5.1+cu124.