Merge pull request #834 from Dstack-TEE/perf/mkosi-parallel-squashfs

perf(os/mkosi): make cached builds reuse what has not changed

Author: kvinwang

os/mkosi/README.md

@@ -88,6 +88,13 @@ DSTACK_DEV_CACHE_DIR="$HOME/.cache/dstack/mkosi-dev" \
 ./os/mkosi/build.sh --no-cache image "$PWD/os/mkosi/build"
 ```
 
+A cached build also skips the two release tarballs, which are roughly a minute
+of gzip over artifacts that already sit unpacked beside them. `disk.raw`, the
+measurements and `metadata.json` are still produced, so QEMU smoke-testing and
+measurement inspection are unaffected; `--archive` asks a cached build for the
+tarballs anyway. A cold build always archives, and `repro-check` archives
+unconditionally because the tarballs are what it compares.
+
 The cache covers dstack Rust, image tools, the container stack, Sysbox, nvattest, the
 kernel build tree, NVIDIA, ZFS and both OVMF variants. Its key conservatively
 includes the inputs, tools, packages and component dependencies declared by

os/mkosi/build.sh

@@ -7,18 +7,26 @@ ROOT=$(cd "$SELF/../.." && pwd)
 # shellcheck source=/dev/null
 source "$SELF/versions.env"
 FLAVORS=${FLAVORS:-prod}
-usage="Usage: $0 [--no-cache] {image|repro-check|lint} [build-dir]"
+usage="Usage: $0 [--no-cache] [--archive] {image|repro-check|lint} [build-dir]"
 # The component cache is keyed on every declared input of each component, so a
 # hit reproduces the same output a cold build would have produced. Reusing it is
 # therefore the sensible default; --no-cache forces the cold path for a release
 # build or when the key itself is what needs auditing. repro-check ignores both
 # and always builds cold, because a cache hit would answer the wrong question.
 cache=${DSTACK_COMPONENT_CACHE:-1}
+# The two release tarballs are ~58 s of gzip over artifacts that already exist
+# unpacked beside them. A cached build is an iteration build, and iterating on
+# the guest wants disk.raw and the measurements, not a redistributable archive
+# -- so a cached build skips them and a cold build still produces them. Pass
+# --archive to get them out of a cached build anyway. Left unset here so the
+# default can be derived from the final value of $cache below.
+archive=${DSTACK_TAR_RELEASE:-}
 action=
 BUILD_DIR=
 while [ $# -gt 0 ]; do
   case "$1" in
     --no-cache) cache=0 ;;
+    --archive) archive=1 ;;
     -h|--help) echo "$usage"; exit 0 ;;
     -*) echo "Unknown option: $1" >&2; echo "$usage" >&2; exit 2 ;;
     *)
@@ -40,6 +48,18 @@ if [[ $action == repro-check ]]; then cache=0; fi
 [[ $cache == 1 || $cache == 0 ]] || {
   echo "DSTACK_COMPONENT_CACHE must be 0 or 1, got: $cache" >&2; exit 2;
 }
+# Derived after repro-check has forced the cold path, so repro-check always
+# archives: it compares the two release tarballs, and skipping them would leave
+# it comparing nothing and passing vacuously.
+archive=${archive:-$(( cache == 1 ? 0 : 1 ))}
+# Not merely defaulted: forced. repro-check compares the two release tarballs,
+# so an environment that switched archiving off would leave it comparing files
+# that do not exist -- a check that fails for the wrong reason, or worse, is
+# read as "no difference found".
+if [[ $action == repro-check ]]; then archive=1; fi
+[[ $archive == 1 || $archive == 0 ]] || {
+  echo "DSTACK_TAR_RELEASE must be 0 or 1, got: $archive" >&2; exit 2;
+}
 if [[ $action == lint ]]; then exec "$SELF/tests/acceptance.sh"; fi
 command -v mkosi >/dev/null || { echo "mkosi $MKOSI_VERSION is required" >&2; exit 1; }
 actual=$(mkosi --version | awk '{print $2}' | cut -d. -f1)
@@ -81,6 +101,39 @@ if [[ $cache == 0 ]]; then
   mkosi --directory "$SELF" --output-directory="$BUILD_DIR/out" clean -f
 fi
 
+# Digest of every input that determines the incrementally cached tree: the
+# package lists and distribution pins in the configs, and the skeleton files
+# copied in before packages are installed. Deliberately not the dstack sources
+# or the component definitions -- those are consumed by the build script, which
+# runs after the cache is restored, and folding them in would defeat the cache
+# on every source edit. Sorted for a stable digest, and the file list itself is
+# hashed too so that deleting a skeleton file also moves the key.
+base_inputs_digest() {
+  local flavor=$1 skeleton
+  local skeletons=("$SELF/mkosi.skeleton")
+  if [[ -d $SELF/mkosi.profiles/$flavor/mkosi.skeleton ]]; then
+    skeletons+=("$SELF/mkosi.profiles/$flavor/mkosi.skeleton")
+  fi
+  {
+    # Hash relative names, types and modes as well as contents. Relative names
+    # let worktrees share a cache; types and modes cover inputs that sha256sum
+    # alone cannot distinguish.
+    for skeleton in "${skeletons[@]}"; do
+      echo "-- skeleton ${skeleton#"$SELF"/}"
+      (cd "$skeleton" && find . -mindepth 1 -printf '%P %y %m\n' | sort)
+      (cd "$skeleton" && find . -type f -print0 | sort -z | xargs -0 -r sha256sum)
+      (cd "$skeleton" && find . -type l -print0 | sort -z | \
+        while IFS= read -r -d '' link; do
+          printf '%s -> %s\n' "${link#./}" "$(readlink "$link")"
+        done)
+    done
+    echo "-- mkosi.conf"
+    cat "$SELF/mkosi.conf" "$SELF/mkosi.tools.conf"
+    echo "-- profile/$flavor/mkosi.conf"
+    cat "$SELF/mkosi.profiles/$flavor/mkosi.conf"
+  } | sha256sum | cut -c1-32
+}
+
 build_one() {
   local out=$1 flavor=$2 jobs=${3:-${JOBS:-$(nproc)}}
   mkdir -p "$out"
@@ -91,6 +144,7 @@ build_one() {
     --profile="$flavor"
     --source-date-epoch="$SOURCE_DATE_EPOCH"
     --environment="DSTACK_COMPONENT_CACHE=$cache"
+    --environment="DSTACK_TAR_RELEASE=$archive"
     --environment="DSTACK_BUILD_GIT_REVISION=$DSTACK_BUILD_GIT_REVISION"
     --environment="DSTACK_SOURCE_REVISION=$revision"
     --environment="JOBS=$jobs"
@@ -109,6 +163,25 @@ build_one() {
     mkosi_args+=(--build-directory="$cache_root/build")
     mkosi_args+=(--build-sources="$cache_root/manifest:component-cache")
     mkosi_args+=(--environment="DSTACK_SOURCE_MANIFEST=/work/src/component-cache/source-manifest")
+    # mkosi's incremental cache captures the tree after the distribution and
+    # build packages are installed and before any build script runs, which is
+    # exactly the ~70 s this build otherwise repeats verbatim every time.
+    #
+    # It cannot be enabled as-is. mkosi keys the cache on CacheKey=, whose
+    # default is &d~&r~&a~&I -- distribution, release, architecture, image id --
+    # and whose specifier set contains no digest of the inputs that actually
+    # determine that tree. Editing Packages= in mkosi.conf, or any file under
+    # mkosi.skeleton/, leaves the key untouched, so mkosi would restore the
+    # stale tree and the change would silently not be in the image. Mixing the
+    # digest below into the key restores the invalidation mkosi does not do.
+    base_digest=$(base_inputs_digest "$flavor")
+    mkosi_args+=(--incremental=yes)
+    mkosi_args+=(--cache-directory="$cache_root/incremental")
+    mkosi_args+=(--cache-key="&d~&r~&a~&I~$base_digest")
+    # Package downloads are content-addressed by the pinned Snapshot=, so this
+    # only avoids refetching identical files. Release builds still take the
+    # cold path and fetch from the snapshot themselves.
+    mkosi_args+=(--package-cache-directory="$cache_root/packages")
   fi
   mkosi "${mkosi_args[@]}" build
   # mkosi's own output is a plain Debian rootfs: unmeasured, not part of the

os/mkosi/components/dstack-rust/dstack-rust-build.sh

@@ -28,10 +28,40 @@ install -m0644 "$ROOT/os/common/rootfs/containerd.service.d/"* \
 if [[ ${DSTACK_SKIP_RUST:-0} != 1 ]]; then
   export CARGO_INCREMENTAL=0 CARGO_NET_OFFLINE=${CARGO_NET_OFFLINE:-false}
   build_root=$(dirname "$DEST")
-  export CARGO_TARGET_DIR="$build_root/dstack-cargo-target"
+  ephemeral_target="$build_root/dstack-cargo-target"
+  target_remap=
+  if [[ -n ${DSTACK_CARGO_TARGET_DIR:-} ]]; then
+    # A cached build keeps the target directory outside the per-build sandbox
+    # so an edit rebuilds only the touched crates. That moves it out from under
+    # the $build_root remap below, and rustc embeds absolute paths from the
+    # target directory (build script OUT_DIRs among them). Without mapping it
+    # back onto the path the throwaway layout produced, a cached build would
+    # differ from a cold one byte for byte -- silently, and only inside the
+    # binaries. Remapped first so it wins over the broader $build_root rule.
+    export CARGO_TARGET_DIR="$DSTACK_CARGO_TARGET_DIR"
+    target_remap="--remap-path-prefix=$CARGO_TARGET_DIR=/usr/src/dstack-build/dstack-cargo-target"
+  else
+    export CARGO_TARGET_DIR="$ephemeral_target"
+  fi
+  # CARGO_HOME needs the same treatment as the target directory, and for a
+  # sharper reason: it holds the git checkouts of dependencies, and rustc
+  # embeds their source paths in the binary (panic locations, tracing call
+  # sites). Moving it without remapping made a cached build's binaries differ
+  # from a cold build's -- the rootfs hash and every measurement downstream
+  # with them -- while every path that was remapped looked identical.
+  # Not `[[ ... ]] && export ...`: under set -e a false test makes the whole
+  # list non-zero and kills the build on the cold path, where this is unset.
+  if [[ -n ${DSTACK_CARGO_HOME:-} ]]; then export CARGO_HOME="$DSTACK_CARGO_HOME"; fi
+  # Remapped unconditionally, not only when it moved: both layouts have to land
+  # on the same string for their binaries to match, so the inherited path needs
+  # the rule as much as the persisted one does.
+  # Keep the cold build's historical canonical path. Cached builds map their
+  # persisted CARGO_HOME onto it, preserving both warm/cold parity and release
+  # artifact bytes from before the cache optimization.
+  home_remap="--remap-path-prefix=${CARGO_HOME:?}=/var/tmp/dstack-cargo-home"
   # A single codegen unit avoids LLVM partition/scheduling differences across
   # hosts with different CPU counts while retaining parallel crate builds.
-  export RUSTFLAGS="${RUSTFLAGS:-} --remap-path-prefix=$ROOT=/usr/src/dstack --remap-path-prefix=$build_root=/usr/src/dstack-build -C codegen-units=1 -C strip=debuginfo"
+  export RUSTFLAGS="${RUSTFLAGS:-} $target_remap $home_remap --remap-path-prefix=$ROOT=/usr/src/dstack --remap-path-prefix=$build_root=/usr/src/dstack-build -C codegen-units=1 -C strip=debuginfo"
   cargo build --locked --release --manifest-path "$ROOT/dstack/Cargo.toml" \
     -p dstack-guest-agent -p dstack-util
   install -m0755 "$CARGO_TARGET_DIR/release/dstack-guest-agent" \

os/mkosi/mkosi.build

@@ -33,6 +33,14 @@ if [[ $cache_enabled == 1 ]]; then
     : "${BUILDDIR:?dev component cache requires mkosi BuildDirectory}"
     export DSTACK_DEV_CACHE_DIR="$BUILDDIR/components"
     toolchain_downloads="$BUILDDIR/toolchains"
+    # A source edit misses the dstack-rust component cache by design, and the
+    # rebuild that follows was recompiling the entire workspace -- every
+    # dependency crate, and re-downloading the registry -- because both the
+    # target directory and CARGO_HOME lived in the per-build sandbox. Keeping
+    # them here leaves only the touched crates to rebuild. Release builds still
+    # get the throwaway layout below.
+    export DSTACK_CARGO_TARGET_DIR="$BUILDDIR/cargo-target"
+    export DSTACK_CARGO_HOME="$BUILDDIR/cargo-home"
 else
     toolchain_downloads=/var/tmp/dstack-toolchain-downloads
 fi
@@ -45,8 +53,45 @@ printf 'cmake: '; cmake --version | head -1
 dpkg-query -W -f='build-package: ${binary:Package}=${Version}\n' \
   gcc g++ binutils cmake make ninja-build dwarves
 
-work=/var/tmp/dstack-component-work
+# component-stages/ is where cached component archives get extracted. Keeping
+# it in the per-build sandbox means every cache hit still pays a full
+# decompression -- the kernel archive alone unpacks 6.4 GiB and costs ~12 s --
+# so a cached build persists it and dev-cache.sh skips extracting a tree that
+# is already staged under the same key.
+#
+# Only that subdirectory. The rest of the work directory holds source checkouts
+# and object trees, and the component scripts assume those start empty: sysbox
+# runs `git submodule update` there and fails outright on a leftover tree.
+# Persisting the staging trees buys the decompression; persisting the build
+# scratch buys nothing and breaks the builds.
+#
+# The kernel staging tree is not a cache either. component_assemble merges into
+# it on every build, cache hit or not, under strict conflict detection, so a
+# persisted one collides with the previous build's merged files.
+#
+# downloads/ and image-tools-target/ are kept for the same reason: the first is
+# fetched tarballs and submodules, the second a cargo target directory, and
+# neither holds git or make state that a rerun could trip over. Everything else
+# is scratch and is deleted by exclusion, so a directory added later is dropped
+# until it is named here deliberately.
+work_keep=(component-stages downloads image-tools-target)
+if [[ $cache_enabled == 1 ]]; then
+    work="$BUILDDIR/component-work"
+    mkdir -p "$work"
+    prune=(-mindepth 1 -maxdepth 1)
+    for keep in "${work_keep[@]}"; do
+        mkdir -p "$work/$keep"
+        prune+=(! -name "$keep")
+    done
+    # Deleting by exclusion rather than symlinking the kept directories
+    # elsewhere: tar replaces a symlink standing where it wants a directory, so
+    # a staging tree would silently stop being the persisted one.
+    find "$work" "${prune[@]}" -exec rm -rf {} +
+else
+    work=/var/tmp/dstack-component-work
+fi
 kernel_stage=/var/tmp/dstack-kernel-stage
+rm -rf "$kernel_stage"
 mkdir -p "$work" "$kernel_stage"
 args=()
 [[ $cache_enabled == 1 ]] && args+=(--dev-cache)

os/mkosi/mkosi.postoutput

@@ -12,6 +12,7 @@ WORK="$OUTPUTDIR/.dstack-$FLAVOR"
 DSTACK_MR_BIN="$OUTPUTDIR/.image-tools/dstack-mr" \
 NITRO_TPM_PCR_COMPUTE_BIN="$OUTPUTDIR/.image-tools/nitro-tpm-pcr-compute" \
 DIST_DIR="$OUTPUTDIR" \
+DSTACK_TAR_RELEASE="${DSTACK_TAR_RELEASE:-1}" \
 TAR_OPTIONS="--sort=name --mtime=@$SOURCE_DATE_EPOCH --owner=0 --group=0 --numeric-owner" \
   "$ROOT/os/image/assemble.sh" --manifest "$WORK/artifacts/artifact-manifest.json"
 

os/mkosi/scripts/dev-cache.sh

@@ -27,23 +27,65 @@ dev_cache_run() {
     checksum="$archive.sha256"
     mkdir -p "$dir" "$base"
 
+    # Records which key the staging tree currently holds. It lives beside the
+    # archive rather than in the staging tree's parent: mkosi.build wipes
+    # everything in that parent except the staging trees themselves, so a stamp
+    # kept there is deleted before every build and the tree is then always
+    # re-extracted -- the persistence buys nothing and the wipe costs extra.
+    local stamp="$dir/staged"
+
+    # Anything staged under a different key -- or under no key we can account
+    # for -- is removed before it can be reused. Skipping this is how files
+    # from a previous build survive into a later image, which is far worse than
+    # the extraction it would have saved. Only the component's declared outputs
+    # are touched, so components sharing a base do not disturb each other.
+    discard_staged() {
+        (cd "$base" && rm -rf -- "${outputs[@]}")
+        rm -f "$stamp"
+    }
+
+    staged_outputs_present() {
+        local out
+        for out in "${outputs[@]}"; do
+            [[ -e $base/$out ]] || return 1
+        done
+    }
+
     (
         flock 9
         if [[ -f $archive && -f $checksum ]] &&
            (cd "$dir" && sha256sum --check --status "${checksum##*/}"); then
+            # The stamp alone is not enough. It outlives the staging tree, so
+            # anything that removes the tree without clearing it -- a manual
+            # wipe, a pruned build directory -- would make this skip an
+            # extraction whose output is not there, and the component would be
+            # silently missing from the image.
+            if [[ -f $stamp && $(cat "$stamp") == "$key" ]] && staged_outputs_present; then
+                echo "development cache hit (already staged): $component"
+                exit
+            fi
+            discard_staged
             tar --zstd -xf "$archive" -C "$base"
+            printf '%s' "$key" > "$stamp"
             echo "development cache hit: $component"
             exit
         fi
 
         local tmp
         rm -f "$archive" "$checksum"
         echo "development cache miss: $component"
+        discard_staged
         "$@"
         tmp=$(mktemp "$dir/.${key}.XXXXXX.tar.zst")
         tar --zstd -cf "$tmp" -C "$base" "${outputs[@]}"
         mv "$tmp" "$archive"
         (cd "$dir" && sha256sum "${archive##*/}" > "${checksum##*/}.tmp")
         mv "$checksum.tmp" "$checksum"
-    ) 9>"$dir/$key.lock"
+        # Written last: a stamp is a claim that the tree matches the archive,
+        # so it must not exist if any step above failed.
+        printf '%s' "$key" > "$stamp"
+    # The staged tree and its stamp are shared by every key for this component.
+    # Lock the component, not an individual archive: otherwise two builds with
+    # different keys can concurrently delete and extract the same outputs.
+    ) 9>"$dir/staging.lock"
 }

os/mkosi/scripts/make-release-artifacts.sh

@@ -25,14 +25,19 @@ truncate -s 0 "$rootfs"
 # Privileged mkosi runs can inherit host default ACLs from their workspace.
 # The guest rootfs does not rely on xattrs, so exclude this host-only metadata.
 # A sorted tar stream also removes backing-filesystem directory/inode order and
-# hardlink topology from the input. A single compressor worker then avoids
-# host-CPU-dependent fragment ordering.
+# hardlink topology from the input, which is what makes the compressor's worker
+# count irrelevant to the output: with -no-tailends and -no-hardlinks there are
+# no cross-file fragments left for workers to pack in a racy order. The worker
+# count therefore tracks the build's job count, and repro-check is what proves
+# it: its two legs deliberately run with different job counts, so a worker-count
+# dependency in this output fails the check rather than hiding in it.
+# Serializing it instead costs ~285 s of a ~500 s incremental build.
 (cd "$TREE" && tar --sort=name --format=gnu \
   --mtime="@$SOURCE_DATE_EPOCH" --owner=0 --group=0 --numeric-owner \
   --mode=g-s --hard-dereference -cf - .) | \
   env -u SOURCE_DATE_EPOCH mksquashfs - "$rootfs" -tar \
     -noappend -all-root -no-progress -exports -no-hardlinks -no-tailends \
-    -no-xattrs -processors 1 -comp zstd -mkfs-time "$SOURCE_DATE_EPOCH" \
+    -no-xattrs -processors "${JOBS:-1}" -comp zstd -mkfs-time "$SOURCE_DATE_EPOCH" \
     -all-time "$SOURCE_DATE_EPOCH" >/dev/null
 data_size=$(stat -c %s "$rootfs")
 data_size=$(( (data_size + 4095) / 4096 * 4096 ))

os/mkosi/tests/acceptance.sh

@@ -183,6 +183,11 @@ grep -q 'Archiving UKI image' "$D/../image/assemble.sh"
 grep -Fq "FLAVORS=\${FLAVORS:-prod}" "$D/build.sh"
 # A cache hit must never be able to answer the question repro-check asks.
 grep -Fq "if [[ \$action == repro-check ]]; then cache=0; fi" "$D/build.sh"
+# repro-check compares the release tarballs, so it must archive unconditionally
+# -- otherwise it compares absent files and reports no difference.
+grep -Fq "if [[ \$action == repro-check ]]; then archive=1; fi" "$D/build.sh"
+grep -Fq -e "--environment=\"DSTACK_TAR_RELEASE=\$archive\"" "$D/build.sh"
+grep -Fq "DSTACK_TAR_RELEASE=\"\${DSTACK_TAR_RELEASE:-1}\"" "$D/mkosi.postoutput"
 grep -Fq -e "--environment=\"DSTACK_COMPONENT_CACHE=\$cache\"" "$D/build.sh"
 grep -q 'write-source-manifest.py' "$D/build.sh"
 grep -q 'DSTACK_SOURCE_REVISION.*revision' "$D/build.sh"