Add vertex sentence embeddings and Dockerfile.cuda

.github/workflows/package_mage.yaml

@@ -27,9 +27,9 @@ on:
         description: "Memgraph package download link. Leave empty to use the official download link."
         default: ""
         type: string
-      push_to_dockerhub:
+      cuda:
         type: boolean
-        description: "Push the image to DockerHub?"
+        description: "Build the image for CUDA"
         default: false
       push_to_s3:
         type: boolean
@@ -60,7 +60,7 @@ jobs:
       memgraph_download_link: ${{ github.event.inputs.memgraph_download_link || '' }}
       shorten_tag: 'true'
       force_release: 'true'
-      push_to_dockerhub: ${{ github.event.inputs.push_to_dockerhub || 'false' }}
+      cuda: ${{ inputs.cuda }}
       push_to_s3: ${{ github.event.inputs.push_to_s3 || 'false' }}
       malloc: ${{ github.event.inputs.malloc == 'true' }}
       run_smoke_tests: ${{ inputs.run_smoke_tests }}
@@ -89,7 +89,7 @@ jobs:
       memgraph_download_link: ${{ github.event.inputs.memgraph_download_link || '' }}
       shorten_tag: 'true'
       force_release: 'true'
-      push_to_dockerhub: ${{ github.event.inputs.push_to_dockerhub || 'false' }}
+      cuda: ${{ inputs.cuda }}
       push_to_s3: ${{ github.event.inputs.push_to_s3 || 'false' }}
       malloc: ${{ matrix.malloc }}
       run_smoke_tests: ${{ inputs.run_smoke_tests }}

.github/workflows/reusable_package_mage.yaml

@@ -66,6 +66,10 @@ on:
         type: boolean
         description: "Print the output URL at the end of the workflow"
         default: false
+      cuda:
+        type: boolean
+        description: "Build the image for CUDA"
+        default: false
 
 env:
   OS: "${{ inputs.arch == 'arm64' && 'ubuntu-24.04-aarch64' || 'ubuntu-24.04' }}${{ inputs.malloc && '-malloc' || '' }}"
@@ -203,7 +207,11 @@ jobs:
 
       - name: Set target dockerfile
         run: |
-          DOCKERFILE="Dockerfile.release"
+          if [[ "${{ inputs.cuda }}" == true ]]; then
+            DOCKERFILE="Dockerfile.cuda"
+          else
+            DOCKERFILE="Dockerfile.release"
+          fi
           echo "DOCKERFILE=${DOCKERFILE}" >> $GITHUB_ENV
 
       - name: Check if specified version tag is already pushed

Dockerfile.cuda

