Dockerfile.worker

ARG CUDA_BUILD_IMAGE=nvidia/cuda:12.8.1-cudnn-devel-ubuntu24.04
ARG CUDA_RUNTIME_IMAGE=nvidia/cuda:12.8.1-cudnn-runtime-ubuntu24.04
ARG NATIVE_BUILD_IMAGE=native-dependencies

FROM ${CUDA_BUILD_IMAGE} AS native-dependencies
ARG NATIVE_BUILD_JOBS=2
RUN case "${NATIVE_BUILD_JOBS}" in 1|2|3|4|5|6|7|8) ;; *) exit 1 ;; esac

ARG MLX_COMMIT=1f8e74e3f12f31365464a6867c6579f0e9b29d85
ARG RUST_VERSION=1.89.0
ARG WHISPER_REVISION=371b5a7561823ab2bb32142d2751e35e7534727b
ARG WHISPER_ARCHIVE_SHA256=89051d8fca516a3ad1f5c2f8f9d2fccb089afbaec338fca3f8731999babc6f81
ARG WHISPER_PATCH_SHA256=e01093ddcab409b7f8f5e2d05df29df566ed5b575dd318c8208719cc398570a6
ENV DEBIAN_FRONTEND=noninteractive
ENV PATH=/root/.cargo/bin:${PATH}
ENV MLX_ROOT=/opt/mlx
ENV WHISPER_SOURCE_DIR=/opt/whisper/source
ENV WHISPER_LIBRARY_DIR=/opt/whisper/build/bin
ENV LD_LIBRARY_PATH=/opt/mlx/lib:/opt/whisper/build/bin:${LD_LIBRARY_PATH}

