torchsig/Dockerfile at main · TorchDSP/torchsig · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
# +----------------------------------------------------------------------------+
# |                       Stage 1: Builder (with CUDA Toolkit)                |
# +----------------------------------------------------------------------------+
FROM nvidia/cuda:12.1.0-devel-ubuntu22.04 AS builder
WORKDIR /workspace

# +----------------------------------------------------------------------------+
# | Install system build tools, Python headers, pip and Git                   |
# +----------------------------------------------------------------------------+
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
      build-essential \
      curl \
      libffi-dev libssl-dev \
      python3.10 python3-pip \
      git && \
    rm -rf /var/lib/apt/lists/*

# +----------------------------------------------------------------------------+
# | Install rustup and set up latest stable Rust (>=1.81)                     |
# +----------------------------------------------------------------------------+
ENV RUSTUP_HOME=/usr/local/rustup \
    CARGO_HOME=/usr/local/cargo \
    PATH=/usr/local/cargo/bin:$PATH

RUN curl https://sh.rustup.rs -sSf | sh -s -- -y --default-toolchain stable \
    && rustc --version

# +----------------------------------------------------------------------------+
# | Upgrade pip and install Python packaging tools (setuptools-rust, wheel)    |
# +----------------------------------------------------------------------------+
RUN python3 -m pip install --upgrade pip setuptools setuptools-rust wheel

# +----------------------------------------------------------------------------+
# | Copy the entire project into the builder image                             |
# +----------------------------------------------------------------------------+
COPY . .

# +----------------------------------------------------------------------------+
# | Install TorchSig (builds the Rust extension in-place)                     |
# +----------------------------------------------------------------------------+
RUN pip install . --no-cache-dir

# +============================================================================+
# |                Stage 2: Runtime (CUDA Runtime Only)                       |
# +============================================================================+
FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
WORKDIR /workspace

# +----------------------------------------------------------------------------+
# | Install minimal system libraries required at runtime (e.g., OpenCV deps)   |
# +----------------------------------------------------------------------------+
RUN apt-get update && \
    apt-get install -y --no-install-recommends \
      python3.10 \
      python-is-python3 \
      libgl1 \
      libsm6 \
      libxrender1 \
      libxext6 && \
    rm -rf /var/lib/apt/lists/*

# +----------------------------------------------------------------------------+
# | Copy installed Python packages from builder into runtime image             |
# +----------------------------------------------------------------------------+
COPY --from=builder /usr/local/lib/python3.10/dist-packages/ \
                     /usr/local/lib/python3.10/dist-packages/

# +----------------------------------------------------------------------------+
# | (Optional) Copy source/tests/scripts if you need them in the runtime       |
# +----------------------------------------------------------------------------+
COPY --from=builder /workspace /workspace

# +----------------------------------------------------------------------------+
# | Default to bash for interactive GPU testing                                |
# +----------------------------------------------------------------------------+
CMD ["bash"]