@@ -0,0 +1,249 @@
+ARG PY_VERSION_DEFAULT=3.12
+ARG MGBUILD_IMAGE
+
+FROM $MGBUILD_IMAGE as builder
+
+USER root
+
+ARG TARGETARCH
+ARG PY_VERSION_DEFAULT
+ARG BUILD_TYPE
+ARG CACHE_PRESENT
+ENV BUILD_TYPE=${BUILD_TYPE}
+ENV PY_VERSION ${PY_VERSION_DEFAULT}
+ARG CUSTOM_MIRROR
+
+
+# This modifies the apt configuration to rety 3 times
+RUN echo 'Acquire::Retries "3";' > /etc/apt/apt.conf.d/80-retries
+
+# If CUSTOM_MIRROR is set, replace the default archive.ubuntu.com
+# and security.ubuntu.com URIs in your .sources file
+RUN if [ -n "$CUSTOM_MIRROR" ]; then \
+      sed -E -i \
+        -e '/^URIs:/ s#https?://[^ ]*archive\.ubuntu\.com#'"$CUSTOM_MIRROR"'#g' \
+        -e '/^URIs:/ s#https?://security\.ubuntu\.com#'"$CUSTOM_MIRROR"'#g' \
+        /etc/apt/sources.list.d/ubuntu.sources; \
+    fi
+
+# Essentials for production/dev
+RUN apt-get update && apt-get install -y \
+    libcurl4        `memgraph` \
+    libpython${PY_VERSION}   `memgraph` \
+    libssl-dev       `memgraph` \
+    openssl         `memgraph` \
+    build-essential `mage-memgraph` \
+    cmake           `mage-memgraph` \
+    curl            `mage-memgraph` \
+    g++             `mage-memgraph` \
+    python3         `mage-memgraph` \
+    python3-pip     `mage-memgraph` \
+    python3-setuptools     `mage-memgraph` \
+    python3-dev     `mage-memgraph` \
+    clang           `mage-memgraph` \
+    git             `mage-memgraph` \
+    unixodbc-dev        `mage-memgraph` \
+    libboost-all-dev `mage-memgraph` \
+    uuid-dev \
+    gdb \
+    procps \
+    libc6-dbg \
+    libxmlsec1-dev xmlsec1 \
+    libatomic1 \
+    --no-install-recommends \
+    && ln -s /usr/bin/$(ls /usr/bin | grep perf) /usr/bin/perf \
+    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+
+COPY memgraph-${TARGETARCH}.deb .
+
+RUN dpkg -i memgraph-${TARGETARCH}.deb && rm memgraph-${TARGETARCH}.deb
+
+# move memgraph HOME so that mounting /var/lib/memgraph as a volume doesn't break Python
+RUN mkdir -pv /home/memgraph && \
+    usermod -d /home/memgraph memgraph && \
+    chown -R memgraph:memgraph /home/memgraph
+
+ENV LD_LIBRARY_PATH /usr/lib/memgraph/query_modules
+
+
+WORKDIR /mage
+COPY . /mage
+
+#hacks so we can do everything under the `memgraph` user
+RUN chown -R memgraph: /mage && \
+    mkdir -pv /usr/lib/memgraph/query_modules && \
+    chown -R memgraph: /usr/lib/memgraph/query_modules
+USER memgraph
+
+# First install requirements
+RUN python3 -m pip install --no-cache-dir -r /mage/python/requirements-gpu.txt --break-system-packages && \
+    python3 -m pip install --no-cache-dir -r /mage/python/tests/requirements.txt --break-system-packages && \
+    python3 -m pip install --no-cache-dir -r /usr/lib/memgraph/auth_module/requirements.txt --break-system-packages && \
+    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; \
+    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; \
+    rm -fr /home/memgraph/.cache/pip
+
+# Build query modules
+SHELL ["/bin/bash", "-c"]
+RUN source /opt/toolchain-v6/activate && curl https://sh.rustup.rs -sSf | sh -s -- -y \
+    && export PATH="$HOME/.cargo/bin:${PATH}" \
+    && rustup toolchain install 1.85 \
+    && rustup default 1.85 \
+    && python3 /mage/setup build -p /usr/lib/memgraph/query_modules/ --cpp-build-flags CMAKE_BUILD_TYPE=${BUILD_TYPE} \
+    && chown -R memgraph: /mage
+
+
+# Memgraph listens for Bolt Protocol on this port by default.
+EXPOSE 7687
+
+
+
+FROM ubuntu:24.04 as base
+
+USER root
+
+ARG TARGETARCH
+ARG PY_VERSION_DEFAULT=3.12
+ARG MAGE_COMMIT
+ARG BUILD_TYPE
+ENV BUILD_TYPE=${BUILD_TYPE}
+ENV PY_VERSION ${PY_VERSION_DEFAULT}
+ENV MAGE_COMMIT=${MAGE_COMMIT}
+ARG CUSTOM_MIRROR
+
+# This modifies the apt configuration to rety 3 times
+RUN echo 'Acquire::Retries "3";' > /etc/apt/apt.conf.d/80-retries
+
+# If CUSTOM_MIRROR is set, replace the default archive.ubuntu.com
+# and security.ubuntu.com URIs in your .sources file
+RUN if [ -n "$CUSTOM_MIRROR" ]; then \
+      sed -E -i \
+        -e '/^URIs:/ s#https?://[^ ]*archive\.ubuntu\.com#'"$CUSTOM_MIRROR"'#g' \
+        -e '/^URIs:/ s#https?://security\.ubuntu\.com#'"$CUSTOM_MIRROR"'#g' \
+        /etc/apt/sources.list.d/ubuntu.sources; \
+    fi
+
+# Essentials for production/dev
+RUN apt-get update && \
+    apt-get install -y curl && \
+    curl -sSL -O https://packages.microsoft.com/config/ubuntu/24.04/packages-microsoft-prod.deb && \
+    dpkg -i packages-microsoft-prod.deb && \
+    rm packages-microsoft-prod.deb && \
+    apt-get update && \
+    ACCEPT_EULA=Y apt-get install -y msodbcsql18 unixodbc-dev && \
+    apt-get install -y \
+    libcurl4        `memgraph` \
+    libpython${PY_VERSION}   `memgraph` \
+    libssl-dev       `memgraph` \
+    openssl         `memgraph` \
+    python3         `mage-memgraph` \
+    python3-pip     `mage-memgraph` \
+    python3-setuptools     `mage-memgraph` \
+    python3-dev     `mage-memgraph` \
+    libc6-dbg \
+    adduser \
+    libgomp1 \
+    libaio1t64 \
+    libatomic1 \
+    --no-install-recommends \
+    && ln -s /usr/bin/$(ls /usr/bin | grep perf) /usr/bin/perf \
+    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \
+    if [ "${TARGETARCH}" = "arm64" ]; then \
+      ln -s /usr/lib/aarch64-linux-gnu/libaio.so.1t64 /usr/lib/aarch64-linux-gnu/libaio.so.1; \
+    else \
+      ln -s /usr/lib/x86_64-linux-gnu/libaio.so.1t64 /usr/lib/x86_64-linux-gnu/libaio.so.1; \
+    fi
+
+# install memgraph
+COPY memgraph-${TARGETARCH}.deb .
+
+# fix `memgraph` UID and GID for compatibility with previous Debian releases
+RUN groupadd -g 103 memgraph && \
+    useradd -u 101 -g memgraph -m -d /home/memgraph -s /bin/bash memgraph && \
+    dpkg -i memgraph-${TARGETARCH}.deb && \
+    rm memgraph-${TARGETARCH}.deb 
+
+ENV LD_LIBRARY_PATH /usr/lib/memgraph/query_modules
+
+# Copy modules
+COPY --from=builder /usr/lib/memgraph/query_modules/ /usr/lib/memgraph/query_modules/
+COPY --from=builder /usr/lib/memgraph/auth_module/ /usr/lib/memgraph/auth_module/
+
+# Copy Python build
+COPY --from=builder /usr/local/lib/python${PY_VERSION}/ /usr/local/lib/python${PY_VERSION}/
+COPY --from=builder --chown=memgraph:memgraph /home/memgraph/.local/ /home/memgraph/.local/ 
+
+# copy script to convert to dev container
+COPY --from=builder /mage/make-dev-container.sh /make-dev-container.sh
+
+FROM base as prod
+
+USER root
+
+ARG TARGETARCH
+ARG PY_VERSION_DEFAULT=3.12
+ARG MAGE_COMMIT
+ARG BUILD_TYPE
+ENV BUILD_TYPE=${BUILD_TYPE}
+ENV PY_VERSION ${PY_VERSION_DEFAULT}
+ENV MAGE_COMMIT=${MAGE_COMMIT}
+ARG CUSTOM_MIRROR
+
+# Copy e2e tests
+COPY --from=builder --chown=memgraph:memgraph /mage/e2e/ /mage/e2e/
+COPY --from=builder --chown=memgraph:memgraph /mage/e2e_correctness/ /mage/e2e_correctness/
+COPY --from=builder --chown=memgraph:memgraph /mage/test_e2e_correctness.py /mage/test_e2e_correctness.py
+COPY --from=builder --chown=memgraph:memgraph /mage/run_e2e_correctness_tests.sh /mage/run_e2e_correctness_tests.sh
+
+# Copy requirements
+COPY --from=builder --chown=memgraph:memgraph /mage/python/requirements.txt /mage/python/requirements.txt
+COPY --from=builder --chown=memgraph:memgraph /mage/python/tests/requirements.txt /mage/python/tests/requirements.txt 
+
+RUN if [ -n "$CUSTOM_MIRROR" ]; then \
+      sed -E -i \
+        -e "/^URIs:/ s#${CUSTOM_MIRROR}/ubuntu/#https://archive.ubuntu.com/ubuntu/#g" \
+        -e "/^URIs:/ s#${CUSTOM_MIRROR}/ubuntu/#https://security.ubuntu.com/ubuntu/#g" \
+        /etc/apt/sources.list.d/ubuntu.sources; \
+    fi
+
+USER memgraph
+EXPOSE 7687
+
+ENTRYPOINT ["/usr/lib/memgraph/memgraph"]
+CMD [""]
+
+
+FROM base as debug
+
+USER root
+
+ARG TARGETARCH
+ARG PY_VERSION_DEFAULT=3.12
+ARG MAGE_COMMIT
+ARG BUILD_TYPE
+ENV BUILD_TYPE=${BUILD_TYPE}
+ENV PY_VERSION ${PY_VERSION_DEFAULT}
+ENV MAGE_COMMIT=${MAGE_COMMIT}
+ARG CUSTOM_MIRROR
+
+# Add gdb
+RUN apt-get update && apt-get install -y \
+    gdb \
+    --no-install-recommends \
+    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* 
+
+# Copy entire mage repo
+COPY --from=builder --chown=memgraph:memgraph /mage/ /mage/
+
+RUN if [ -n "$CUSTOM_MIRROR" ]; then \
+      sed -E -i \
+        -e "/^URIs:/ s#${CUSTOM_MIRROR}/ubuntu/#https://archive.ubuntu.com/ubuntu/#g" \
+        -e "/^URIs:/ s#${CUSTOM_MIRROR}/ubuntu/#https://security.ubuntu.com/ubuntu/#g" \
+        /etc/apt/sources.list.d/ubuntu.sources; \
+    fi
+
+USER memgraph
+EXPOSE 7687
+
+ENTRYPOINT ["/usr/lib/memgraph/memgraph"]
+CMD [""]

