Split dockerfile into multi-stage build to separate pre-req builds, runner, and full layers

Author: joshhvulcan

.dockerignore

@@ -7,3 +7,4 @@ lightning_logs
 wandb
 **/test_data/**/**/*.tif
 **/project_data
+one_off_projects

Dockerfile

@@ -1,37 +1,71 @@
-FROM pytorch/pytorch:2.7.0-cuda12.8-cudnn9-runtime@sha256:7db0e1bf4b1ac274ea09cf6358ab516f8a5c7d3d0e02311bed445f7e236a5d80
+###
+### docker build --ssh default=/path/to/private.key -f Dockerfile .
 
-RUN apt update
-RUN apt install -y libpq-dev ffmpeg libsm6 libxext6 git wget
+ARG BASE=ubuntu:22.04
+ARG BASE_PYTORCH=pytorch/pytorch:2.7.0-cuda12.8-cudnn9-runtime@sha256:7db0e1bf4b1ac274ea09cf6358ab516f8a5c7d3d0e02311bed445f7e236a5d80
+ARG PLATFORM=linux/amd64
 
-# Install Go (used for Satlas smooth_point_labels_viterbi.go).
-RUN wget https://go.dev/dl/go1.22.12.linux-amd64.tar.gz
-RUN rm -rf /usr/local/go && tar -C /usr/local -xzf go1.22.12.linux-amd64.tar.gz
-ENV PATH="${PATH}:/usr/local/go/bin"
+FROM --platform=${PLATFORM} ${BASE} AS tippecanoe
 
-# Install tippecanoe (used by forest loss driver).
-RUN apt install -y build-essential libsqlite3-dev zlib1g-dev
-RUN git clone https://github.com/mapbox/tippecanoe /opt/tippecanoe
-WORKDIR /opt/tippecanoe
+RUN apt-get update && apt-get install -y --no-install-recommends build-essential ca-certificates curl libsqlite3-dev zlib1g-dev
+
+RUN mkdir -p /tmp/tippecanoe && curl -L https://github.com/mapbox/tippecanoe/archive/refs/tags/1.36.0.tar.gz | tar -xvz --strip 1 -C /tmp/tippecanoe
+WORKDIR /tmp/tippecanoe
 RUN make -j
-RUN make install
+RUN PREFIX=/opt/tippecanoe make install
+
+# To use this:
+#  COPY --from=tippecanoe /opt/tippecanoe /opt/tippecanoe
+#  ENV PATH="/opt/tippecanoe/bin:${PATH}"
+
+FROM --platform=${PLATFORM} ${BASE} AS golang
+
+## Build Satlas smooth_point_labels_viterbi.go (Requires golang)
+RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates curl
+RUN curl -L https://go.dev/dl/go1.22.12.linux-amd64.tar.gz | tar -xvz -C /usr/local
+ENV PATH="/usr/local/go/bin:${PATH}"
+
+COPY rslp/satlas/scripts /tmp/smooth_point_labels_viterbi/
+WORKDIR /tmp/smooth_point_labels_viterbi/
+RUN go build smooth_point_labels_viterbi.go
+
+# To use this:
+#  COPY --from=golang /tmp/smooth_point_labels_viterbi/smooth_point_labels_viterbi /usr/local/bin/smooth_point_labels_viterbi
+
+FROM --platform=${PLATFORM} pytorch/pytorch:2.7.0-cuda12.8-cudnn9-runtime@sha256:7db0e1bf4b1ac274ea09cf6358ab516f8a5c7d3d0e02311bed445f7e236a5d80 AS base
+
+RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates git``
+
+COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
+
+COPY . /opt/rslearn_projects/
+# We assume that the caller has already cloned olmoearth_run and olmoearth_pretrain to /opt/rslearn_projects/docker_build/
+RUN echo 'olmoearth_run[runner] @ /opt/rslearn_projects/docker_build/olmoearth_run/' > /opt/rslearn_projects/requirements-olmoearth_run.txt
+RUN echo 'olmoearth_pretrain @ /opt/rslearn_projects/docker_build/olmoearth_pretrain/' > /opt/rslearn_projects/requirements-olmoearth_pretrain.txt
+RUN uv pip install --system /opt/rslearn_projects[olmoearth_run,olmoearth_pretrain]
+
+FROM base AS full
+
+RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates git openssh-client
 
+# Pin GitHub’s host keys so SSH won’t prompt
+RUN mkdir -p -m 700 /root/.ssh && \
+    ssh-keyscan -t rsa,ecdsa,ed25519 github.com >> /root/.ssh/known_hosts \
+
+COPY --from=tippecanoe /opt/tippecanoe /opt/tippecanoe
+COPY --from=golang /tmp/smooth_point_labels_viterbi/smooth_point_labels_viterbi /usr/local/bin/smooth_point_labels_viterbi
 COPY --from=ghcr.io/astral-sh/uv:latest /uv /usr/local/bin/uv
 
-# Install rslearn.
-# We use git clone and then git checkout instead of git clone -b so that the user could
-# specify a commit name or branch instead of only accepting a branch.
+ENV PATH="/opt/tippecanoe/bin:${PATH}"
+
+## Install rslearn.
+## We use git clone and then git checkout instead of git clone -b so that the user could
+## specify a commit name or branch instead of only accepting a branch.
 ARG RSLEARN_BRANCH=master
-RUN git clone https://github.com/allenai/rslearn.git /opt/rslearn
+RUN --mount=type=ssh git clone git@github.com:allenai/rslearn.git /opt/rslearn
 WORKDIR /opt/rslearn
 RUN git checkout $RSLEARN_BRANCH
 RUN uv pip install --system /opt/rslearn[extra]
 
-# Install rslearn_projects.
 COPY . /opt/rslearn_projects/
 RUN uv pip install --system /opt/rslearn_projects[dev,extra]
-
-# Build Satlas smooth_point_labels_viterbi.go program.
-WORKDIR /opt/rslearn_projects/rslp/satlas/scripts
-RUN go build smooth_point_labels_viterbi.go
-
-WORKDIR /opt/rslearn_projects

pyproject.toml

@@ -23,9 +23,12 @@ file = ["requirements.txt"]
 [tool.setuptools.dynamic.optional-dependencies]
 dev = { file = ["requirements-dev.txt"] }
 extra = { file = ["requirements-extra.txt"] }
-helios = { file = ["requirements-helios.txt"] }
+olmoearth_pretrain = { file = ["requirements-olmoearth_pretrain.txt"] }
 olmoearth_run = { file = ["requirements-olmoearth_run.txt"] }
 
+# Kept for backward compatibility
+helios = { file = ["requirements-olmoearth_pretrain.txt"] }
+
 [tool.ruff]
 fix = true
 

requirements-helios.txt renamed to requirements-olmoearth_pretrain.txt

@@ -1 +1 @@
-olmoearth_pretrain @ git+https://github.com/allenai/olmoearth_pretrain.git@0cd82b4784bc2f246a23f6da98a9bab27761c1ba
+olmoearth_pretrain @ git+https://github.com/allenai/olmoearth_pretrain.git@f0a63f190b0f99d9c503249daf7e3e47bbd4792a

requirements-olmoearth_run.txt

@@ -1 +1,2 @@
 olmoearth_run @ git+https://github.com/allenai/olmoearth_run.git@develop
+# olmoearth_runner >= 0.1.0 # Once we have packages publishing.