# The pinned CUDA devel image already installs and holds NCCL runtime/dev.
# Do not request unversioned NCCL packages: apt would try to upgrade those holds.
RUN apt-get update \
    && apt-get install -y --no-install-recommends \
       build-essential clang cmake curl ca-certificates ffmpeg git patch pkg-config \
       libblas-dev liblapack-dev liblapacke-dev libopenblas-dev \
    && test "$(dpkg-query -W -f='${db:Status-Status}' libnccl2)" = installed \
    && test "$(dpkg-query -W -f='${db:Status-Status}' libnccl-dev)" = installed \
    && test "libnccl2=$(dpkg-query -W -f='${Version}' libnccl2)" = "${NV_LIBNCCL_PACKAGE}" \
    && test "libnccl-dev=$(dpkg-query -W -f='${Version}' libnccl-dev)" = "${NV_LIBNCCL_DEV_PACKAGE}" \
    && test -f /usr/include/nccl.h \
    && test -f /usr/lib/x86_64-linux-gnu/libnccl.so \
    && rm -rf /var/lib/apt/lists/* \
    && curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \
       | sh -s -- -y --default-toolchain ${RUST_VERSION} --profile minimal --no-modify-path

# Build the C++ API directly from the exact MLX release. No Python bindings or
# Python runtime are involved. SM 86 is the NVIDIA A10G architecture used by g5.
# CUDA JIT runs after installation, so retain a dereferenced copy of the
# toolkit headers independently of the build tree and its symlinks.
RUN git clone --filter=blob:none --no-checkout https://github.com/ml-explore/mlx.git /tmp/mlx \
    && git -C /tmp/mlx fetch --depth=1 origin ${MLX_COMMIT} \
    && git -C /tmp/mlx checkout --detach ${MLX_COMMIT} \
    && cmake -S /tmp/mlx -B /tmp/mlx/build \
       -DCMAKE_BUILD_TYPE=MinSizeRel \
       -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
       -DBUILD_SHARED_LIBS=ON \
       -DMLX_BUILD_CUDA=ON \
       -DMLX_CUDA_ARCHITECTURES=86 \
       -DMLX_BUILD_METAL=OFF \
       -DMLX_BUILD_CPU=ON \
       -DMLX_BUILD_PYTHON_BINDINGS=OFF \
       -DMLX_BUILD_TESTS=OFF \
       -DMLX_BUILD_EXAMPLES=OFF \
       -DMLX_BUILD_BENCHMARKS=OFF \
       -DMLX_BUILD_SAFETENSORS=ON \
       -DMLX_BUILD_GGUF=OFF \
       -DCMAKE_INSTALL_PREFIX=/opt/mlx \
    && cmake --build /tmp/mlx/build --parallel "${NATIVE_BUILD_JOBS}" \
    && cmake --install /tmp/mlx/build \
    && mkdir -p /opt/build-evidence \
    && cp /tmp/mlx/build/compile_commands.json /opt/build-evidence/mlx-compile-commands.json \
    && mkdir -p /opt/mlx-cuda-headers/include \
    && cp -aL /usr/local/cuda/include/. /opt/mlx-cuda-headers/include/ \
    && rm -rf /tmp/mlx

# Preserve JIT libraries as well as headers: CUDA runtime images need not ship
# NVRTC development dependencies. Never copy libcuda, which the host supplies.
RUN mkdir -p /opt/mlx-cuda-jit-libs \
    && cp -aL /usr/local/cuda/lib64/libnvrtc.so* \
       /usr/local/cuda/lib64/libnvrtc-builtins.so* \
       /usr/local/cuda/lib64/libnvJitLink.so* /opt/mlx-cuda-jit-libs/

# Match tools/prepare_whisper.ts, including its pinned source and Sandchest
# extensions. An explicit SM86 target allows compilation without a live GPU.
COPY workers/inference-rs/native/whisper-sandchest.patch /tmp/whisper-sandchest.patch
RUN curl --fail --location --proto '=https' --proto-redir '=https' \
       --connect-timeout 30 --max-time 300 --max-filesize 67108864 \
       "https://codeload.github.com/ggml-org/whisper.cpp/tar.gz/${WHISPER_REVISION}" \
       --output /tmp/whisper.tar.gz \
    && printf '%s  %s\n' "${WHISPER_ARCHIVE_SHA256}" /tmp/whisper.tar.gz | sha256sum --check --strict \
    && printf '%s  %s\n' "${WHISPER_PATCH_SHA256}" /tmp/whisper-sandchest.patch | sha256sum --check --strict \
    && mkdir -p "${WHISPER_SOURCE_DIR}" \
    && tar -xzf /tmp/whisper.tar.gz --strip-components=1 -C "${WHISPER_SOURCE_DIR}" \
    && patch --batch --fuzz=0 -p1 -d "${WHISPER_SOURCE_DIR}" -i /tmp/whisper-sandchest.patch \
    && cmake -S "${WHISPER_SOURCE_DIR}" -B /opt/whisper/build \
       -DCMAKE_BUILD_TYPE=Release \
       -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
       -DCMAKE_BUILD_RPATH_USE_ORIGIN=ON \
       -DCMAKE_CUDA_ARCHITECTURES=86 \
       -DBUILD_SHARED_LIBS=ON \
       -DGGML_METAL=OFF -DGGML_CUDA=ON -DGGML_BLAS=OFF \
       -DGGML_CCACHE=OFF -DGGML_NATIVE=OFF \
       -DWHISPER_BUILD_TESTS=OFF -DWHISPER_BUILD_EXAMPLES=OFF \
       -DWHISPER_SDL2=OFF -DWHISPER_CURL=OFF \
       -DWHISPER_COMMON_FFMPEG=OFF -DWHISPER_BUILD_SERVER=OFF \
       -DWHISPER_BUILD_NUMBER=4938 \
       "-DWHISPER_BUILD_COMMIT=${WHISPER_REVISION}" \
    && cmake --build /opt/whisper/build --target whisper --parallel "${NATIVE_BUILD_JOBS}" \
    && cp /opt/whisper/build/compile_commands.json /opt/build-evidence/whisper-compile-commands.json \
    && if grep -E -- "(-march=native|-mtune=native|-mavx512|-mamx|-mavxvnni)" /opt/build-evidence/*-compile-commands.json; then exit 1; fi \
    && mkdir -p /opt/whisper/lib \
    && for name in whisper ggml ggml-base ggml-cpu ggml-cuda; do \
         test -f "${WHISPER_LIBRARY_DIR}/lib${name}.so" \
         && cp -aL "${WHISPER_LIBRARY_DIR}/lib${name}.so"* /opt/whisper/lib/ || exit 1; \
       done \
    && rm /tmp/whisper.tar.gz /tmp/whisper-sandchest.patch

# Publish this stage independently, keyed by its source/toolchain inputs. Later
# Rust builds can use that immutable image digest without rebuilding MLX/Whisper.
FROM ${NATIVE_BUILD_IMAGE} AS build
ARG NATIVE_BUILD_JOBS=2
RUN case "${NATIVE_BUILD_JOBS}" in 1|2|3|4|5|6|7|8) ;; *) exit 1 ;; esac

WORKDIR /src
COPY workers/inference-rs/Cargo.toml workers/inference-rs/Cargo.lock workers/inference-rs/build.rs ./workers/inference-rs/
COPY workers/inference-rs/src ./workers/inference-rs/src
COPY workers/inference-rs/native ./workers/inference-rs/native
COPY workers/inference-rs/vendor/uroman ./workers/inference-rs/vendor/uroman
# CPU-only builders lack the driver supplied by the GPU host. Resolve indirect
# libcuda.so.1 symbols against the toolkit stub at link time only. Never add this
# directory to a runtime RPATH or copy the stub into the serving image.
RUN mkdir -p /opt/cuda-link-stubs \
    && test -f /usr/local/cuda/lib64/stubs/libcuda.so \
    && ln -s /usr/local/cuda/lib64/stubs/libcuda.so /opt/cuda-link-stubs/libcuda.so.1 \
    && RUSTFLAGS="-C link-arg=-Wl,-rpath-link,/opt/cuda-link-stubs" \
       cargo build --locked --release --features whisper-mms \
       --manifest-path workers/inference-rs/Cargo.toml --jobs "${NATIVE_BUILD_JOBS}" \
    && ! readelf -d workers/inference-rs/target/release/sandchest-inference \
       | grep -E '(RPATH|RUNPATH).*(stubs|cuda-link)' \
    && rm -rf /opt/cuda-link-stubs

# The build image has no NVIDIA driver, so libcuda is expected to be supplied
# by the container runtime. Every other native dependency must be resolvable
# before producing an image; do not defer missing MLX/CUDA/BLAS libraries to
# first request handling.
RUN set -eu; \
    for elf in /opt/mlx/lib/libmlx.so /src/workers/inference-rs/target/release/sandchest-inference; do \
      missing="$(ldd "${elf}" | awk '/not found$/ { print $1 }')"; \
      bad="$(printf '%s\n' "${missing}" | awk 'NF && $0 !~ /^libcuda[.]so([.]|$)/')"; \
      if [ -n "${bad}" ]; then \
        echo "Unresolved native dependencies for ${elf}: ${bad}" >&2; \
        exit 1; \
      fi; \
    done

# Keep media updates after the native cache boundary. Ubuntu's universe FFmpeg
# security fixes require ESM; build the current signed upstream release instead.
# Only these statically linked FFmpeg libraries/binaries reach the serving image.
FROM ${NATIVE_BUILD_IMAGE} AS media-tools
ARG FFMPEG_VERSION=9.0.1
ARG FFMPEG_ARCHIVE_SHA256=cf38e0e28c7e5605942c4a77755349b0145804a397af37eb1fb4c77cb237f635
ARG FFMPEG_SIGNING_FINGERPRINT=FCF986EA15E6E293A5644F10B4322F04D67658D8
ARG NATIVE_BUILD_JOBS=2
RUN case "${NATIVE_BUILD_JOBS}" in 1|2|3|4|5|6|7|8) ;; *) exit 1 ;; esac
RUN apt-get update \
    && apt-get install -y --no-install-recommends gnupg nasm xz-utils zlib1g-dev \
    && rm -rf /var/lib/apt/lists/* \
    && mkdir -p /tmp/ffmpeg-source /tmp/ffmpeg-gnupg /opt/media-evidence /opt/ffmpeg/share/source \
    && chmod 700 /tmp/ffmpeg-gnupg \
    && curl --fail --location --proto '=https' --proto-redir '=https' \
       --connect-timeout 30 --max-time 300 --max-filesize 67108864 \
       "https://ffmpeg.org/releases/ffmpeg-${FFMPEG_VERSION}.tar.xz" -o /tmp/ffmpeg.tar.xz \
    && curl --fail --location --proto '=https' --proto-redir '=https' \
       --connect-timeout 30 --max-time 60 --max-filesize 16384 \
       "https://ffmpeg.org/releases/ffmpeg-${FFMPEG_VERSION}.tar.xz.asc" -o /tmp/ffmpeg.tar.xz.asc \
    && curl --fail --location --proto '=https' --proto-redir '=https' \
       --connect-timeout 30 --max-time 60 --max-filesize 16384 \
       https://ffmpeg.org/ffmpeg-devel.asc -o /tmp/ffmpeg-devel.asc \
    && printf '%s  %s\n' "${FFMPEG_ARCHIVE_SHA256}" /tmp/ffmpeg.tar.xz | sha256sum --check --strict \
    && gpg --batch --homedir /tmp/ffmpeg-gnupg --import /tmp/ffmpeg-devel.asc \
    && gpg --batch --no-auto-key-retrieve --homedir /tmp/ffmpeg-gnupg --status-fd 1 \
       --verify /tmp/ffmpeg.tar.xz.asc /tmp/ffmpeg.tar.xz > /opt/media-evidence/ffmpeg-signature.txt \
    && grep -F "[GNUPG:] VALIDSIG ${FFMPEG_SIGNING_FINGERPRINT} " /opt/media-evidence/ffmpeg-signature.txt \
    && tar -xJf /tmp/ffmpeg.tar.xz --strip-components=1 -C /tmp/ffmpeg-source \
    && test "$(cat /tmp/ffmpeg-source/VERSION)" = "${FFMPEG_VERSION}" \
    && cd /tmp/ffmpeg-source \
    && ./configure --prefix=/opt/ffmpeg --cpu=x86-64 \
       --disable-autodetect --disable-network --disable-ffplay --disable-doc --disable-debug \
       --disable-shared --enable-static --enable-zlib \
    && make -j "${NATIVE_BUILD_JOBS}" \
    && make install \
    && cp ffbuild/config.mak /opt/media-evidence/ffmpeg-config.mak \
    && cp LICENSE.md COPYING.LGPLv2.1 /opt/ffmpeg/share/source/ \
    && cp /tmp/ffmpeg.tar.xz /opt/ffmpeg/share/source/ffmpeg-${FFMPEG_VERSION}.tar.xz \
    && cp /tmp/ffmpeg.tar.xz.asc /tmp/ffmpeg-devel.asc /opt/ffmpeg/share/source/ \
    && printf '%s  %s\n' "${FFMPEG_ARCHIVE_SHA256}" "ffmpeg-${FFMPEG_VERSION}.tar.xz" > /opt/media-evidence/ffmpeg-source.sha256 \
    && /opt/ffmpeg/bin/ffmpeg -version > /opt/media-evidence/ffmpeg-version.txt \
    && /opt/ffmpeg/bin/ffprobe -version > /opt/media-evidence/ffprobe-version.txt \
    && rm -rf /tmp/ffmpeg-source /tmp/ffmpeg-gnupg /tmp/ffmpeg.tar.xz /tmp/ffmpeg.tar.xz.asc /tmp/ffmpeg-devel.asc

# Generate only owned, benign media for actual-runtime decoder checks. The
# serving tools must still handle seekable fd: input, resampling and channels.
RUN mkdir -p /opt/media-fixtures \
    && /usr/bin/ffmpeg -v error -f lavfi -i 'sine=frequency=440:sample_rate=48000:duration=0.25' \
       -ac 2 -c:a pcm_s16le /opt/media-fixtures/stereo.wav \
    && /usr/bin/ffmpeg -v error -i /opt/media-fixtures/stereo.wav -c:a flac /opt/media-fixtures/stereo.flac \
    && /usr/bin/ffmpeg -v error -i /opt/media-fixtures/stereo.wav -c:a libmp3lame /opt/media-fixtures/stereo.mp3 \
    && /usr/bin/ffmpeg -v error -i /opt/media-fixtures/stereo.wav -c:a aac /opt/media-fixtures/stereo.mp4 \
    && /usr/bin/ffmpeg -v error -i /opt/media-fixtures/stereo.wav -c:a libopus /opt/media-fixtures/stereo.webm \
    && /usr/bin/ffmpeg -v error -i /opt/media-fixtures/stereo.wav -c:a pcm_s16le /opt/media-fixtures/stereo.caf

FROM ${CUDA_RUNTIME_IMAGE} AS runtime-base
COPY --from=build /opt/build-evidence /opt/sandchest/build-evidence
COPY --from=media-tools /opt/media-evidence /opt/sandchest/build-evidence
COPY --from=media-tools /opt/ffmpeg/bin/ffmpeg /opt/ffmpeg/bin/ffprobe /usr/local/bin/
COPY --from=media-tools /opt/ffmpeg/share/source /opt/sandchest/third-party/ffmpeg
COPY --from=media-tools /opt/media-fixtures /tmp/sandchest-media-fixtures

ARG MLX_COMMIT=1f8e74e3f12f31365464a6867c6579f0e9b29d85
ENV DEBIAN_FRONTEND=noninteractive
ENV HOME=/home/sandchest
ENV HF_HOME=/home/sandchest/.cache/huggingface
ENV CUDA_HOME=/opt/sandchest/cuda
# Keep the NVIDIA base image's host-driver locations while adding MLX/CUDA
# user-space libraries. The final ${LD_LIBRARY_PATH} preserves any image-level
# additions supplied by a future CUDA base-image update.
ENV LD_LIBRARY_PATH=/opt/sandchest/lib:/usr/local/cuda/lib64:/usr/local/nvidia/lib:/usr/local/nvidia/lib64:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH}
ENV SANDCHEST_DEVICE=cuda
ENV SANDCHEST_INFERENCE_HOST=0.0.0.0
ENV SANDCHEST_INFERENCE_PORT=8765
ENV SANDCHEST_MLX_COMMIT=${MLX_COMMIT}

# Keep the runtime image's held NCCL version rather than selecting apt's latest.
RUN apt-get update \
    && apt-get install -y --no-install-recommends ca-certificates curl libopenblas0 liblapack3 libgomp1 zlib1g \
    && test "$(dpkg-query -W -f='${db:Status-Status}' libnccl2)" = installed \
    && test "libnccl2=$(dpkg-query -W -f='${Version}' libnccl2)" = "${NV_LIBNCCL_PACKAGE}" \
    && test -f /usr/lib/x86_64-linux-gnu/libnccl.so.2 \
    && rm -rf /var/lib/apt/lists/* \
    && useradd --create-home --uid 10001 --shell /usr/sbin/nologin sandchest \
    && install -d -o sandchest -g sandchest /home/sandchest/.cache/huggingface

WORKDIR /opt/sandchest
COPY --from=build /opt/mlx/lib/ ./lib/
COPY --from=build /opt/whisper/lib/ ./lib/
COPY --from=build /opt/mlx-cuda-jit-libs/ ./lib/
COPY --from=build /opt/mlx/include/ ./include/
COPY --from=build /opt/mlx-cuda-headers/include/ ./cuda/include/
COPY --from=build --chown=sandchest:sandchest /src/workers/inference-rs/target/release/sandchest-inference ./sandchest-inference
COPY workers/inference-rs/LICENSE workers/inference-rs/NOTICE ./

# Validate the actual runtime filesystem, not just the richer devel image.
# All packaged shared objects must resolve; only the host's libcuda may be absent.
RUN set -eu; \
    for name in mlx whisper ggml ggml-base ggml-cpu ggml-cuda; do \
      test -f "/opt/sandchest/lib/lib${name}.so"; \
    done; \
    test ! -e /usr/bin/ffmpeg; test ! -e /usr/bin/ffprobe; \
    for elf in /opt/sandchest/sandchest-inference /opt/sandchest/lib/*.so* /usr/local/bin/ffmpeg /usr/local/bin/ffprobe; do \
      ldd "${elf}" > /tmp/sandchest-ldd.txt; \
      bad="$(awk '/not found$/ && $1 !~ /^libcuda[.]so([.][0-9]+)*$/ { print $1 }' /tmp/sandchest-ldd.txt)"; \
      if [ -n "${bad}" ]; then \
        echo "Unresolved runtime dependencies for ${elf}: ${bad}" >&2; \
        exit 1; \
      fi; \
    done; \
    rm /tmp/sandchest-ldd.txt

RUN set -eu; \
    ffmpeg -version | grep '^ffmpeg version 9[.]0[.]1 '; \
    ffprobe -version | grep '^ffprobe version 9[.]0[.]1 '; \
    ffmpeg -hide_banner -protocols > /opt/sandchest/build-evidence/ffmpeg-protocols.txt; \
    if grep -E '^  (http|https|tcp|udp|tls)$' /opt/sandchest/build-evidence/ffmpeg-protocols.txt; then exit 1; fi; \
    for input in /tmp/sandchest-media-fixtures/*; do \
      ffprobe -v error -protocol_whitelist fd -format_whitelist aac,ac3,aiff,amr,asf,avi,caf,flac,matroska,webm,mov,mp3,mpeg,mpegts,ogg,wav \
        -select_streams a:0 -show_entries stream=channels -of json -i fd: < "${input}" > /tmp/media-probe.json; \
      grep -q '"channels": 2' /tmp/media-probe.json; \
      for pcm in s16le f32le; do \
        ffmpeg -v error -max_alloc 67108864 -threads 1 -probesize 8388608 -analyzeduration 10000000 \
          -protocol_whitelist fd -format_whitelist aac,ac3,aiff,amr,asf,avi,caf,flac,matroska,webm,mov,mp3,mpeg,mpegts,ogg,wav \
          -nostdin -filter_threads 1 -i fd: -map 0:a:0 -vn -sn -dn -threads 1 \
          -f "${pcm}" -ac 1 -acodec "pcm_${pcm}" -ar 16000 -protocol_whitelist pipe pipe:1 \
          < "${input}" > "/tmp/media-${pcm}.pcm"; \
      done; \
      integer_bytes="$(wc -c < /tmp/media-s16le.pcm)"; float_bytes="$(wc -c < /tmp/media-f32le.pcm)"; \
      test "${integer_bytes}" -ge 7000; test "${integer_bytes}" -le 12000; \
      test "${float_bytes}" -eq "$((integer_bytes * 2))"; \
      printf '%s channels=2 pcm16_bytes=%s float32_bytes=%s PASS\n' "${input##*/}" "${integer_bytes}" "${float_bytes}" \
        >> /opt/sandchest/build-evidence/ffmpeg-runtime-checks.txt; \
    done; \
    sha256sum /usr/local/bin/ffmpeg /usr/local/bin/ffprobe >> /opt/sandchest/build-evidence/ffmpeg-runtime-checks.txt; \
    rm -rf /tmp/sandchest-media-fixtures /tmp/media-probe.json /tmp/media-s16le.pcm /tmp/media-f32le.pcm

# Explicit opt-in target: SageMaker supplies the `serve` argument to ENTRYPOINT.
# Mount separately verified model assets read-only; this target never fetches them.
FROM runtime-base AS sagemaker
USER root
ENV SANDCHEST_INFERENCE_PORT=8080
ENV SANDCHEST_SAGEMAKER_MODE=true
ENV SANDCHEST_SHUTDOWN_GRACE_SECONDS=20
EXPOSE 8080
HEALTHCHECK --interval=15s --timeout=5s --start-period=900s --retries=5 \
    CMD curl -fsS http://localhost:8080/ping || exit 1
ENTRYPOINT ["/opt/sandchest/sandchest-inference"]
CMD ["serve"]

# Keep the ordinary worker last so an untargeted docker build retains its
# original non-root startup, model installation, port and readiness behavior.
FROM runtime-base AS runtime
USER sandchest
EXPOSE 8765
HEALTHCHECK --interval=15s --timeout=5s --start-period=900s --retries=5 \
    CMD curl -fsS http://localhost:8765/ready || exit 1
CMD ["sh", "-ec", "/opt/sandchest/sandchest-inference fetch-models && exec /opt/sandchest/sandchest-inference serve"]