Dockerfile.release

@@ -260,4 +260,4 @@ USER memgraph
 EXPOSE 7687
 
 ENTRYPOINT ["/usr/lib/memgraph/memgraph"]
-CMD [""]
+CMD [""]

python/embed_worker/embed_worker.py

@@ -0,0 +1,34 @@
+import os
+
+
+def encode_chunk(visible_gpu: int | None, model_id: str, texts: list[str], batch_size: int):
+    """
+    Runs in a spawned process. Returns (count, embeddings_as_list)
+    """
+    # Set device visibility BEFORE importing torch/transformers
+    if visible_gpu is None:
+        os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    else:
+        os.environ["CUDA_VISIBLE_DEVICES"] = str(visible_gpu)
+
+    os.environ.setdefault("HF_HUB_DISABLE_TELEMETRY", "1")
+    os.environ.setdefault("TRANSFORMERS_NO_TORCHVISION", "1")
+    os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
+
+    import torch
+    from sentence_transformers import SentenceTransformer
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = SentenceTransformer(model_id, device=device)
+
+    if not texts:
+        return 0, []
+
+    embs = model.encode(
+        texts,
+        batch_size=min(batch_size, len(texts)),
+        convert_to_numpy=True,
+        normalize_embeddings=True,
+        show_progress_bar=False,
+    )
+    return len(texts), embs.tolist()