Skip to content

Commit fa79c8c

Browse files
mattkjames7mattkjames7
authored
Add CUDA build + packaging workflow + embeddings QM (#672)
1 parent c365f70 commit fa79c8c

File tree

8 files changed

+482
-24
lines changed

8 files changed

+482
-24
lines changed

‎.github/workflows/package_mage.yaml‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -27,9 +27,9 @@ on:
2727
description: "Memgraph package download link. Leave empty to use the official download link."
2828
default: ""
2929
type: string
30-
push_to_dockerhub:
30+
cuda:
3131
type: boolean
32-
description: "Push the image to DockerHub?"
32+
description: "Build the image for CUDA"
3333
default: false
3434
push_to_s3:
3535
type: boolean
@@ -60,7 +60,7 @@ jobs:
6060
memgraph_download_link: ${{ github.event.inputs.memgraph_download_link || '' }}
6161
shorten_tag: 'true'
6262
force_release: 'true'
63-
push_to_dockerhub: ${{ github.event.inputs.push_to_dockerhub || 'false' }}
63+
cuda: ${{ inputs.cuda }}
6464
push_to_s3: ${{ github.event.inputs.push_to_s3 || 'false' }}
6565
malloc: ${{ github.event.inputs.malloc == 'true' }}
6666
run_smoke_tests: ${{ inputs.run_smoke_tests }}
@@ -89,7 +89,7 @@ jobs:
8989
memgraph_download_link: ${{ github.event.inputs.memgraph_download_link || '' }}
9090
shorten_tag: 'true'
9191
force_release: 'true'
92-
push_to_dockerhub: ${{ github.event.inputs.push_to_dockerhub || 'false' }}
92+
cuda: ${{ inputs.cuda }}
9393
push_to_s3: ${{ github.event.inputs.push_to_s3 || 'false' }}
9494
malloc: ${{ matrix.malloc }}
9595
run_smoke_tests: ${{ inputs.run_smoke_tests }}

‎.github/workflows/reusable_package_mage.yaml‎

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,10 @@ on:
6666
type: boolean
6767
description: "Print the output URL at the end of the workflow"
6868
default: false
69+
cuda:
70+
type: boolean
71+
description: "Build the image for CUDA"
72+
default: false
6973

7074
env:
7175
OS: "${{ inputs.arch == 'arm64' && 'ubuntu-24.04-aarch64' || 'ubuntu-24.04' }}${{ inputs.malloc && '-malloc' || '' }}"
@@ -200,7 +204,11 @@ jobs:
200204
201205
- name: Set target dockerfile
202206
run: |
203-
DOCKERFILE="Dockerfile.release"
207+
if [[ "${{ inputs.cuda }}" == true ]]; then
208+
DOCKERFILE="Dockerfile.cuda"
209+
else
210+
DOCKERFILE="Dockerfile.release"
211+
fi
204212
echo "DOCKERFILE=${DOCKERFILE}" >> $GITHUB_ENV
205213
206214
- name: Check if specified version tag is already pushed

‎Dockerfile.cuda‎

Lines changed: 249 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,249 @@
1+
ARG PY_VERSION_DEFAULT=3.12
2+
ARG MGBUILD_IMAGE
3+
4+
FROM $MGBUILD_IMAGE as builder
5+
6+
USER root
7+
8+
ARG TARGETARCH
9+
ARG PY_VERSION_DEFAULT
10+
ARG BUILD_TYPE
11+
ARG CACHE_PRESENT
12+
ENV BUILD_TYPE=${BUILD_TYPE}
13+
ENV PY_VERSION ${PY_VERSION_DEFAULT}
14+
ARG CUSTOM_MIRROR
15+
16+
17+
# This modifies the apt configuration to rety 3 times
18+
RUN echo 'Acquire::Retries "3";' > /etc/apt/apt.conf.d/80-retries
19+
20+
# If CUSTOM_MIRROR is set, replace the default archive.ubuntu.com
21+
# and security.ubuntu.com URIs in your .sources file
22+
RUN if [ -n "$CUSTOM_MIRROR" ]; then \
23+
sed -E -i \
24+
-e '/^URIs:/ s#https?://[^ ]*archive\.ubuntu\.com#'"$CUSTOM_MIRROR"'#g' \
25+
-e '/^URIs:/ s#https?://security\.ubuntu\.com#'"$CUSTOM_MIRROR"'#g' \
26+
/etc/apt/sources.list.d/ubuntu.sources; \
27+
fi
28+
29+
# Essentials for production/dev
30+
RUN apt-get update && apt-get install -y \
31+
libcurl4 `memgraph` \
32+
libpython${PY_VERSION} `memgraph` \
33+
libssl-dev `memgraph` \
34+
openssl `memgraph` \
35+
build-essential `mage-memgraph` \
36+
cmake `mage-memgraph` \
37+
curl `mage-memgraph` \
38+
g++ `mage-memgraph` \
39+
python3 `mage-memgraph` \
40+
python3-pip `mage-memgraph` \
41+
python3-setuptools `mage-memgraph` \
42+
python3-dev `mage-memgraph` \
43+
clang `mage-memgraph` \
44+
git `mage-memgraph` \
45+
unixodbc-dev `mage-memgraph` \
46+
libboost-all-dev `mage-memgraph` \
47+
uuid-dev \
48+
gdb \
49+
procps \
50+
libc6-dbg \
51+
libxmlsec1-dev xmlsec1 \
52+
libatomic1 \
53+
--no-install-recommends \
54+
&& ln -s /usr/bin/$(ls /usr/bin | grep perf) /usr/bin/perf \
55+
&& rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
56+
57+
COPY memgraph-${TARGETARCH}.deb .
58+
59+
RUN dpkg -i memgraph-${TARGETARCH}.deb && rm memgraph-${TARGETARCH}.deb
60+
61+
# move memgraph HOME so that mounting /var/lib/memgraph as a volume doesn't break Python
62+
RUN mkdir -pv /home/memgraph && \
63+
usermod -d /home/memgraph memgraph && \
64+
chown -R memgraph:memgraph /home/memgraph
65+
66+
ENV LD_LIBRARY_PATH /usr/lib/memgraph/query_modules
67+
68+
69+
WORKDIR /mage
70+
COPY . /mage
71+
72+
#hacks so we can do everything under the `memgraph` user
73+
RUN chown -R memgraph: /mage && \
74+
mkdir -pv /usr/lib/memgraph/query_modules && \
75+
chown -R memgraph: /usr/lib/memgraph/query_modules
76+
USER memgraph
77+
78+
# First install requirements
79+
RUN python3 -m pip install --no-cache-dir -r /mage/python/requirements-gpu.txt --break-system-packages && \
80+
python3 -m pip install --no-cache-dir -r /mage/python/tests/requirements.txt --break-system-packages && \
81+
python3 -m pip install --no-cache-dir -r /usr/lib/memgraph/auth_module/requirements.txt --break-system-packages && \
82+
python3 -m pip install --no-cache-dir torch-sparse torch-cluster torch-spline-conv torch-geometric torch-scatter --break-system-packages -f https://data.pyg.org/whl/torch-2.6.0+cu126.html; \
83+
python3 -m pip install --no-cache-dir dgl==2.5.0 -f https://data.dgl.ai/wheels/torch-2.6/repo.html --break-system-packages; \
84+
rm -fr /home/memgraph/.cache/pip
85+
86+
# Build query modules
87+
SHELL ["/bin/bash", "-c"]
88+
RUN source /opt/toolchain-v6/activate && curl https://sh.rustup.rs -sSf | sh -s -- -y \
89+
&& export PATH="$HOME/.cargo/bin:${PATH}" \
90+
&& rustup toolchain install 1.85 \
91+
&& rustup default 1.85 \
92+
&& python3 /mage/setup build -p /usr/lib/memgraph/query_modules/ --cpp-build-flags CMAKE_BUILD_TYPE=${BUILD_TYPE} \
93+
&& chown -R memgraph: /mage
94+
95+
96+
# Memgraph listens for Bolt Protocol on this port by default.
97+
EXPOSE 7687
98+
99+
100+
101+
FROM ubuntu:24.04 as base
102+
103+
USER root
104+
105+
ARG TARGETARCH
106+
ARG PY_VERSION_DEFAULT=3.12
107+
ARG MAGE_COMMIT
108+
ARG BUILD_TYPE
109+
ENV BUILD_TYPE=${BUILD_TYPE}
110+
ENV PY_VERSION ${PY_VERSION_DEFAULT}
111+
ENV MAGE_COMMIT=${MAGE_COMMIT}
112+
ARG CUSTOM_MIRROR
113+
114+
# This modifies the apt configuration to rety 3 times
115+
RUN echo 'Acquire::Retries "3";' > /etc/apt/apt.conf.d/80-retries
116+
117+
# If CUSTOM_MIRROR is set, replace the default archive.ubuntu.com
118+
# and security.ubuntu.com URIs in your .sources file
119+
RUN if [ -n "$CUSTOM_MIRROR" ]; then \
120+
sed -E -i \
121+
-e '/^URIs:/ s#https?://[^ ]*archive\.ubuntu\.com#'"$CUSTOM_MIRROR"'#g' \
122+
-e '/^URIs:/ s#https?://security\.ubuntu\.com#'"$CUSTOM_MIRROR"'#g' \
123+
/etc/apt/sources.list.d/ubuntu.sources; \
124+
fi
125+
126+
# Essentials for production/dev
127+
RUN apt-get update && \
128+
apt-get install -y curl && \
129+
curl -sSL -O https://packages.microsoft.com/config/ubuntu/24.04/packages-microsoft-prod.deb && \
130+
dpkg -i packages-microsoft-prod.deb && \
131+
rm packages-microsoft-prod.deb && \
132+
apt-get update && \
133+
ACCEPT_EULA=Y apt-get install -y msodbcsql18 unixodbc-dev && \
134+
apt-get install -y \
135+
libcurl4 `memgraph` \
136+
libpython${PY_VERSION} `memgraph` \
137+
libssl-dev `memgraph` \
138+
openssl `memgraph` \
139+
python3 `mage-memgraph` \
140+
python3-pip `mage-memgraph` \
141+
python3-setuptools `mage-memgraph` \
142+
python3-dev `mage-memgraph` \
143+
libc6-dbg \
144+
adduser \
145+
libgomp1 \
146+
libaio1t64 \
147+
libatomic1 \
148+
--no-install-recommends \
149+
&& ln -s /usr/bin/$(ls /usr/bin | grep perf) /usr/bin/perf \
150+
&& rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \
151+
if [ "${TARGETARCH}" = "arm64" ]; then \
152+
ln -s /usr/lib/aarch64-linux-gnu/libaio.so.1t64 /usr/lib/aarch64-linux-gnu/libaio.so.1; \
153+
else \
154+
ln -s /usr/lib/x86_64-linux-gnu/libaio.so.1t64 /usr/lib/x86_64-linux-gnu/libaio.so.1; \
155+
fi
156+
157+
# install memgraph
158+
COPY memgraph-${TARGETARCH}.deb .
159+
160+
# fix `memgraph` UID and GID for compatibility with previous Debian releases
161+
RUN groupadd -g 103 memgraph && \
162+
useradd -u 101 -g memgraph -m -d /home/memgraph -s /bin/bash memgraph && \
163+
dpkg -i memgraph-${TARGETARCH}.deb && \
164+
rm memgraph-${TARGETARCH}.deb
165+
166+
ENV LD_LIBRARY_PATH /usr/lib/memgraph/query_modules
167+
168+
# Copy modules
169+
COPY --from=builder /usr/lib/memgraph/query_modules/ /usr/lib/memgraph/query_modules/
170+
COPY --from=builder /usr/lib/memgraph/auth_module/ /usr/lib/memgraph/auth_module/
171+
172+
# Copy Python build
173+
COPY --from=builder /usr/local/lib/python${PY_VERSION}/ /usr/local/lib/python${PY_VERSION}/
174+
COPY --from=builder --chown=memgraph:memgraph /home/memgraph/.local/ /home/memgraph/.local/
175+
176+
# copy script to convert to dev container
177+
COPY --from=builder /mage/make-dev-container.sh /make-dev-container.sh
178+
179+
FROM base as prod
180+
181+
USER root
182+
183+
ARG TARGETARCH
184+
ARG PY_VERSION_DEFAULT=3.12
185+
ARG MAGE_COMMIT
186+
ARG BUILD_TYPE
187+
ENV BUILD_TYPE=${BUILD_TYPE}
188+
ENV PY_VERSION ${PY_VERSION_DEFAULT}
189+
ENV MAGE_COMMIT=${MAGE_COMMIT}
190+
ARG CUSTOM_MIRROR
191+
192+
# Copy e2e tests
193+
COPY --from=builder --chown=memgraph:memgraph /mage/e2e/ /mage/e2e/
194+
COPY --from=builder --chown=memgraph:memgraph /mage/e2e_correctness/ /mage/e2e_correctness/
195+
COPY --from=builder --chown=memgraph:memgraph /mage/test_e2e_correctness.py /mage/test_e2e_correctness.py
196+
COPY --from=builder --chown=memgraph:memgraph /mage/run_e2e_correctness_tests.sh /mage/run_e2e_correctness_tests.sh
197+
198+
# Copy requirements
199+
COPY --from=builder --chown=memgraph:memgraph /mage/python/requirements.txt /mage/python/requirements.txt
200+
COPY --from=builder --chown=memgraph:memgraph /mage/python/tests/requirements.txt /mage/python/tests/requirements.txt
201+
202+
RUN if [ -n "$CUSTOM_MIRROR" ]; then \
203+
sed -E -i \
204+
-e "/^URIs:/ s#${CUSTOM_MIRROR}/ubuntu/#https://archive.ubuntu.com/ubuntu/#g" \
205+
-e "/^URIs:/ s#${CUSTOM_MIRROR}/ubuntu/#https://security.ubuntu.com/ubuntu/#g" \
206+
/etc/apt/sources.list.d/ubuntu.sources; \
207+
fi
208+
209+
USER memgraph
210+
EXPOSE 7687
211+
212+
ENTRYPOINT ["/usr/lib/memgraph/memgraph"]
213+
CMD [""]
214+
215+
216+
FROM base as debug
217+
218+
USER root
219+
220+
ARG TARGETARCH
221+
ARG PY_VERSION_DEFAULT=3.12
222+
ARG MAGE_COMMIT
223+
ARG BUILD_TYPE
224+
ENV BUILD_TYPE=${BUILD_TYPE}
225+
ENV PY_VERSION ${PY_VERSION_DEFAULT}
226+
ENV MAGE_COMMIT=${MAGE_COMMIT}
227+
ARG CUSTOM_MIRROR
228+
229+
# Add gdb
230+
RUN apt-get update && apt-get install -y \
231+
gdb \
232+
--no-install-recommends \
233+
&& rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
234+
235+
# Copy entire mage repo
236+
COPY --from=builder --chown=memgraph:memgraph /mage/ /mage/
237+
238+
RUN if [ -n "$CUSTOM_MIRROR" ]; then \
239+
sed -E -i \
240+
-e "/^URIs:/ s#${CUSTOM_MIRROR}/ubuntu/#https://archive.ubuntu.com/ubuntu/#g" \
241+
-e "/^URIs:/ s#${CUSTOM_MIRROR}/ubuntu/#https://security.ubuntu.com/ubuntu/#g" \
242+
/etc/apt/sources.list.d/ubuntu.sources; \
243+
fi
244+
245+
USER memgraph
246+
EXPOSE 7687
247+
248+
ENTRYPOINT ["/usr/lib/memgraph/memgraph"]
249+
CMD [""]

‎Dockerfile.release‎

Lines changed: 20 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -79,22 +79,27 @@ USER memgraph
7979
# First install requirements
8080
RUN python3 -m pip install --no-cache-dir -r /mage/python/requirements.txt --break-system-packages && \
8181
python3 -m pip install --no-cache-dir -r /mage/python/tests/requirements.txt --break-system-packages && \
82-
python3 -m pip install --no-cache-dir -r /usr/lib/memgraph/auth_module/requirements.txt --break-system-packages
83-
84-
# RUN python3 -m pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu124 --break-system-packages && \
85-
# python3 -m pip install --no-cache-dir "transformers>=4.43.0" "sentence-transformers>=3,<5" "filelock" "huggingface_hub>=0.23.0" --break-system-packages && \
86-
# rm -fr /home/memgraph/.cache/pip
87-
88-
# RUN python3 -m pip install --no-cache-dir torch torchvision --index-url https://download.pytorch.org/whl/cu129 --break-system-packages && \
89-
# python3 -m pip install --no-cache-dir "transformers" "sentence-transformers" "filelock" "huggingface_hub" --break-system-packages && \
90-
# rm -fr /home/memgraph/.cache/pip
91-
92-
# TODO: embedding works witho those versions, but pytorch 2.4 has vulnerabilities...
93-
RUN python3 -m pip install --no-cache-dir torch==2.4.1 torchvision==0.19.1 --index-url https://download.pytorch.org/whl/cu124 --break-system-packages && \
94-
python3 -m pip install --no-cache-dir "transformers" "sentence-transformers" "filelock" "huggingface_hub" --break-system-packages && \
82+
python3 -m pip install --no-cache-dir -r /usr/lib/memgraph/auth_module/requirements.txt --break-system-packages && \
83+
if [ "$TARGETARCH" = "arm64" ]; then \
84+
if [ "$CACHE_PRESENT" = "true" ]; then \
85+
echo "Using cached torch packages"; \
86+
python3 -m pip install --no-index --find-links=/mage/wheels/ torch-sparse torch-cluster torch-spline-conv torch-geometric torch-scatter --break-system-packages; \
87+
else \
88+
python3 -m pip install --no-cache-dir torch-sparse torch-cluster torch-spline-conv torch-geometric torch-scatter -f https://data.pyg.org/whl/torch-2.6.0+cpu.html --break-system-packages; \
89+
fi && \
90+
curl -o dgl-2.5.0-cp312-cp312-linux_aarch64.whl https://s3.eu-west-1.amazonaws.com/deps.memgraph.io/wheels/arm64/dgl-2.5.0-cp312-cp312-linux_aarch64.whl && \
91+
python3 -m pip install --no-cache-dir dgl-2.5.0-cp312-cp312-linux_aarch64.whl --break-system-packages; \
92+
else \
93+
if [ "$CACHE_PRESENT" = "true" ]; then \
94+
echo "Using cached torch packages"; \
95+
python3 -m pip install --no-index --find-links=/mage/wheels/ torch-sparse torch-cluster torch-spline-conv torch-geometric torch-scatter --break-system-packages; \
96+
else \
97+
python3 -m pip install --no-cache-dir torch-sparse torch-cluster torch-spline-conv torch-geometric torch-scatter -f https://data.pyg.org/whl/torch-2.6.0+cpu.html --break-system-packages; \
98+
fi && \
99+
python3 -m pip install --no-cache-dir dgl==2.5.0 -f https://data.dgl.ai/wheels/torch-2.6/repo.html --break-system-packages; \
100+
fi && \
95101
rm -fr /home/memgraph/.cache/pip
96102

97-
98103
# Build query modules
99104
SHELL ["/bin/bash", "-c"]
100105
RUN source /opt/toolchain-v6/activate && curl https://sh.rustup.rs -sSf | sh -s -- -y \
@@ -146,13 +151,10 @@ RUN apt-get update && \
146151
apt-get install -y \
147152
libcurl4 `memgraph` \
148153
libpython${PY_VERSION} `memgraph` \
149-
libssl-dev `memgraph` \
150154
openssl `memgraph` \
151155
python3 `mage-memgraph` \
152156
python3-pip `mage-memgraph` \
153157
python3-setuptools `mage-memgraph` \
154-
python3-dev `mage-memgraph` \
155-
libc6-dbg \
156158
adduser \
157159
libgomp1 \
158160
libaio1t64 \
@@ -240,7 +242,7 @@ ARG CUSTOM_MIRROR
240242

241243
# Add gdb
242244
RUN apt-get update && apt-get install -y \
243-
gdb \
245+
gdb libc6-dbg \
244246
--no-install-recommends \
245247
&& rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
246248

0 commit comments

Comments
 (0)