From 534442e480090615f9a5b2bdee9a84a4a9718f73 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 03:04:33 +0000 Subject: [PATCH] experimental/air: warm snapshot cache + parallel gzip for plain_tar The plain_tar snapshot path (dirty working tree / no git ref) re-walked, re-tarred and re-gzipped the whole tree on every submission. Add a local warm cache keyed by (repo, config, include_paths) under $TMPDIR/databricks/.air/: an uncompressed snapshot.tar plus a manifest of each file's size+mtime and byte range. Later runs copy unchanged members verbatim from the warm tar and re-read only the changed files; gzip is parallelised with klauspost/pgzip (drop-in, ~18x faster on a 470 MiB tar). --no-cache bypasses the cache and re-packs from scratch. The cache engages only above 64 MiB, where a plain re-pack is slow. Local packaging latency only (file walk + tar + gzip; upload and API round trips are not measured and are unchanged by this PR): research (7.4k files): before 2456 ms -> warm hit 557 ms research+js+spark (27.8k/476MB): before 8525 ms -> warm hit 934 ms whole universe (534k files/~4GB): before ~2-3 min -> warm hit 6.6 s Co-authored-by: Isaac --- NOTICE | 4 + experimental/air/cmd/run.go | 4 +- experimental/air/cmd/runsubmit.go | 4 +- experimental/air/cmd/runsubmit_test.go | 24 +- experimental/air/cmd/snapshot_cache.go | 390 ++++++++++++++++++++ experimental/air/cmd/snapshot_cache_test.go | 183 +++++++++ experimental/air/cmd/snapshot_dabs.go | 26 +- experimental/air/cmd/snapshot_package.go | 39 +- go.mod | 2 + go.sum | 4 + 10 files changed, 651 insertions(+), 29 deletions(-) create mode 100644 experimental/air/cmd/snapshot_cache.go create mode 100644 experimental/air/cmd/snapshot_cache_test.go diff --git a/NOTICE b/NOTICE index 8ef4a7a1bd0..4b9c9991f20 100644 --- a/NOTICE +++ b/NOTICE @@ -131,6 +131,10 @@ jackc/pgx - https://github.com/jackc/pgx Copyright (c) 2013-2021 Jack Christensen License - https://github.com/jackc/pgx/blob/master/LICENSE +klauspost/pgzip - https://github.com/klauspost/pgzip +Copyright (c) 2014 Klaus Post +License - https://github.com/klauspost/pgzip/blob/master/LICENSE + charmbracelet/bubbles - https://github.com/charmbracelet/bubbles Copyright (c) 2020-2025 Charmbracelet, Inc License - https://github.com/charmbracelet/bubbles/blob/master/LICENSE diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 91b2e76065f..46f5debdb6e 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -35,6 +35,7 @@ func newRunCommand() *cobra.Command { overrides []string dryRun bool idempotencyKey string + noCache bool ) cmd := &cobra.Command{ @@ -78,6 +79,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so cmd.Flags().StringArrayVar(&overrides, "override", nil, "Override a YAML field, e.g. compute.num_accelerators=8 (repeatable)") cmd.Flags().BoolVar(&dryRun, "dry-run", false, "Validate the config without submitting") cmd.Flags().StringVar(&idempotencyKey, "idempotency-key", "", "Return the existing run if this key was already used") + cmd.Flags().BoolVar(&noCache, "no-cache", false, "Bypass the local snapshot cache and re-pack the code tarball from scratch") _ = cmd.MarkFlagRequired("file") // --dry-run only validates the config locally, so it needs no workspace. @@ -114,7 +116,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so } w := cmdctx.WorkspaceClient(ctx) - runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey, !jsonOut) + runID, dashboardURL, err := submitWorkload(ctx, w, cfg, file, idempotencyKey, !jsonOut, noCache) if err != nil { return err } diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 2456ef3eca4..d4179b0a5de 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -212,7 +212,7 @@ func withSpinner(ctx context.Context, show bool, msg string, fn func() error) er // upload the launch artifacts, assemble the Jobs payload, and submit it. It // returns the new run_id and its dashboard URL. showProgress enables the // stderr upload/packaging spinners (text mode only). -func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (int64, string, error) { +func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress, noCache bool) (int64, string, error) { // Compute the launch dir and command_path up front — a read-only workspace lookup plus a // local path build, no writes yet — so the pre-flight validates the real command_path. The // same path is reused for the upload and submit below, so the validated path is the submitted @@ -297,7 +297,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run // Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh. err = withSpinner(ctx, showProgress, "Packaging code snapshot…", func() error { var e error - snap, e = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir) + snap, e = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir, noCache) return e }) if err != nil { diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 4c0c68735d1..c14ccca74c3 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -211,7 +211,7 @@ func TestSubmitWorkload(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) + runID, dashboardURL, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false, false) require.NoError(t, err) assert.Equal(t, int64(777), runID) assert.Contains(t, dashboardURL, "/jobs/runs/777") @@ -255,7 +255,7 @@ func TestSubmitWorkloadHonorsOverride(t *testing.T) { cfg, err := loadRunConfigWithOverrides(t.Context(), cfgPath, []string{"compute.num_accelerators=4"}) require.NoError(t, err) - _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) + _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false, false) require.NoError(t, err) require.Len(t, got.Tasks, 1) @@ -298,7 +298,7 @@ code_source: // The DABs upload path logs via cmdio; the real `air run` context carries it. ctx := cmdio.MockDiscard(t.Context()) - _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false, false) require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask @@ -342,7 +342,7 @@ code_source: require.NoError(t, err) ctx := cmdio.MockDiscard(t.Context()) - _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false, false) require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask @@ -392,7 +392,7 @@ code_source: // The uploaded name carries a discriminator (timestamp), not the bare dir name. ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false) require.NoError(t, err) base := path.Base(snap.CodeSourcePath) assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be unique, not the bare dir name") @@ -444,9 +444,9 @@ code_source: ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false) require.NoError(t, err) - second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false) require.NoError(t, err) // Same pinned commit → identical content-addressed remote path, uploaded once @@ -488,7 +488,7 @@ code_source: ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase, false) require.NoError(t, err) assert.Empty(t, snap.GitStatePath) @@ -537,7 +537,7 @@ code_source: require.NoError(t, err) ctx := cmdio.MockDiscard(t.Context()) - _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem", false, false) require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask @@ -572,7 +572,7 @@ func TestSubmitWorkloadGuards(t *testing.T) { cfg := *base cfg.UsagePolicyName = new("nope") - _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false) + _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false, false) require.ErrorContains(t, err, `no usage policy named "nope"`) for _, p := range paths { assert.NotContains(t, p, "/workspace/", "no workspace write may precede policy resolution") @@ -609,7 +609,7 @@ func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false) + _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false, false) require.NoError(t, err) assert.Equal(t, policyID, got.BudgetPolicyId) }) @@ -620,7 +620,7 @@ func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false) + _, _, err = submitWorkload(cmdio.MockDiscard(t.Context()), w, cfg, cfgPath, "idem", false, false) require.NoError(t, err) assert.Equal(t, policyID, got.BudgetPolicyId) }) diff --git a/experimental/air/cmd/snapshot_cache.go b/experimental/air/cmd/snapshot_cache.go new file mode 100644 index 00000000000..82ce0c64ad7 --- /dev/null +++ b/experimental/air/cmd/snapshot_cache.go @@ -0,0 +1,390 @@ +package aircmd + +// Warm snapshot cache for the plain_tar path (working tree, no git ref). The Python +// CLI and the git_archive path re-pack the whole tree every run; for a large repo the +// file walk + read + gzip dominates submit latency. This keeps a warm, uncompressed +// tar of the tree on local disk plus a manifest of each member's identity (size+mtime) +// and byte range. On the next run we stat the file set, and rebuild the tarball by +// copying unchanged members verbatim from the warm tar — reading only changed files +// from disk — before gzipping the upload. Nothing changed is the degenerate case: we +// just recompress the warm tar. The cache is keyed by (repo path, config path, +// include_paths), so distinct repos or configs never share an entry, and it is gated +// to large trees (below the threshold a plain re-pack is cheap enough not to bother). + +import ( + "archive/tar" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "fmt" + "io" + "os" + "path" + "path/filepath" + "slices" + "strings" + "time" + + "github.com/databricks/cli/libs/log" + "github.com/klauspost/pgzip" +) + +const ( + // snapshotCacheVersion invalidates on-disk caches when the layout below changes. + snapshotCacheVersion = "v1" + snapshotCacheManifestName = "manifest.json" + snapshotCacheTarName = "snapshot.tar" + + // snapshotCacheMinBytes gates the cache to trees large enough that the walk+read + // cost dominates. Below it a plain re-pack is cheap and the cache bookkeeping + // isn't worth it. Rough heuristic, not exact — tune from measurement. + snapshotCacheMinBytes = 64 << 20 // 64 MiB +) + +// tarTrailer is the two zero blocks that mark end-of-archive. writeSnapshot appends +// it because members are streamed without the tar.Writer's own Close (which would +// embed a trailer between members). +var tarTrailer = make([]byte, 2*512) + +// cacheEntry records a member's identity for change detection (size+mtime) and its +// byte range within the warm snapshot.tar, so an unchanged member can be copied +// verbatim instead of re-read from disk. +type cacheEntry struct { + Size int64 `json:"size"` + ModTime int64 `json:"mtime_ns"` + Offset int64 `json:"offset"` + Length int64 `json:"length"` +} + +// snapshotManifest is the on-disk index of a warm snapshot.tar. +type snapshotManifest struct { + Version string `json:"version"` + DirName string `json:"dir_name"` + Entries map[string]cacheEntry `json:"entries"` // keyed by slash-separated relative path +} + +// snapshotCacheKey is a stable digest of the inputs that determine the tar's content +// set. Different repos, config files, or include_paths get different cache folders. +func snapshotCacheKey(absRepo, absConfig string, includePaths []string) string { + paths := slices.Clone(includePaths) + slices.Sort(paths) + material := strings.Join(append([]string{absRepo, absConfig, snapshotCacheVersion}, paths...), "\x00") + sum := sha256.Sum256([]byte(material)) + return hex.EncodeToString(sum[:]) +} + +func snapshotCacheDir(absRepo, absConfig string, includePaths []string) string { + return filepath.Join(os.TempDir(), "databricks", ".air", snapshotCacheKey(absRepo, absConfig, includePaths)) +} + +// packagePlainTarWithCache writes the working-tree tarball to outputTarball, using the +// warm cache when the tree is large enough. Small trees are packed fresh without +// touching the cache. +func packagePlainTarWithCache(ctx context.Context, repoPath, configPath string, includePaths []string, isGitRepo bool, outputTarball string) (err error) { + start := time.Now() + files, err := snapshotFiles(ctx, repoPath, includePaths, isGitRepo) + if err != nil { + return err + } + listDone := time.Now() + var total int64 + for _, f := range files { + total += f.size + } + dirName := filepath.Base(repoPath) + + // One timing line comparable to the shell path's "snapshot profile", so a + // cache-on vs --no-cache run can be compared directly. Debug-only. + mode := "rebuild-cold" + defer func() { + log.Debugf(ctx, "air snapshot cache: mode=%s files=%d uncompressed_bytes=%d list=%s pack=%s", + mode, len(files), total, listDone.Sub(start), time.Since(listDone)) + }() + + if total < snapshotCacheMinBytes { + mode = "skip-small" + return writeGzOnly(repoPath, dirName, files, outputTarball) + } + + absRepo, err := filepath.Abs(repoPath) + if err != nil { + return err + } + absConfig, err := filepath.Abs(configPath) + if err != nil { + return err + } + cacheDir := snapshotCacheDir(absRepo, absConfig, includePaths) + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + manifestPath := filepath.Join(cacheDir, snapshotCacheManifestName) + + old := loadSnapshotManifest(manifestPath) + if old != nil && old.DirName == dirName && fileExists(tarPath) && !snapshotChanged(files, old) { + mode = "hit-nochange" + return gzipFile(tarPath, outputTarball) + } + + if err := os.MkdirAll(cacheDir, 0o700); err != nil { + return fmt.Errorf("failed to create snapshot cache dir: %w", err) + } + if old != nil { + mode = "rebuild-warm" + } + return rebuildWarmSnapshot(repoPath, dirName, files, old, tarPath, outputTarball) +} + +// snapshotChanged reports whether the current file set differs from the manifest by +// any add, delete, or modification (mtime+size). Equal length plus every current file +// matching an entry means the sets are identical. +func snapshotChanged(files []snapshotFile, old *snapshotManifest) bool { + if len(files) != len(old.Entries) { + return true + } + for _, f := range files { + e, ok := old.Entries[filepath.ToSlash(f.rel)] + if !ok || e.Size != f.size || e.ModTime != f.modTime { + return true + } + } + return false +} + +// rebuildWarmSnapshot writes a fresh warm tar (copying unchanged members from the old +// one when available) and its gzipped upload copy in a single pass, then atomically +// replaces the cached tar and manifest. +func rebuildWarmSnapshot(repoPath, dirName string, files []snapshotFile, old *snapshotManifest, tarPath, outputTarball string) (err error) { + gz, closeGz, err := newGzFile(outputTarball) + if err != nil { + return err + } + defer func() { err = firstErr(err, closeGz()) }() + + newTarPath := tarPath + ".tmp" + tarFile, err := os.Create(newTarPath) + if err != nil { + return fmt.Errorf("failed to create warm tar: %w", err) + } + defer func() { + if err != nil { + os.Remove(newTarPath) + } + }() + + var oldTar io.ReaderAt + if old != nil { + if f, e := os.Open(tarPath); e == nil { + defer f.Close() + oldTar = f + } else { + old = nil // warm tar gone; rebuild every member from disk + } + } + + manifest, err := writeSnapshot(repoPath, dirName, files, old, oldTar, tarFile, gz) + if err != nil { + tarFile.Close() + return err + } + if err = tarFile.Close(); err != nil { + return fmt.Errorf("failed to finalize warm tar: %w", err) + } + if err = closeGz(); err != nil { + return err + } + if err = os.Rename(newTarPath, tarPath); err != nil { + return fmt.Errorf("failed to install warm tar: %w", err) + } + return saveSnapshotManifest(filepath.Join(filepath.Dir(tarPath), snapshotCacheManifestName), manifest) +} + +// writeGzOnly packs files straight to a gzipped tarball without persisting a cache, +// used for trees below the cache threshold. +func writeGzOnly(repoPath, dirName string, files []snapshotFile, outputTarball string) (err error) { + gz, closeGz, err := newGzFile(outputTarball) + if err != nil { + return err + } + defer func() { err = firstErr(err, closeGz()) }() + _, err = writeSnapshot(repoPath, dirName, files, nil, nil, nil, gz) + return err +} + +// writeSnapshot streams every member (sorted for determinism) to gzDst, and to tarDst +// too when non-nil, returning the manifest that indexes each member's byte range in +// the tarDst stream. When old+oldTar are set, an unchanged member (matching size+mtime) +// is copied verbatim from oldTar rather than re-read from disk. +func writeSnapshot(repoPath, dirName string, files []snapshotFile, old *snapshotManifest, oldTar io.ReaderAt, tarDst, gzDst io.Writer) (snapshotManifest, error) { + dst := gzDst + if tarDst != nil { + dst = io.MultiWriter(tarDst, gzDst) + } + cw := &countWriter{w: dst} + + sorted := slices.Clone(files) + slices.SortFunc(sorted, func(a, b snapshotFile) int { + return strings.Compare(filepath.ToSlash(a.rel), filepath.ToSlash(b.rel)) + }) + + entries := make(map[string]cacheEntry, len(sorted)) + for _, f := range sorted { + rel := filepath.ToSlash(f.rel) + start := cw.n + + reused := false + if old != nil && oldTar != nil { + if e, ok := old.Entries[rel]; ok && e.Size == f.size && e.ModTime == f.modTime { + if _, err := io.Copy(cw, io.NewSectionReader(oldTar, e.Offset, e.Length)); err != nil { + return snapshotManifest{}, fmt.Errorf("failed to copy cached member %q: %w", rel, err) + } + reused = true + } + } + if !reused { + if err := writeMember(cw, repoPath, path.Join(dirName, rel), f.rel); err != nil { + return snapshotManifest{}, err + } + } + entries[rel] = cacheEntry{Size: f.size, ModTime: f.modTime, Offset: start, Length: cw.n - start} + } + if _, err := cw.Write(tarTrailer); err != nil { + return snapshotManifest{}, err + } + return snapshotManifest{Version: snapshotCacheVersion, DirName: dirName, Entries: entries}, nil +} + +// writeMember streams one framed tar member (header + content + block padding) for the +// file at repoPath/rel, named name inside the archive. It deliberately Flushes rather +// than Closes the tar.Writer, so no end-of-archive trailer is written between members. +func writeMember(w io.Writer, repoPath, name, rel string) error { + full := filepath.Join(repoPath, filepath.FromSlash(rel)) + info, err := os.Lstat(full) + if err != nil { + return fmt.Errorf("failed to stat %q: %w", rel, err) + } + link := "" + if info.Mode()&os.ModeSymlink != 0 { + if link, err = os.Readlink(full); err != nil { + return fmt.Errorf("failed to read symlink %q: %w", rel, err) + } + } + hdr, err := tar.FileInfoHeader(info, link) + if err != nil { + return fmt.Errorf("failed to build tar header for %q: %w", rel, err) + } + hdr.Name = name + + tw := tar.NewWriter(w) + if err := tw.WriteHeader(hdr); err != nil { + return fmt.Errorf("failed to write tar header for %q: %w", rel, err) + } + if info.Mode().IsRegular() { + f, err := os.Open(full) + if err != nil { + return fmt.Errorf("failed to open %q: %w", rel, err) + } + defer f.Close() + if _, err := io.Copy(tw, f); err != nil { + return fmt.Errorf("failed to archive %q: %w", rel, err) + } + } + return tw.Flush() +} + +// countWriter counts the bytes written through it, to record member byte offsets. +type countWriter struct { + w io.Writer + n int64 +} + +func (c *countWriter) Write(p []byte) (int, error) { + n, err := c.w.Write(p) + c.n += int64(n) + return n, err +} + +// newGzFile creates outputTarball and returns a BestSpeed gzip writer over it plus an +// idempotent close that flushes the gzip stream and the file. Compression level is +// BestSpeed because the uploaded size does not matter for this workflow — only latency. +// +// gzip is parallel (klauspost/pgzip): compressing the whole tar is the dominant +// packaging cost on a large tree and is paid on every run — even a no-change cache hit +// re-gzips the warm tar — so it is spread across cores (measured ~18x faster than +// compress/gzip on a 470 MiB tar). pgzip buffers its own blocks, so the tar writer's +// small writes parallelize fine without extra buffering. Its output is an ordinary gzip +// stream any gunzip/tar reads, and it falls back to serial below one block — fine, since +// the cache only engages above snapshotCacheMinBytes. +func newGzFile(outputTarball string) (io.Writer, func() error, error) { + f, err := os.Create(outputTarball) + if err != nil { + return nil, nil, fmt.Errorf("failed to create tarball: %w", err) + } + gz, err := pgzip.NewWriterLevel(f, pgzip.BestSpeed) + if err != nil { + f.Close() + return nil, nil, err + } + closed := false + closeFn := func() error { + if closed { + return nil + } + closed = true + return firstErr(gz.Close(), f.Close()) + } + return gz, closeFn, nil +} + +// gzipFile writes a BestSpeed gzip of src to outputTarball. Used on a no-change cache +// hit to recompress the warm tar without re-reading the working tree. +func gzipFile(src, outputTarball string) (err error) { + in, err := os.Open(src) + if err != nil { + return fmt.Errorf("failed to open warm tar: %w", err) + } + defer in.Close() + gz, closeGz, err := newGzFile(outputTarball) + if err != nil { + return err + } + defer func() { err = firstErr(err, closeGz()) }() + _, err = io.Copy(gz, in) + return err +} + +func loadSnapshotManifest(path string) *snapshotManifest { + data, err := os.ReadFile(path) + if err != nil { + return nil + } + var m snapshotManifest + if err := json.Unmarshal(data, &m); err != nil || m.Version != snapshotCacheVersion { + return nil + } + return &m +} + +func saveSnapshotManifest(path string, m snapshotManifest) error { + data, err := json.Marshal(m) + if err != nil { + return err + } + if err := os.WriteFile(path, data, 0o600); err != nil { + return fmt.Errorf("failed to write snapshot manifest: %w", err) + } + return nil +} + +func fileExists(path string) bool { + _, err := os.Stat(path) + return err == nil +} + +func firstErr(errs ...error) error { + for _, e := range errs { + if e != nil { + return e + } + } + return nil +} diff --git a/experimental/air/cmd/snapshot_cache_test.go b/experimental/air/cmd/snapshot_cache_test.go new file mode 100644 index 00000000000..134f3418c0e --- /dev/null +++ b/experimental/air/cmd/snapshot_cache_test.go @@ -0,0 +1,183 @@ +package aircmd + +import ( + "archive/tar" + "compress/gzip" + "os" + "path/filepath" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// extractTarball returns entry name -> content for a .tar.gz. +func extractTarball(t *testing.T, path string) map[string]string { + t.Helper() + f, err := os.Open(path) + require.NoError(t, err) + defer f.Close() + gz, err := gzip.NewReader(f) + require.NoError(t, err) + defer gz.Close() + + out := map[string]string{} + tr := tar.NewReader(gz) + for { + hdr, err := tr.Next() + if err != nil { + break + } + buf := make([]byte, hdr.Size) + _, _ = tr.Read(buf) + out[hdr.Name] = string(buf) + } + return out +} + +func statFiles(t *testing.T, repo string, rels ...string) []snapshotFile { + t.Helper() + var files []snapshotFile + for _, rel := range rels { + info, err := os.Lstat(filepath.Join(repo, filepath.FromSlash(rel))) + require.NoError(t, err) + files = append(files, snapshotFile{rel: filepath.FromSlash(rel), size: info.Size(), modTime: info.ModTime().UnixNano()}) + } + return files +} + +func TestSnapshotCacheKey(t *testing.T) { + base := snapshotCacheKey("/repo", "/repo/air.yaml", nil) + assert.Equal(t, base, snapshotCacheKey("/repo", "/repo/air.yaml", nil), "stable for identical inputs") + assert.NotEqual(t, base, snapshotCacheKey("/other", "/repo/air.yaml", nil), "repo path matters") + assert.NotEqual(t, base, snapshotCacheKey("/repo", "/repo/other.yaml", nil), "config path matters") + assert.NotEqual(t, base, snapshotCacheKey("/repo", "/repo/air.yaml", []string{"src"}), "include_paths matter") + // include_paths order must not matter. + assert.Equal(t, + snapshotCacheKey("/repo", "/repo/air.yaml", []string{"a", "b"}), + snapshotCacheKey("/repo", "/repo/air.yaml", []string{"b", "a"})) +} + +func TestSnapshotChanged(t *testing.T) { + old := &snapshotManifest{Entries: map[string]cacheEntry{ + "a.txt": {Size: 1, ModTime: 100}, + "src/b.py": {Size: 2, ModTime: 200}, + }} + unchanged := []snapshotFile{{rel: "a.txt", size: 1, modTime: 100}, {rel: filepath.FromSlash("src/b.py"), size: 2, modTime: 200}} + assert.False(t, snapshotChanged(unchanged, old)) + + modified := []snapshotFile{{rel: "a.txt", size: 1, modTime: 100}, {rel: filepath.FromSlash("src/b.py"), size: 2, modTime: 999}} + assert.True(t, snapshotChanged(modified, old), "mtime change detected") + + resized := []snapshotFile{{rel: "a.txt", size: 5, modTime: 100}, {rel: filepath.FromSlash("src/b.py"), size: 2, modTime: 200}} + assert.True(t, snapshotChanged(resized, old), "size change detected") + + removed := []snapshotFile{{rel: "a.txt", size: 1, modTime: 100}} + assert.True(t, snapshotChanged(removed, old), "deletion detected") + + added := append(append([]snapshotFile(nil), unchanged...), snapshotFile{rel: "c.txt", size: 3, modTime: 300}) + assert.True(t, snapshotChanged(added, old), "addition detected") +} + +func TestWarmSnapshotColdBuild(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + writeRepoFile(t, repo, "src/model.py", "print()") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + out := filepath.Join(t.TempDir(), "snap.tar.gz") + + files := statFiles(t, repo, "a.txt", "src/model.py") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, out)) + + contents := extractTarball(t, out) + assert.Equal(t, "alpha", contents[dirName+"/a.txt"]) + assert.Equal(t, "print()", contents[dirName+"/src/model.py"]) + + // The warm tar and manifest are persisted for the next run. + assert.FileExists(t, tarPath) + m := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, m) + assert.Equal(t, dirName, m.DirName) + assert.Len(t, m.Entries, 2) +} + +func TestWarmSnapshotReusesUnchangedFromCache(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + writeRepoFile(t, repo, "keep.py", "keep") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + files := statFiles(t, repo, "a.txt", "keep.py") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, filepath.Join(t.TempDir(), "cold.tar.gz"))) + + old := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, old) + + // Delete keep.py from disk but keep it in the file list with its original + // size+mtime: a correct rebuild must copy its bytes from the warm tar, proving + // unchanged members are not re-read from disk. + require.NoError(t, os.Remove(filepath.Join(repo, "keep.py"))) + + out := filepath.Join(t.TempDir(), "warm.tar.gz") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, old, tarPath, out)) + + contents := extractTarball(t, out) + assert.Equal(t, "keep", contents[dirName+"/keep.py"], "unchanged member copied from warm tar") + assert.Equal(t, "alpha", contents[dirName+"/a.txt"]) +} + +func TestWarmSnapshotRebuildEditAddDelete(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + writeRepoFile(t, repo, "b.txt", "bravo") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + files := statFiles(t, repo, "a.txt", "b.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, filepath.Join(t.TempDir(), "cold.tar.gz"))) + old := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, old) + + // Edit a.txt (changed), delete b.txt, add c.txt. + writeRepoFile(t, repo, "a.txt", "alpha-v2") + require.NoError(t, os.Remove(filepath.Join(repo, "b.txt"))) + writeRepoFile(t, repo, "c.txt", "charlie") + + out := filepath.Join(t.TempDir(), "warm.tar.gz") + newFiles := statFiles(t, repo, "a.txt", "c.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, newFiles, old, tarPath, out)) + + contents := extractTarball(t, out) + assert.Equal(t, "alpha-v2", contents[dirName+"/a.txt"], "edited file updated") + assert.Equal(t, "charlie", contents[dirName+"/c.txt"], "added file present") + _, hasB := contents[dirName+"/b.txt"] + assert.False(t, hasB, "deleted file dropped") + + // The refreshed manifest reflects the new set. + updated := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, updated) + assert.Len(t, updated.Entries, 2) +} + +func TestGzipFileRoundTrip(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "alpha") + dirName := filepath.Base(repo) + + cacheDir := t.TempDir() + tarPath := filepath.Join(cacheDir, snapshotCacheTarName) + files := statFiles(t, repo, "a.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, tarPath, filepath.Join(t.TempDir(), "cold.tar.gz"))) + + // gzipFile recompresses the warm tar directly (the no-change hit path). + out := filepath.Join(t.TempDir(), "reuse.tar.gz") + require.NoError(t, gzipFile(tarPath, out)) + contents := extractTarball(t, out) + assert.Equal(t, "alpha", contents[dirName+"/a.txt"]) +} diff --git a/experimental/air/cmd/snapshot_dabs.go b/experimental/air/cmd/snapshot_dabs.go index ed809b09e01..2176bf2cce1 100644 --- a/experimental/air/cmd/snapshot_dabs.go +++ b/experimental/air/cmd/snapshot_dabs.go @@ -38,7 +38,7 @@ const uploadProvenanceSidecars = false // not reimplement workspace/volume upload. A minimal in-memory bundle carries the // local tarball path as code_source_path; ReplaceWithRemotePath rewrites it to the // artifact .internal path and Upload pushes the bytes. -func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath string, sidecarStore filer.Filer, sidecarBase string) (snapshotResult, error) { +func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath string, sidecarStore filer.Filer, sidecarBase string, noCache bool) (snapshotResult, error) { repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) if err != nil { return snapshotResult{}, err @@ -57,7 +57,7 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s if snap.RemoteVolume != nil { remoteVolume = *snap.RemoteVolume } - result, err := uploadSnapshotViaDABs(ctx, w, repoPath, plan, remoteVolume) + result, err := uploadSnapshotViaDABs(ctx, w, repoPath, configPath, plan, remoteVolume, noCache) if err != nil { return snapshotResult{}, err } @@ -141,13 +141,17 @@ func snapshotTarballName(plan snapshotPlan, dirName string) string { } // packageSnapshot writes the snapshot to tarball per the resolved plan: `git archive` -// of the pinned commit for git_archive, else a plain tar of the working tree. -func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, tarball string) error { - dirName := filepath.Base(repoPath) +// of the pinned commit for git_archive, else a plain tar of the working tree. The +// working-tree path uses the warm snapshot cache unless noCache is set, in which case +// it falls back to the shell `tar` path. +func packageSnapshot(ctx context.Context, repoPath, configPath string, plan snapshotPlan, tarball string, noCache bool) error { if plan.mode == modeGitArchive { - return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths) + return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, filepath.Base(repoPath), plan.includePaths) } - return createPlainTarball(ctx, repoPath, tarball, plan.includePaths, plan.isGitRepo) + if noCache { + return createPlainTarball(ctx, repoPath, tarball, plan.includePaths, plan.isGitRepo) + } + return packagePlainTarWithCache(ctx, repoPath, configPath, plan.includePaths, plan.isGitRepo, tarball) } // uploadSnapshotViaDABs uploads the snapshot through DABs' artifact-upload machinery @@ -159,7 +163,7 @@ func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, ta // git_archive snapshots are cacheable: the tarball name is content-addressed by // (commit, include_paths), so if the identical object is already uploaded we skip // packaging and upload entirely and just reuse the remote path. -func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath string, plan snapshotPlan, remoteVolume string) (snapshotResult, error) { +func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath, configPath string, plan snapshotPlan, remoteVolume string, noCache bool) (snapshotResult, error) { // artifactPath is where DABs uploads the tarball; GetFilerForLibraries routes to // a Workspace or Volume filer based on its prefix, then appends /.internal. artifactPath := remoteVolume @@ -236,7 +240,7 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r } // Cache miss (or plain_tar): package the tarball locally, then upload the bytes. - if err := packageSnapshot(ctx, repoPath, plan, filepath.Join(tmp, tarName)); err != nil { + if err := packageSnapshot(ctx, repoPath, configPath, plan, filepath.Join(tmp, tarName), noCache); err != nil { return snapshotResult{}, err } @@ -244,9 +248,13 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r if diags.HasError() { return snapshotResult{}, diags.Error() } + // Phase 0 profiling: isolate the upload stage from packaging so we know whether + // the cost is local (walk/gzip) or network. Debug-only; enable with -v/--debug. + uploadStart := time.Now() if diags := bundle.Apply(ctx, b, libraries.Upload(libs)); diags.HasError() { return snapshotResult{}, diags.Error() } + log.Debugf(ctx, "air snapshot profile: upload=%s", time.Since(uploadStart)) remote, err := readCodeSourcePath(b) if err != nil { diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index c22f7dd61af..8ff47ed4f72 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -9,6 +9,9 @@ import ( "os/exec" "path/filepath" "strings" + "time" + + "github.com/databricks/cli/libs/log" ) // Tar builders ported from cli/utils/snapshot.py. Both shell out (git archive / tar) @@ -58,33 +61,59 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc } outName := filepath.Base(outputTarball) + listStart := time.Now() files, err := snapshotFiles(ctx, repoPath, includePaths, isGitRepo) if err != nil { return err } + listDur := time.Since(listStart) + args := []string{"-czf", outName, "-C", parent, "--null", "--no-recursion", "-T", "-"} cmd := exec.CommandContext(ctx, "tar", args...) // Run tar in the output directory so the bare -f basename lands there. cmd.Dir = outDirAbs var stdin bytes.Buffer + var uncompressedBytes int64 for _, file := range files { - stdin.WriteString(filepath.ToSlash(filepath.Join(dirName, file))) + stdin.WriteString(filepath.ToSlash(filepath.Join(dirName, file.rel))) stdin.WriteByte(0) + uncompressedBytes += file.size } cmd.Stdin = &stdin var stderr bytes.Buffer cmd.Stderr = &stderr + packStart := time.Now() if err := cmd.Run(); err != nil { if msg := strings.TrimSpace(stderr.String()); msg != "" { return fmt.Errorf("failed to create plain tarball: %w: %s", err, msg) } return fmt.Errorf("failed to create plain tarball: %w", err) } + + // Phase 0 profiling: decompose plain_tar cost into file walk (list) vs tar+gzip + // (pack_gzip) so we can judge whether a warm-tar cache would pay off before + // building one. Debug-only; enable with -v/--debug. + var compressedBytes int64 = -1 + if fi, statErr := os.Stat(outputTarball); statErr == nil { + compressedBytes = fi.Size() + } + log.Debugf(ctx, "air snapshot profile: mode=plain_tar files=%d uncompressed_bytes=%d compressed_bytes=%d list=%s pack_gzip=%s", + len(files), uncompressedBytes, compressedBytes, listDur, time.Since(packStart)) return nil } -func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, isGitRepo bool) ([]string, error) { +// snapshotFile is a file selected for the snapshot: its repo-relative path (native +// separators) with the size and mtime the warm cache uses to detect changes. +type snapshotFile struct { + rel string + size int64 + modTime int64 // Unix nanoseconds +} + +// snapshotFiles returns the files to archive (git-tracked and untracked, honoring +// .gitignore, minus .git and AppleDouble files), each with its size and mtime. +func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, isGitRepo bool) ([]snapshotFile, error) { args := []string{"-C", repoPath, "ls-files", "-z", "--cached", "--others", "--exclude-standard"} if !isGitRepo { gitDir, err := os.MkdirTemp("", "air-snapshot-git-") @@ -108,7 +137,7 @@ func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, return nil, fmt.Errorf("failed to evaluate git ignore rules: %w", err) } - var files []string + var files []snapshotFile for raw := range bytes.SplitSeq(output, []byte{0}) { if len(raw) == 0 { continue @@ -118,14 +147,14 @@ func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, if name == ".git" || strings.HasPrefix(name, ".git/") || strings.HasPrefix(base, "._") { continue } - _, err := os.Lstat(filepath.Join(repoPath, filepath.FromSlash(name))) + info, err := os.Lstat(filepath.Join(repoPath, filepath.FromSlash(name))) if errors.Is(err, os.ErrNotExist) { continue } if err != nil { return nil, fmt.Errorf("failed to inspect snapshot path %q: %w", name, err) } - files = append(files, filepath.FromSlash(name)) + files = append(files, snapshotFile{rel: filepath.FromSlash(name), size: info.Size(), modTime: info.ModTime().UnixNano()}) } return files, nil } diff --git a/go.mod b/go.mod index f8eee11b2b5..aa9a310befb 100644 --- a/go.mod +++ b/go.mod @@ -23,6 +23,7 @@ require ( github.com/hashicorp/terraform-json v0.28.0 // MPL-2.0 github.com/hexops/gotextdiff v1.0.3 // BSD-3-Clause github.com/jackc/pgx/v5 v5.10.0 // MIT + github.com/klauspost/pgzip v1.2.6 // MIT github.com/mattn/go-isatty v0.0.24 // MIT github.com/muesli/termenv v0.16.0 // MIT github.com/palantir/pkg/yamlpatch v1.5.0 // BSD-3-Clause @@ -87,6 +88,7 @@ require ( github.com/inconshreveable/mousetrap v1.1.0 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect + github.com/klauspost/compress v1.19.2 // indirect github.com/lucasb-eyer/go-colorful v1.4.0 // indirect github.com/mattn/go-colorable v0.1.14 // indirect github.com/mattn/go-localereader v0.0.1 // indirect diff --git a/go.sum b/go.sum index b78e7271bba..f2dc83afaa7 100644 --- a/go.sum +++ b/go.sum @@ -156,6 +156,10 @@ github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 h1:BQSFePA1RWJOl github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99/go.mod h1:1lJo3i6rXxKeerYnT8Nvf0QmHCRC1n8sfWVwXF2Frvo= github.com/kevinburke/ssh_config v1.2.0 h1:x584FjTGwHzMwvHx18PXxbBVzfnxogHaAReU4gf13a4= github.com/kevinburke/ssh_config v1.2.0/go.mod h1:CT57kijsi8u/K/BOFA39wgDQJ9CxiF4nAY/ojJ6r6mM= +github.com/klauspost/compress v1.19.2 h1:hMRETovs/pu/dVWN7zIT1PGG8t509MwT6bO7XSi26R8= +github.com/klauspost/compress v1.19.2/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/klauspost/pgzip v1.2.6 h1:8RXeL5crjEUFnR2/Sn6GJNWtSQ3Dk8pq4CL3jvdDyjU= +github.com/klauspost/pgzip v1.2.6/go.mod h1:Ch1tH69qFZu15pkjo5kYi6mth2Zzwzt50oCQKQE9RUs= github.com/kr/pretty v0.3.1 h1:flRD4NNwYAUpkphVc1HcthR4KEIFJ65n8Mw5qdRn3LE= github.com/kr/pretty v0.3.1/go.mod h1:hoEshYVHaxMs3cyo3Yncou5ZscifuDolrwPKZanG3xk= github.com/kr/text v0.2.0 h1:5Nx0Ya0ZqY2ygV366QzturHI13Jq95ApcVaJBhpS+AY=