Merge pull request #2946 from VWS-Python/adamtheturtle/cpu-only-torch

Use CPU-only torch for Docker builds to reduce image size

Author: adamtheturtle

pyproject.toml

@@ -44,6 +44,7 @@ dependencies = [
     "responses>=0.25.3",
     "torch>=2.5.1",
     "torchmetrics>=1.5.1",
+    "torchvision>=0.20.1",
     "tzdata; sys_platform=='win32'",
     "vws-auth-tools>=2024.7.12",
     "werkzeug>=3.1.2",
@@ -138,6 +139,11 @@ fallback_version = "0.0.0"
 # Code to match this is in ``conf.py``.
 version_scheme = "post-release"
 
+[tool.uv]
+sources.torch = { index = "pytorch-cpu" }
+sources.torchvision = { index = "pytorch-cpu" }
+index = [ { name = "pytorch-cpu", url = "https://download.pytorch.org/whl/cpu", explicit = true } ]
+
 [tool.ruff]
 line-length = 79
 lint.select = [
@@ -312,6 +318,9 @@ per_rule_ignores.DEP002 = [
     # tzdata is needed on Windows for zoneinfo to work.
     # See https://docs.python.org/3/library/zoneinfo.html#data-sources.
     "tzdata",
+    # torchvision is used transitively via piq, but must be a direct dependency
+    # so that tool.uv.sources can route it to the CPU-only PyTorch index.
+    "torchvision",
 ]
 
 [tool.pyproject-fmt]

src/mock_vws/_flask_server/Dockerfile

@@ -12,12 +12,13 @@ COPY --chown=myuser:myuser . /app
 # See https://pythonspeed.com/articles/activate-virtualenv-dockerfile/
 # For why we use this method of activating the virtual environment.
 ENV VIRTUAL_ENV=/app/docker_venvs/.venv
-RUN python3 -m venv $VIRTUAL_ENV
+ENV UV_PROJECT_ENVIRONMENT=/app/docker_venvs/.venv
 ENV PATH="$VIRTUAL_ENV/bin:$PATH"
 
 WORKDIR /app
-RUN pip install --no-cache-dir uv==0.10.4 && \
-    uv pip install --no-cache-dir --upgrade --editable .
+RUN python3 -m venv $VIRTUAL_ENV && \
+    pip install --no-cache-dir uv==0.10.4 && \
+    uv sync --no-cache
 EXPOSE 5000
 ENTRYPOINT ["python"]
 HEALTHCHECK --interval=1s --timeout=10s --start-period=5s --retries=3 CMD ["python", "/app/src/mock_vws/_flask_server/healthcheck.py"]