Skip to content

Commit 6281137

Browse files
JCHAVEROTJCHAVEROT
authored
Fix torch import crash in Docker images (#322)
1 parent e3890d1 commit 6281137

9 files changed

Lines changed: 64 additions & 49 deletions

File tree

‎.dockerignore‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,2 @@
11
.venv
2-
tests
3-
test_data
2+
.git

‎.github/workflows/tests.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ jobs:
4444
- name: Run tests - process
4545
run: |
4646
source .venv/bin/activate
47-
pytest --ignore=tests/test_colvision.py
47+
pytest
4848
4949
- name: Install dependencies - colvision
5050
run: |

‎docker/arch/Dockerfile‎

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,17 +3,15 @@ ARG UV_OVERRIDE
33

44
# --- Base images ---
55
FROM archlinux:latest AS gpu
6-
ENV UV_ARGUMENTS="--extra all,cu126"
7-
RUN echo "Using GPU image"
6+
ENV SYNC_ARGS="--extra all --extra cu126"
87

98
FROM archlinux:latest AS cpu
10-
ENV UV_ARGUMENTS="--extra all,cpu"
11-
RUN echo "Using CPU-only image"
9+
ENV SYNC_ARGS="--extra all --extra cpu"
1210

1311
# --- Build ---
1412
FROM ${DEVICE:-gpu} AS build
1513
ARG UV_OVERRIDE
16-
ENV UV_ARGUMENTS=${UV_OVERRIDE:-$UV_ARGUMENTS}
14+
ENV SYNC_ARGS=${UV_OVERRIDE:-$SYNC_ARGS}
1715
RUN echo "DisableSandbox" >> /etc/pacman.conf
1816

1917
ARG USER_UID=1000
@@ -22,15 +20,22 @@ ARG USER_GID=1000
2220
# Install system dependencies
2321
RUN pacman -Syu --noconfirm && \
2422
pacman -S --noconfirm --needed \
25-
base-devel python uv \
26-
tzdata curl ffmpeg \
23+
base-devel \
24+
tzdata curl git ffmpeg \
2725
libsm libxext nss \
2826
libxi libxrandr libxcomposite libxcursor libxdamage libxfixes libxrender \
2927
alsa-lib atk gtk3 libreoffice-fresh libjpeg-turbo pango && \
3028
ln -fs /usr/share/zoneinfo/Europe/Zurich /etc/localtime && \
3129
echo "Europe/Zurich" > /etc/timezone && \
3230
pacman -Sc --noconfirm
3331

32+
# Install uv
33+
COPY --from=ghcr.io/astral-sh/uv:0.7.13 /uv /uvx /bin/
34+
35+
ENV UV_PYTHON=3.11 \
36+
UV_PROJECT_ENVIRONMENT=/app/.venv \
37+
UV_LINK_MODE=copy
38+
3439
# Create non-root user
3540
RUN groupadd --gid ${USER_GID} mmoreuser \
3641
&& useradd --uid ${USER_UID} --gid ${USER_GID} -m mmoreuser
@@ -39,14 +44,14 @@ WORKDIR /app
3944
RUN chown mmoreuser:mmoreuser /app
4045
USER mmoreuser
4146

42-
# Set up Python virtual environment
43-
ENV VIRTUAL_ENV=/app/.venv
44-
RUN uv venv --python 3.11 .venv \
45-
&& uv pip install --no-cache weasyprint
46-
4747
# Install dependencies (cached unless uv.lock changes)
4848
COPY --chown=mmoreuser:mmoreuser pyproject.toml uv.lock /app/
49-
RUN uv sync --frozen --no-install-project ${UV_ARGUMENTS}
49+
RUN --mount=type=cache,target=/home/mmoreuser/.cache/uv,uid=${USER_UID},gid=${USER_GID} \
50+
uv sync --frozen --no-install-project ${SYNC_ARGS}
51+
52+
# weasyprint needs to be installed using pip
53+
RUN --mount=type=cache,target=/home/mmoreuser/.cache/uv,uid=${USER_UID},gid=${USER_GID} \
54+
uv pip install weasyprint
5055

5156
# Install mmore from local source code
5257
COPY --chown=mmoreuser:mmoreuser . /app

‎docker/arch/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@ CPU-only:
1919
sudo docker build -f docker/arch/Dockerfile --build-arg DEVICE=cpu -t mmore:arch-cpu .
2020
```
2121

22-
Custom extras (overrides the default `--extra all,cu126` or `--extra all,cpu`):
22+
Custom extras (overrides the default `--extra all --extra cu126` or `--extra all --extra cpu`):
2323
```bash
24-
sudo docker build -f docker/arch/Dockerfile --build-arg UV_OVERRIDE="--extra all,cu126" -t mmore:arch .
24+
sudo docker build -f docker/arch/Dockerfile --build-arg UV_OVERRIDE="--extra process --extra rag --extra cpu" -t mmore:arch .
2525
```
2626

2727
Custom user UID/GID (e.g. for RCP):

‎docker/leap/Dockerfile‎

Lines changed: 18 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -3,43 +3,48 @@ ARG UV_OVERRIDE
33

44
# --- Base images ---
55
FROM opensuse/leap:15.6 AS gpu
6-
ENV UV_ARGUMENTS="--extra all,cu126"
7-
RUN echo "Using GPU image"
6+
ENV SYNC_ARGS="--extra all --extra cu126"
87

98
FROM opensuse/leap:15.6 AS cpu
10-
ENV UV_ARGUMENTS="--extra all,cpu"
11-
RUN echo "Using CPU-only image"
9+
ENV SYNC_ARGS="--extra all --extra cpu"
1210

1311
# --- Build ---
1412
FROM ${DEVICE:-gpu} AS build
1513
ARG UV_OVERRIDE
16-
ENV UV_ARGUMENTS=${UV_OVERRIDE:-$UV_ARGUMENTS}
14+
ENV SYNC_ARGS=${UV_OVERRIDE:-$SYNC_ARGS}
1715

1816
# Install system dependencies
1917
RUN zypper --non-interactive refresh && \
2018
zypper --non-interactive install -y \
21-
python311 python311-pip \
22-
nano curl ffmpeg \
19+
nano curl git ffmpeg \
2320
libSM6 libXext6 mozilla-nss \
2421
libXi6 libXrandr2 libXcomposite1 libXcursor1 libXdamage1 libXfixes3 libXrender1 \
2522
libasound2 libatk-1_0-0 gtk3 libreoffice libjpeg8-devel libpango-1_0-0 && \
2623
ln -fs /usr/share/zoneinfo/Europe/Zurich /etc/localtime && \
2724
echo "Europe/Zurich" > /etc/timezone && \
2825
zypper clean --all
2926

30-
WORKDIR /app
27+
# Install uv
28+
COPY --from=ghcr.io/astral-sh/uv:0.7.13 /uv /uvx /bin/
29+
30+
ENV UV_PYTHON=3.11 \
31+
UV_PROJECT_ENVIRONMENT=/app/.venv \
32+
UV_LINK_MODE=copy
3133

32-
# Set up Python virtual environment
33-
RUN python3.11 -m venv .venv \
34-
&& .venv/bin/pip install --no-cache-dir uv weasyprint
34+
WORKDIR /app
3535

3636
# Install dependencies (cached unless uv.lock changes)
3737
COPY pyproject.toml uv.lock /app/
38-
RUN .venv/bin/uv sync --frozen --no-install-project ${UV_ARGUMENTS}
38+
RUN --mount=type=cache,target=/root/.cache/uv \
39+
uv sync --frozen --no-install-project ${SYNC_ARGS}
40+
41+
# weasyprint needs to be installed using pip
42+
RUN --mount=type=cache,target=/root/.cache/uv \
43+
uv pip install weasyprint
3944

4045
# Install mmore from local source code
4146
COPY . /app
42-
RUN .venv/bin/uv pip install --no-cache --no-deps -e .
47+
RUN uv pip install --no-cache --no-deps -e .
4348

4449
# --- Runtime ---
4550
ENV PATH="/app/.venv/bin:$PATH"

‎docker/leap/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@ CPU-only:
1919
sudo docker build -f docker/leap/Dockerfile --build-arg DEVICE=cpu -t mmore:leap-cpu .
2020
```
2121

22-
Custom extras (overrides the default `--extra all,cu126` or `--extra all,cpu`):
22+
Custom extras (overrides the default `--extra all --extra cu126` or `--extra all --extra cpu`):
2323
```bash
24-
sudo docker build -f docker/leap/Dockerfile --build-arg UV_OVERRIDE="--extra all,cu126" -t mmore:leap .
24+
sudo docker build -f docker/leap/Dockerfile --build-arg UV_OVERRIDE="--extra process --extra rag --extra cpu" -t mmore:leap .
2525
```
2626

2727
## Run

‎docker/ubuntu/Dockerfile‎

Lines changed: 15 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -3,33 +3,37 @@ ARG UV_OVERRIDE
33

44
# --- Base images ---
55
FROM nvidia/cuda:12.6.3-base-ubuntu22.04 AS gpu
6-
ENV UV_ARGUMENTS="--extra all,cu126"
7-
RUN echo "Using GPU image"
6+
ENV SYNC_ARGS="--extra all --extra cu126"
87

98
FROM ubuntu:22.04 AS cpu
10-
ENV UV_ARGUMENTS="--extra all,cpu"
11-
RUN echo "Using CPU-only image"
9+
ENV SYNC_ARGS="--extra all --extra cpu"
1210

1311
# --- Build ---
1412
FROM ${DEVICE:-gpu} AS build
1513
ARG UV_OVERRIDE
16-
ENV UV_ARGUMENTS=${UV_OVERRIDE:-$UV_ARGUMENTS}
14+
ENV SYNC_ARGS=${UV_OVERRIDE:-$SYNC_ARGS}
1715

1816
ARG USER_UID=1000
1917
ARG USER_GID=1000
2018

2119
# Install system dependencies
2220
RUN apt-get update && \
2321
DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends \
24-
python3.11 python3.11-venv python3-pip \
25-
tzdata nano curl ffmpeg libsm6 libxext6 libnss3 \
22+
tzdata nano curl git ffmpeg libsm6 libxext6 libnss3 \
2623
libxi6 libxrandr2 libxcomposite1 libxcursor1 libxdamage1 libxfixes3 libxrender1 \
2724
libasound2 libatk1.0-0 libgtk-3-0 libreoffice libjpeg-dev libpango-1.0-0 \
2825
libpangoft2-1.0-0 weasyprint && \
2926
ln -fs /usr/share/zoneinfo/Europe/Zurich /etc/localtime && \
3027
dpkg-reconfigure --frontend noninteractive tzdata && \
3128
apt-get clean && rm -rf /var/lib/apt/lists/*
3229

30+
# Install uv
31+
COPY --from=ghcr.io/astral-sh/uv:0.7.13 /uv /uvx /bin/
32+
33+
ENV UV_PYTHON=3.11 \
34+
UV_PROJECT_ENVIRONMENT=/app/.venv \
35+
UV_LINK_MODE=copy
36+
3337
# Create non-root user
3438
RUN groupadd --gid ${USER_GID} mmoreuser \
3539
&& useradd --uid ${USER_UID} --gid ${USER_GID} -m mmoreuser
@@ -38,17 +42,14 @@ WORKDIR /app
3842
RUN chown mmoreuser:mmoreuser /app
3943
USER mmoreuser
4044

41-
# Set up Python virtual environment
42-
RUN python3.11 -m venv .venv \
43-
&& .venv/bin/pip install --no-cache-dir uv
44-
4545
# Install dependencies (cached unless uv.lock changes)
4646
COPY --chown=mmoreuser:mmoreuser pyproject.toml uv.lock /app/
47-
RUN .venv/bin/uv sync --frozen --no-install-project ${UV_ARGUMENTS}
47+
RUN --mount=type=cache,target=/home/mmoreuser/.cache/uv,uid=${USER_UID},gid=${USER_GID} \
48+
uv sync --frozen --no-install-project ${SYNC_ARGS}
4849

4950
# Install mmore from local source code
5051
COPY --chown=mmoreuser:mmoreuser . /app
51-
RUN .venv/bin/uv pip install --no-cache --no-deps -e .
52+
RUN uv pip install --no-cache --no-deps -e .
5253

5354
# --- Runtime ---
5455
ENV PATH="/app/.venv/bin:$PATH"
@@ -57,4 +58,4 @@ ENV HF_HOME="/home/mmoreuser/.cache/huggingface"
5758
ENV TORCH_HOME="/home/mmoreuser/.cache/torch"
5859
ENV XDG_CACHE_HOME="/home/mmoreuser/.cache"
5960

60-
ENTRYPOINT ["/bin/bash"]
61+
ENTRYPOINT ["/bin/bash"]

‎docker/ubuntu/README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,9 +26,9 @@ CPU-only:
2626
sudo docker build -f docker/ubuntu/Dockerfile --build-arg DEVICE=cpu -t mmore:cpu .
2727
```
2828

29-
Custom extras (overrides the default `--extra all,cu126` or `--extra all,cpu`):
29+
Custom extras (overrides the default `--extra all --extra cu126` or `--extra all --extra cpu`):
3030
```bash
31-
sudo docker build -f docker/ubuntu/Dockerfile --build-arg UV_OVERRIDE="--extra all,cu126" -t mmore .
31+
sudo docker build -f docker/ubuntu/Dockerfile --build-arg UV_OVERRIDE="--extra process --extra rag --extra cpu" -t mmore .
3232
```
3333

3434
Custom user UID/GID (e.g. for RCP):

‎tests/conftest.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import importlib.util
12
import json
23
from typing import Dict, List
34

@@ -6,6 +7,10 @@
67

78
from mmore.type import DocumentMetadata, MultimodalSample
89

10+
# colvision is a conflicting extra installed in its own venv (see tests.yml)
11+
if importlib.util.find_spec("colpali_engine") is None:
12+
collect_ignore = ["test_colvision.py"]
13+
914

1015
class FakeSparseEmbedding(BaseSparseEmbedding):
1116
"""Fake sparse embedder for test purposes."""

0 commit comments

Comments
 (0)