Dockerfile.preflight

# Dockerfile for 3DGS Preflight Check
#
# Detects GPU capabilities and validates the environment before
# long-running training jobs. Uses NVIDIA CUDA runtime base so
# nvidia-smi is available when a GPU host is present; gracefully
# reports CPU-only when no GPU is detected.
#
# Build:
#   docker build -f Dockerfile.preflight -t 3dgs-preflight:latest .
#
# Run on GPU host (detect + report):
#   docker run --rm --gpus all 3dgs-preflight
#
# Run on CPU host (still works — reports CPU-only):
#   docker run --rm 3dgs-preflight
#
# Assert a specific backend:
#   docker run --rm --gpus all 3dgs-preflight --expect gsplat
#   docker run --rm 3dgs-preflight --expect mock
#
# Exit codes: 0 = pass, 1 = fail

# Global ARG — must be before the first FROM to be visible across stages
ARG CUDA_VERSION=12.6.3

# ============================================================================
# Stage 1: Build the preflight binary
# ============================================================================
FROM rust:1.93-bookworm AS builder

WORKDIR /build

RUN apt-get update && apt-get install -y --no-install-recommends \
    pkg-config \
    libssl-dev \
    && rm -rf /var/lib/apt/lists/*

COPY Cargo.toml Cargo.lock ./
COPY src ./src

RUN cargo build --release --locked --bin 3dgs-preflight && \
    strip target/release/3dgs-preflight

# ============================================================================
# Stage 2: Runtime with GPU detection (NVIDIA CUDA base)
# ============================================================================
ARG CUDA_VERSION
FROM nvidia/cuda:${CUDA_VERSION}-runtime-ubuntu24.04

ENV DEBIAN_FRONTEND=noninteractive

# NVIDIA env vars — enable GPU discovery when host has a GPU.
# On CPU-only hosts these are harmless; nvidia-smi simply won't find a device.
ENV NVIDIA_VISIBLE_DEVICES=all \
    NVIDIA_DRIVER_CAPABILITIES=compute,utility

# Install only what the preflight check needs to probe the environment
RUN apt-get update && \
    apt-get install -y --no-install-recommends software-properties-common && \
    add-apt-repository -y universe && \
    apt-get update && \
    apt-get install -y --no-install-recommends \
        ca-certificates \
        python3 \
        python3-venv \
        libpython3.12 \
        ffmpeg \
        colmap \
    && apt-get clean \
    && rm -rf /var/lib/apt/lists/*

# Copy the preflight binary
COPY --from=builder /build/target/release/3dgs-preflight /usr/local/bin/3dgs-preflight
RUN chmod +x /usr/local/bin/3dgs-preflight

ENTRYPOINT ["/usr/local/bin/3dgs-preflight"]