From 2945d4588fb2b2da8183cf4313b6bfc061636bab Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 22:42:04 +0000 Subject: [PATCH 1/7] experimental/air: content-address the plain_tar upload git_archive snapshots are already content-addressed: a repeat submission at the same commit reuses the uploaded tarball and skips packaging + upload. plain_tar (dirty working tree) used a timestamped name, so it re-packaged and re-uploaded the full tarball on every submission, even when nothing changed. Name the plain_tar tarball by a working-tree fingerprint (sha256 over each file's path, size and mtime) and run the same snapshotExists skip for both modes. An unchanged resubmit now reuses the remote object and moves no bytes. The listing is captured once and threaded into packaging, so the tree is walked only once. The fingerprint is size+mtime, not content, matching DABs file-sync. Verified on df1: a second submission of an unchanged tree logs "snapshot upload skipped; reusing ..." and returns the identical remote path. Co-authored-by: Isaac --- experimental/air/cmd/runsubmit_test.go | 35 +++++-- experimental/air/cmd/snapshot_cachekey.go | 25 +++++ .../air/cmd/snapshot_cachekey_test.go | 23 +++++ experimental/air/cmd/snapshot_dabs.go | 98 ++++++++++--------- experimental/air/cmd/snapshot_package.go | 46 +++++---- experimental/air/cmd/snapshot_package_test.go | 20 +++- 6 files changed, 168 insertions(+), 79 deletions(-) diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index cf781c42f59..92dc69af481 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -555,15 +555,27 @@ func testSidecarStore(t *testing.T, w *databricks.WorkspaceClient) (filer.Filer, const testSnapshotArtifactPath = "/Workspace/Users/tester@databricks.com/.air/repo_snapshots" -// A plain-tar (working-tree) snapshot is uploaded under a unique, timestamped name so -// two concurrent submissions of the same root_path don't clobber each other's upload. -func TestSubmitWorkloadPlainTarNameIsUnique(t *testing.T) { +// A plain-tar (working-tree) snapshot is content-addressed by its file fingerprint +// (path+size+mtime): submitting the same unchanged tree twice reuses the already-uploaded +// tarball and skips the second upload, resolving to the identical remote path. +func TestSubmitWorkloadPlainTarContentAddressed(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { return jobs.SubmitRunResponse{RunId: 555} }) + // Track snapshot uploads, preserving fake-workspace persistence so the second + // submit's existence Stat sees the first upload. Dedupe by path: the DABs uploader + // mkdirs-and-retries the import, so one logical upload can hit this route twice. + uploaded := map[string]bool{} + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { + p := req.Vars["path"] + if strings.Contains(p, "/.air/repo_snapshots/") { + uploaded[p] = true + } + return req.Workspace.WorkspaceFilesImportFile(p, req.Body, req.URL.Query().Get("overwrite") == "true") + }) stubValidateConfig(server) testserver.AddDefaultHandlers(server) w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) @@ -584,14 +596,21 @@ code_source: loaded, err := loadRunConfig(cfgPath) require.NoError(t, err) - // 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, testSnapshotArtifactPath, sidecarStore, sidecarBase) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) + require.NoError(t, err) + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) 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") - assert.Regexp(t, `^src_\d{8}_\d{6}\.tar\.gz$`, base) + + // Content-addressed name: not the bare dir, but a 16-hex-char fingerprint. + base := path.Base(first.CodeSourcePath) + assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be content-addressed, not the bare dir name") + assert.Regexp(t, `^src_[0-9a-f]{16}\.tar\.gz$`, base) + + // Same unchanged tree → identical remote path, uploaded once (second submit is a hit). + assert.Equal(t, first.CodeSourcePath, second.CodeSourcePath) + assert.Len(t, uploaded, 1, "unchanged plain_tar should skip the second upload") } // A git_archive snapshot is content-addressed by (commit, include_paths): submitting diff --git a/experimental/air/cmd/snapshot_cachekey.go b/experimental/air/cmd/snapshot_cachekey.go index 44c5cb8f351..fd680b880a2 100644 --- a/experimental/air/cmd/snapshot_cachekey.go +++ b/experimental/air/cmd/snapshot_cachekey.go @@ -7,6 +7,8 @@ package aircmd import ( "crypto/sha256" "encoding/hex" + "fmt" + "path/filepath" "slices" "strings" ) @@ -14,6 +16,29 @@ import ( // snapshotPackagingVersion is bumped when packaging logic changes in a way that invalidates existing caches const snapshotPackagingVersion = "v1" +// plainTarKeyVersion namespaces the plain_tar working-tree key (so it can never collide +// with a git_archive key) and lets us invalidate it if the fingerprint scheme changes. +const plainTarKeyVersion = "plaintar-v1" + +// computePlainTarKey returns a content-addressed key for a working-tree snapshot: the +// SHA-256 over every file's path, size and mtime (sorted for stability). An unchanged +// tree yields the same key, so an already-uploaded tarball can be reused instead of +// re-packaged and re-uploaded. The fingerprint is size+mtime, not content — the same +// trade-off DABs file-sync makes — so an edit preserving both size and mtime is not seen. +func computePlainTarKey(files []snapshotFile) string { + sorted := slices.Clone(files) + slices.SortFunc(sorted, func(a, b snapshotFile) int { + return strings.Compare(a.rel, b.rel) + }) + + h := sha256.New() + for _, f := range sorted { + fmt.Fprintf(h, "%s\x00%d\x00%d\n", filepath.ToSlash(f.rel), f.size, f.modTime) + } + fmt.Fprint(h, plainTarKeyVersion) + return hex.EncodeToString(h.Sum(nil)) +} + // computeSnapshotCacheKey returns a stable cache key for a snapshot tarball: the // SHA-256 digest of (commitSHA, normalized includePaths, snapshotPackagingVersion, // subtreePrefix). Changing any input yields a different entry. An empty subtree diff --git a/experimental/air/cmd/snapshot_cachekey_test.go b/experimental/air/cmd/snapshot_cachekey_test.go index 2f933aa40bd..2607bc8dd89 100644 --- a/experimental/air/cmd/snapshot_cachekey_test.go +++ b/experimental/air/cmd/snapshot_cachekey_test.go @@ -70,3 +70,26 @@ func TestComputeSnapshotCacheKeyProperties(t *testing.T) { // The version constant participates: a different version is a different key. assert.NotEqual(t, snapshotPackagingVersion, "") } + +// TestComputePlainTarKeyProperties pins the working-tree fingerprint behavior: it is +// order-independent and reacts to any change in a file's path, size, or mtime. +func TestComputePlainTarKeyProperties(t *testing.T) { + base := []snapshotFile{ + {rel: "a.txt", size: 10, modTime: 100}, + {rel: "src/b.py", size: 20, modTime: 200}, + } + + // Order-independent: the files are sorted by path before hashing. + assert.Equal(t, + computePlainTarKey(base), + computePlainTarKey([]snapshotFile{base[1], base[0]}), + ) + + // A changed size, mtime, or path each yields a different key. + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 11, modTime: 100}, base[1]})) + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 10, modTime: 101}, base[1]})) + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "renamed.txt", size: 10, modTime: 100}, base[1]})) + + // Adding or dropping a file changes the key. + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey(base[:1])) +} diff --git a/experimental/air/cmd/snapshot_dabs.go b/experimental/air/cmd/snapshot_dabs.go index fd753d050d8..706b3c2102a 100644 --- a/experimental/air/cmd/snapshot_dabs.go +++ b/experimental/air/cmd/snapshot_dabs.go @@ -44,8 +44,9 @@ func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, s return snapshotResult{}, err } - // Resolve how to package before touching the tarball: git_archive (pinned commit, - // cacheable) vs plain_tar (working tree, not cacheable). + // Resolve how to package before touching the tarball: git_archive (pinned commit) vs + // plain_tar (working tree). Both are content-addressed, so an unchanged input reuses + // the already-uploaded tarball instead of re-packaging and re-uploading. plan, err := resolveSnapshotPlan(ctx, newGitRepo(repoPath), snap.Git, snap.IncludePaths) if err != nil { return snapshotResult{}, err @@ -126,28 +127,33 @@ func uploadSnapshotSidecars(ctx context.Context, sidecarStore filer.Filer, sidec return path.Join(sidecarBase, gitStateName), diffPath } -// snapshotTarballName is the uploaded filename for the snapshot. It is deterministic -// for git_archive — _.tar.gz keyed on (commit, include_paths, -// root_path subtree) — so an identical snapshot reuses the same remote object (see -// the cache check below). For plain_tar it is timestamped so concurrent submissions -// of the same directory don't clobber each other's upload (working-tree content -// isn't pinned to a SHA, so it can't be content-addressed). -func snapshotTarballName(plan snapshotPlan, dirName string) string { - if plan.mode == modeGitArchive { - key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths, plan.subtreePrefix) - return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]) +// snapshotTarName resolves the content-addressed upload filename for the snapshot and, +// for plain_tar, the working-tree file listing used to build both the key and the tarball +// (nil for git_archive, which lists nothing locally). The name is _.tar.gz, +// keyed on (commit, include_paths, root_path subtree) for git_archive and on the +// working-tree fingerprint (path+size+mtime) for plain_tar, so an identical input +// reuses the same remote object (see the skip in uploadSnapshotViaDABs). +func snapshotTarName(ctx context.Context, repoPath string, plan snapshotPlan) (string, []snapshotFile, error) { + dirName := filepath.Base(repoPath) + if plan.mode == modeGitArchive { + key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths, plan.subtreePrefix) + return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]), nil, nil + } + files, err := snapshotFiles(ctx, repoPath, plan.includePaths, plan.isGitRepo) + if err != nil { + return "", nil, err } - return fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) + return fmt.Sprintf("%s_%s.tar.gz", dirName, computePlainTarKey(files)[:16]), files, nil } // 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) - if plan.mode == modeGitArchive { - return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths, plan.subtreePrefix) - } - return createPlainTarball(ctx, repoPath, tarball, plan.includePaths, plan.isGitRepo) +// of the pinned commit for git_archive, else a plain tar of the pre-listed working-tree +// files (nil for git_archive). +func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, files []snapshotFile, tarball string) error { + if plan.mode == modeGitArchive { + return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, filepath.Base(repoPath), plan.includePaths, plan.subtreePrefix) + } + return createPlainTarball(ctx, repoPath, tarball, files) } // uploadSnapshotViaDABs uploads the snapshot through DABs' artifact-upload machinery @@ -156,9 +162,10 @@ func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, ta // the remote .internal path, and uploads the bytes. artifactPath is either the // already-resolved user's repo_snapshots directory or a configured UC Volume. // -// git_archive snapshots are cacheable: the tarball name is content-addressed by -// (commit, include_paths, root_path subtree), so if the identical object is already -// uploaded we skip packaging and upload entirely and just reuse the remote path. +// The tarball name is content-addressed — by (commit, include_paths, root_path subtree) +// for git_archive and by the working-tree fingerprint (path+size+mtime) for plain_tar — +// so if the identical object is already uploaded we skip packaging and upload entirely +// and reuse the remote path. func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath string, plan snapshotPlan, artifactPath string) (snapshotResult, error) { // artifactPath is where DABs uploads the tarball; GetFilerForLibraries routes to // a Workspace or Volume filer based on its prefix, then appends /.internal. @@ -168,7 +175,10 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r } defer os.RemoveAll(tmp) - tarName := snapshotTarballName(plan, filepath.Base(repoPath)) + tarName, files, err := snapshotTarName(ctx, repoPath, plan) + if err != nil { + return snapshotResult{}, err + } b := &bundle.Bundle{ BundleRootPath: tmp, @@ -199,33 +209,31 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r return snapshotResult{}, err } - // git_archive is cacheable by (commit, include_paths, root_path subtree): if the - // identical tarball is already uploaded, skip packaging + upload and reuse it. - // Only the config-path rewrite (ReplaceWithRemotePath) runs — no bytes move. - if plan.mode == modeGitArchive { - f, uploadPath, diags := libraries.GetFilerForLibraries(ctx, b) - if diags.HasError() { - return snapshotResult{}, diags.Error() + // Both modes are content-addressed by tarName: if the identical tarball is already + // uploaded, skip packaging + upload and reuse it. Only the config-path rewrite + // (ReplaceWithRemotePath) runs — no bytes move. + f, uploadPath, diags := libraries.GetFilerForLibraries(ctx, b) + if diags.HasError() { + return snapshotResult{}, diags.Error() + } + exists, err := snapshotExists(ctx, f, tarName) + if err != nil { + return snapshotResult{}, err + } + if exists { + if _, diags := libraries.ReplaceWithRemotePath(ctx, b); diags.HasError() { + return snapshotResult{}, diags.Error() } - exists, err := snapshotExists(ctx, f, tarName) + remote, err := readCodeSourcePath(b) if err != nil { return snapshotResult{}, err } - if exists { - if _, diags := libraries.ReplaceWithRemotePath(ctx, b); diags.HasError() { - return snapshotResult{}, diags.Error() - } - remote, err := readCodeSourcePath(b) - if err != nil { - return snapshotResult{}, err - } - log.Debugf(ctx, "snapshot cache hit for %s at %s", shortSHA(plan.commitSHA), path.Join(uploadPath, tarName)) - return snapshotResult{CodeSourcePath: remote}, nil - } + log.Debugf(ctx, "snapshot upload skipped; reusing %s", path.Join(uploadPath, tarName)) + return snapshotResult{CodeSourcePath: remote}, nil } - // 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 { + // Miss: package the tarball locally, then upload the bytes. + if err := packageSnapshot(ctx, repoPath, plan, files, filepath.Join(tmp, tarName)); err != nil { return snapshotResult{}, err } diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index 131580c0aa6..979d8c91e84 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -49,25 +49,21 @@ func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outpu return nil } -// createPlainTarball writes a gzipped tar of repoPath's working tree to +// createPlainTarball writes a gzipped tar of the given working-tree files to // outputTarball. The archive preserves repoPath's directory name as the top-level -// entry. When includePaths is set, only those paths (nested under the directory -// name) are archived. .git and macOS AppleDouble files are always excluded; a -// .gitignore at repoPath is honored. -func createPlainTarball(ctx context.Context, repoPath, outputTarball string, includePaths []string, isGitRepo bool) error { - dirName := filepath.Base(repoPath) +// entry. files come pre-resolved from snapshotFiles (.gitignore honored, .git and +// macOS AppleDouble excluded), so the caller can reuse the same listing to +// content-address the upload. +func createPlainTarball(ctx context.Context, repoPath, outputTarball string, files []snapshotFile) error { + dirName := filepath.Base(repoPath) - files, err := snapshotFiles(ctx, repoPath, includePaths, isGitRepo) - if err != nil { - return err - } - entries := make([]tarpack.Entry, len(files)) - for i, rel := range files { - entries[i] = tarpack.Entry{ - Name: filepath.ToSlash(filepath.Join(dirName, rel)), - Path: filepath.Join(repoPath, rel), - } - } + entries := make([]tarpack.Entry, len(files)) + for i, f := range files { + entries[i] = tarpack.Entry{ + Name: filepath.ToSlash(filepath.Join(dirName, f.rel)), + Path: filepath.Join(repoPath, f.rel), + } + } out, err := os.Create(outputTarball) if err != nil { @@ -95,7 +91,15 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc return out.Close() } -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) plus the size and mtime used to content-address the plain_tar upload. +type snapshotFile struct { + rel string + size int64 + modTime int64 // Unix nanoseconds +} + +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-") @@ -119,7 +123,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 @@ -129,14 +133,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/experimental/air/cmd/snapshot_package_test.go b/experimental/air/cmd/snapshot_package_test.go index 9c4f07c8794..9bda10dad43 100644 --- a/experimental/air/cmd/snapshot_package_test.go +++ b/experimental/air/cmd/snapshot_package_test.go @@ -132,7 +132,9 @@ func TestCreatePlainTarball(t *testing.T) { writeRepoFile(t, repo, ".git/config", "x") out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(ctx, repo, out, nil, false)) + files, err := snapshotFiles(ctx, repo, nil, false) + require.NoError(t, err) + require.NoError(t, createPlainTarball(ctx, repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) @@ -152,7 +154,9 @@ func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { writeRepoFile(t, repo, ".gitignore", "*.log\n") out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(ctx, repo, out, nil, false)) + files, err := snapshotFiles(ctx, repo, nil, false) + require.NoError(t, err) + require.NoError(t, createPlainTarball(ctx, repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) @@ -167,7 +171,9 @@ func TestCreatePlainTarball_IncludePaths(t *testing.T) { writeRepoFile(t, repo, "src/model.py", "print()") out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(ctx, repo, out, []string{"src"}, false)) + files, err := snapshotFiles(ctx, repo, []string{"src"}, false) + require.NoError(t, err) + require.NoError(t, createPlainTarball(ctx, repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) @@ -187,7 +193,9 @@ func TestCreatePlainTarball_HonorsNestedGitignoreAndNegation(t *testing.T) { writeRepoFile(t, repo, "nested/keep.tmp", "keep") out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(t.Context(), repo, out, nil, false)) + files, err := snapshotFiles(t.Context(), repo, nil, false) + require.NoError(t, err) + require.NoError(t, createPlainTarball(t.Context(), repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) @@ -207,7 +215,9 @@ func TestCreatePlainTarball_SkipsDeletedTrackedFiles(t *testing.T) { require.NoError(t, os.Remove(filepath.Join(repo, "deleted.txt"))) out := filepath.Join(t.TempDir(), "snap.tar.gz") - require.NoError(t, createPlainTarball(t.Context(), repo, out, nil, true)) + files, err := snapshotFiles(t.Context(), repo, nil, true) + require.NoError(t, err) + require.NoError(t, createPlainTarball(t.Context(), repo, out, files)) dirName := filepath.Base(repo) entries := tarballEntries(t, out) From 560d24828b3e8216425717564f8a3a8b89035cdb Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Wed, 9 Sep 2026 17:22:28 +0000 Subject: [PATCH 2/7] experimental/air: address review feedback on content-addressed plain_tar - Strengthen the dedup test: count import-file calls and assert the second (unchanged) submit adds zero, instead of asserting a path-keyed set has one entry. The set couldn't distinguish a skipped submit from a re-upload to the same content-addressed name; the counter can (verified it fails when the skip is disabled). - Fold snapshotPackagingVersion into computePlainTarKey so a packaging-logic bump invalidates plain_tar keys too, not just plainTarKeyVersion. - Fix stale comments now that plain_tar is content-addressed: modePlainTar ("not cacheable") and snapshotExists ("git_archive" only). Co-authored-by: Isaac --- experimental/air/cmd/runsubmit_test.go | 19 ++++++++++++------- experimental/air/cmd/snapshot_cachekey.go | 4 +++- experimental/air/cmd/snapshot_dabs.go | 4 ++-- experimental/air/cmd/snapshot_resolve.go | 3 ++- 4 files changed, 19 insertions(+), 11 deletions(-) diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 92dc69af481..f079b543d1f 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -565,14 +565,16 @@ func TestSubmitWorkloadPlainTarContentAddressed(t *testing.T) { server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { return jobs.SubmitRunResponse{RunId: 555} }) - // Track snapshot uploads, preserving fake-workspace persistence so the second - // submit's existence Stat sees the first upload. Dedupe by path: the DABs uploader - // mkdirs-and-retries the import, so one logical upload can hit this route twice. - uploaded := map[string]bool{} + // Count snapshot import-file calls, preserving fake-workspace persistence so the + // second submit's existence Stat sees the first upload. A count (not a set keyed by + // path) is what proves the skip: both submits resolve to the same content-addressed + // name, so a set could not tell a skipped second submit from one that re-uploaded to + // that same path. + snapshotUploads := 0 server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { p := req.Vars["path"] if strings.Contains(p, "/.air/repo_snapshots/") { - uploaded[p] = true + snapshotUploads++ } return req.Workspace.WorkspaceFilesImportFile(p, req.Body, req.URL.Query().Get("overwrite") == "true") }) @@ -600,6 +602,8 @@ code_source: sidecarStore, sidecarBase := testSidecarStore(t, w) first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) require.NoError(t, err) + require.NotZero(t, snapshotUploads, "first submit should upload the tarball") + afterFirst := snapshotUploads second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) require.NoError(t, err) @@ -608,9 +612,10 @@ code_source: assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be content-addressed, not the bare dir name") assert.Regexp(t, `^src_[0-9a-f]{16}\.tar\.gz$`, base) - // Same unchanged tree → identical remote path, uploaded once (second submit is a hit). + // Same unchanged tree → identical remote path, and the second submit moved no bytes: + // zero new import-file calls (a real skip), not a re-upload to the same name. assert.Equal(t, first.CodeSourcePath, second.CodeSourcePath) - assert.Len(t, uploaded, 1, "unchanged plain_tar should skip the second upload") + assert.Equal(t, afterFirst, snapshotUploads, "unchanged plain_tar should skip the second upload") } // A git_archive snapshot is content-addressed by (commit, include_paths): submitting diff --git a/experimental/air/cmd/snapshot_cachekey.go b/experimental/air/cmd/snapshot_cachekey.go index fd680b880a2..4bebf2b7643 100644 --- a/experimental/air/cmd/snapshot_cachekey.go +++ b/experimental/air/cmd/snapshot_cachekey.go @@ -18,6 +18,8 @@ const snapshotPackagingVersion = "v1" // plainTarKeyVersion namespaces the plain_tar working-tree key (so it can never collide // with a git_archive key) and lets us invalidate it if the fingerprint scheme changes. +// computePlainTarKey also folds in the shared snapshotPackagingVersion, so a +// packaging-logic bump invalidates both modes' keys. const plainTarKeyVersion = "plaintar-v1" // computePlainTarKey returns a content-addressed key for a working-tree snapshot: the @@ -35,7 +37,7 @@ func computePlainTarKey(files []snapshotFile) string { for _, f := range sorted { fmt.Fprintf(h, "%s\x00%d\x00%d\n", filepath.ToSlash(f.rel), f.size, f.modTime) } - fmt.Fprint(h, plainTarKeyVersion) + fmt.Fprintf(h, "%s\x00%s", plainTarKeyVersion, snapshotPackagingVersion) return hex.EncodeToString(h.Sum(nil)) } diff --git a/experimental/air/cmd/snapshot_dabs.go b/experimental/air/cmd/snapshot_dabs.go index 706b3c2102a..bc157cd25a8 100644 --- a/experimental/air/cmd/snapshot_dabs.go +++ b/experimental/air/cmd/snapshot_dabs.go @@ -253,8 +253,8 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r } // snapshotExists reports whether name already exists in the artifact store, used to -// short-circuit a cacheable git_archive upload. A not-found is a clean miss (false, -// nil); any other error is surfaced. +// short-circuit a content-addressed upload (either mode). A not-found is a clean miss +// (false, nil); any other error is surfaced. func snapshotExists(ctx context.Context, store filer.Filer, name string) (bool, error) { _, err := store.Stat(ctx, name) if err == nil { diff --git a/experimental/air/cmd/snapshot_resolve.go b/experimental/air/cmd/snapshot_resolve.go index d799dd96a77..0fcd3e82657 100644 --- a/experimental/air/cmd/snapshot_resolve.go +++ b/experimental/air/cmd/snapshot_resolve.go @@ -20,7 +20,8 @@ const ( // root_path subtree). modeGitArchive snapshotMode = iota // modePlainTar packages the working tree (including uncommitted changes) via - // `tar`. Not cacheable — working-tree content isn't pinned to a SHA. + // `tar`. Content-addressed by the working-tree fingerprint (path+size+mtime), so an + // unchanged tree reuses the uploaded tarball; a same-size, same-mtime edit is missed. modePlainTar ) From edecb8a0d2716601077918fb57eda2eab7484203 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 14:39:50 +0000 Subject: [PATCH 3/7] experimental/air: warm snapshot cache for the plain_tar path Add a local warm cache for the plain_tar snapshot path, 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 stat the file set, copy unchanged members verbatim from the warm tar, and re-read only changed files before gzipping the upload. --no-cache bypasses it and re-packs from scratch. The cache engages only above 64 MiB. The cache's payoff is largest when the working set does not fit the OS page cache: cold, scattered small-file reads cost seconds for a large folder versus ~ms to read the warm tar sequentially. When the tree is already warm in RAM, parallel gzip accounts for most of the gain and the cache adds little. Co-authored-by: Isaac --- experimental/air/cmd/run.go | 4 +- experimental/air/cmd/runsubmit.go | 4 +- experimental/air/cmd/runsubmit_test.go | 26 +- experimental/air/cmd/snapshot_cache.go | 390 ++++++++++++++++++++ experimental/air/cmd/snapshot_cache_test.go | 183 +++++++++ experimental/air/cmd/snapshot_dabs.go | 49 +-- experimental/air/cmd/snapshot_package.go | 19 +- 7 files changed, 626 insertions(+), 49 deletions(-) create mode 100644 experimental/air/cmd/snapshot_cache.go create mode 100644 experimental/air/cmd/snapshot_cache_test.go diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index 2e789869121..ae50cfb19bf 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 72e128b8f34..fb2a7d13235 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -273,7 +273,7 @@ func stageRunArtifacts(ctx context.Context, launchWriter fileWriter, items []upl // 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 // staging spinner (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 @@ -346,7 +346,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run snapshotArtifactPath := path.Join(base, ".air", "repo_snapshots") uploadSnapshot = func(ctx context.Context) (snapshotResult, error) { // Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh. - return snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, snapshotArtifactPath, fc, funcDir) + return snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, snapshotArtifactPath, fc, funcDir, noCache) } } diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index f079b543d1f..839264d8673 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -337,7 +337,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") @@ -408,7 +408,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) @@ -490,7 +490,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 @@ -535,7 +535,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 @@ -600,11 +600,11 @@ code_source: ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase, false) require.NoError(t, err) require.NotZero(t, snapshotUploads, "first submit should upload the tarball") afterFirst := snapshotUploads - second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase, false) require.NoError(t, err) // Content-addressed name: not the bare dir, but a 16-hex-char fingerprint. @@ -663,9 +663,9 @@ code_source: ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase, false) require.NoError(t, err) - second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase, false) require.NoError(t, err) // Same pinned commit → identical content-addressed remote path, uploaded once @@ -707,7 +707,7 @@ code_source: ctx := cmdio.MockDiscard(t.Context()) sidecarStore, sidecarBase := testSidecarStore(t, w) - snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, testSnapshotArtifactPath, sidecarStore, sidecarBase, false) require.NoError(t, err) assert.Empty(t, snap.GitStatePath) @@ -756,7 +756,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 @@ -791,7 +791,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") @@ -828,7 +828,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) }) @@ -839,7 +839,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 bc157cd25a8..8c54a87fc90 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, snapshotArtifactPath string, sidecarStore filer.Filer, sidecarBase string) (snapshotResult, error) { +func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath, snapshotArtifactPath 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 { snapshotArtifactPath = *snap.RemoteVolume } - result, err := uploadSnapshotViaDABs(ctx, w, repoPath, plan, snapshotArtifactPath) + result, err := uploadSnapshotViaDABs(ctx, w, repoPath, configPath, plan, snapshotArtifactPath, noCache) if err != nil { return snapshotResult{}, err } @@ -134,11 +134,11 @@ func uploadSnapshotSidecars(ctx context.Context, sidecarStore filer.Filer, sidec // working-tree fingerprint (path+size+mtime) for plain_tar, so an identical input // reuses the same remote object (see the skip in uploadSnapshotViaDABs). func snapshotTarName(ctx context.Context, repoPath string, plan snapshotPlan) (string, []snapshotFile, error) { - dirName := filepath.Base(repoPath) - if plan.mode == modeGitArchive { - key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths, plan.subtreePrefix) - return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]), nil, nil - } + dirName := filepath.Base(repoPath) + if plan.mode == modeGitArchive { + key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths, plan.subtreePrefix) + return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]), nil, nil + } files, err := snapshotFiles(ctx, repoPath, plan.includePaths, plan.isGitRepo) if err != nil { return "", nil, err @@ -146,14 +146,17 @@ func snapshotTarName(ctx context.Context, repoPath string, plan snapshotPlan) (s return fmt.Sprintf("%s_%s.tar.gz", dirName, computePlainTarKey(files)[:16]), files, nil } -// 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 pre-listed working-tree -// files (nil for git_archive). -func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, files []snapshotFile, tarball string) error { - if plan.mode == modeGitArchive { - return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, filepath.Base(repoPath), plan.includePaths, plan.subtreePrefix) - } - return createPlainTarball(ctx, repoPath, tarball, files) +// packageSnapshot writes the snapshot to tarball per the resolved plan. The working-tree +// path uses the warm snapshot cache unless noCache is set. +func packageSnapshot(ctx context.Context, repoPath, configPath string, plan snapshotPlan, files []snapshotFile, tarball string, noCache bool) error { + dirName := filepath.Base(repoPath) + if plan.mode == modeGitArchive { + return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths, plan.subtreePrefix) + } + if noCache { + return createPlainTarball(ctx, repoPath, tarball, files) + } + return packagePlainTarWithCache(ctx, repoPath, configPath, plan.includePaths, plan.isGitRepo, tarball) } // uploadSnapshotViaDABs uploads the snapshot through DABs' artifact-upload machinery @@ -166,7 +169,7 @@ func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, fi // for git_archive and by the working-tree fingerprint (path+size+mtime) for plain_tar — // so if the identical object is already uploaded we skip packaging and upload entirely // and reuse the remote path. -func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath string, plan snapshotPlan, artifactPath string) (snapshotResult, error) { +func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath, configPath string, plan snapshotPlan, artifactPath 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. tmp, err := os.MkdirTemp("", "air-snapshot-*") @@ -209,9 +212,9 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r return snapshotResult{}, err } - // Both modes are content-addressed by tarName: if the identical tarball is already - // uploaded, skip packaging + upload and reuse it. Only the config-path rewrite - // (ReplaceWithRemotePath) runs — no bytes move. + // Both modes are content-addressed by tarName: if the identical tarball is already + // uploaded, skip packaging + upload and reuse it. Only the config-path rewrite + // (ReplaceWithRemotePath) runs — no bytes move. f, uploadPath, diags := libraries.GetFilerForLibraries(ctx, b) if diags.HasError() { return snapshotResult{}, diags.Error() @@ -220,9 +223,9 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r if err != nil { return snapshotResult{}, err } - if exists { - if _, diags := libraries.ReplaceWithRemotePath(ctx, b); diags.HasError() { - return snapshotResult{}, diags.Error() + if exists { + if _, diags := libraries.ReplaceWithRemotePath(ctx, b); diags.HasError() { + return snapshotResult{}, diags.Error() } remote, err := readCodeSourcePath(b) if err != nil { @@ -233,7 +236,7 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r } // Miss: package the tarball locally, then upload the bytes. - if err := packageSnapshot(ctx, repoPath, plan, files, filepath.Join(tmp, tarName)); err != nil { + if err := packageSnapshot(ctx, repoPath, configPath, plan, files, filepath.Join(tmp, tarName), noCache); err != nil { return snapshotResult{}, err } diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index 979d8c91e84..cce694d5ddb 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -55,15 +55,14 @@ func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outpu // macOS AppleDouble excluded), so the caller can reuse the same listing to // content-address the upload. func createPlainTarball(ctx context.Context, repoPath, outputTarball string, files []snapshotFile) error { - dirName := filepath.Base(repoPath) - - entries := make([]tarpack.Entry, len(files)) - for i, f := range files { - entries[i] = tarpack.Entry{ - Name: filepath.ToSlash(filepath.Join(dirName, f.rel)), - Path: filepath.Join(repoPath, f.rel), - } - } + dirName := filepath.Base(repoPath) + entries := make([]tarpack.Entry, len(files)) + for i, f := range files { + entries[i] = tarpack.Entry{ + Name: filepath.ToSlash(filepath.Join(dirName, f.rel)), + Path: filepath.Join(repoPath, f.rel), + } + } out, err := os.Create(outputTarball) if err != nil { @@ -92,7 +91,7 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, fil } // snapshotFile is a file selected for the snapshot: its repo-relative path (native -// separators) plus the size and mtime used to content-address the plain_tar upload. +// separators) plus the size and mtime used by content addressing and the warm cache. type snapshotFile struct { rel string size int64 From 72534e2e12a381e03004ad0e1f0cc1eea2542b5e Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 17:09:33 +0000 Subject: [PATCH 4/7] experimental/air: use default gzip level in the warm cache too Match the parent PR: DefaultCompression rather than BestSpeed in newGzFile, so the cached tarball is re-gzipped at the same level as the --no-cache path and the upload stays small. Parallel compression makes the higher level nearly free. Co-authored-by: Isaac --- experimental/air/cmd/snapshot_cache.go | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/experimental/air/cmd/snapshot_cache.go b/experimental/air/cmd/snapshot_cache.go index 82ce0c64ad7..60a0a793896 100644 --- a/experimental/air/cmd/snapshot_cache.go +++ b/experimental/air/cmd/snapshot_cache.go @@ -303,23 +303,24 @@ func (c *countWriter) Write(p []byte) (int, error) { 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. +// newGzFile creates outputTarball and returns a parallel gzip writer over it plus an +// idempotent close that flushes the gzip stream and the file. // // 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. +// re-gzips the warm tar — so it is spread across cores (~18x faster than compress/gzip +// on a 470 MiB tar). The level is DefaultCompression, not BestSpeed: the tarball is +// re-uploaded every run so its size matters, and with parallel compression a normal +// level is nearly free (a few hundred ms for ~15-18% fewer bytes). 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) + gz, err := pgzip.NewWriterLevel(f, pgzip.DefaultCompression) if err != nil { f.Close() return nil, nil, err From aa05c10e77203df20b91a4802bad00c4d11322c4 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 17:47:02 +0000 Subject: [PATCH 5/7] experimental/air: update config-help golden for --no-cache flag The new --no-cache flag on `air run` adds a line to its --help output, which the experimental/air/config-help acceptance test pins. Regenerate the golden. Co-authored-by: Isaac --- acceptance/experimental/air/config-help/output.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/acceptance/experimental/air/config-help/output.txt b/acceptance/experimental/air/config-help/output.txt index a17388562ad..ee0ce7f6687 100644 --- a/acceptance/experimental/air/config-help/output.txt +++ b/acceptance/experimental/air/config-help/output.txt @@ -22,6 +22,7 @@ Flags: -f, --file string Path to the workload YAML config -h, --help help for run --idempotency-key string Return the existing run if this key was already used + --no-cache Bypass the local snapshot cache and re-pack the code tarball from scratch --override stringArray Override a YAML field, e.g. compute.num_accelerators=8 (repeatable) --watch Stream logs until the run completes From de856467f47689f5bc334eff865341061d238fc1 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 8 Sep 2026 19:17:31 +0000 Subject: [PATCH 6/7] experimental/air: make the warm cache crash- and concurrency-safe Isaac Review flagged two MAJOR correctness bugs in the warm cache: - Concurrent `air run` on the same cache key wrote the same snapshot.tar.tmp and raced the rename plus a non-atomic manifest write, interleaving into a corrupt tar/manifest pair. - rebuildWarmSnapshot renamed the new tar into place before saving the manifest, so a crash between the two left a manifest whose byte offsets described a different tar layout -- silently corrupting a later verbatim-reuse rebuild. Fix both by binding the manifest to a per-build, uniquely named tar (snapshot..tar) that is never overwritten, and installing the manifest atomically (unique temp + rename) only after its tar is durable. A manifest and the tar it indexes are therefore always a consistent pair: there is no window where offsets describe a mismatched tar, and concurrent rebuilds are last-writer-wins on the manifest rather than interleaving, so no lock is needed. Superseded and orphaned tars are cleaned up best-effort. Adds a rotation test. Co-authored-by: Isaac --- experimental/air/cmd/snapshot_cache.go | 110 ++++++++++++++++---- experimental/air/cmd/snapshot_cache_test.go | 65 +++++++++--- 2 files changed, 140 insertions(+), 35 deletions(-) diff --git a/experimental/air/cmd/snapshot_cache.go b/experimental/air/cmd/snapshot_cache.go index 60a0a793896..e319c84be90 100644 --- a/experimental/air/cmd/snapshot_cache.go +++ b/experimental/air/cmd/snapshot_cache.go @@ -14,6 +14,7 @@ package aircmd import ( "archive/tar" "context" + "crypto/rand" "crypto/sha256" "encoding/hex" "encoding/json" @@ -34,7 +35,10 @@ const ( // snapshotCacheVersion invalidates on-disk caches when the layout below changes. snapshotCacheVersion = "v1" snapshotCacheManifestName = "manifest.json" - snapshotCacheTarName = "snapshot.tar" + // snapshotTarPrefix begins each warm tar's filename; the rest is a per-build random + // id (snapshot..tar). A build never overwrites another build's tar, so a manifest + // and the tar it names stay a consistent pair — see rebuildWarmSnapshot. + snapshotTarPrefix = "snapshot." // 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 @@ -57,9 +61,11 @@ type cacheEntry struct { Length int64 `json:"length"` } -// snapshotManifest is the on-disk index of a warm snapshot.tar. +// snapshotManifest is the on-disk index of a warm snapshot tar. TarName binds it to the +// exact tar its byte offsets describe, so a stale or half-written pair is never reused. type snapshotManifest struct { Version string `json:"version"` + TarName string `json:"tar_name"` // the snapshot..tar this manifest indexes DirName string `json:"dir_name"` Entries map[string]cacheEntry `json:"entries"` // keyed by slash-separated relative path } @@ -116,13 +122,16 @@ func packagePlainTarWithCache(ctx context.Context, repoPath, configPath string, 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) { + var oldTarPath string + if old != nil { + oldTarPath = filepath.Join(cacheDir, old.TarName) + } + if old != nil && old.DirName == dirName && fileExists(oldTarPath) && !snapshotChanged(files, old) { mode = "hit-nochange" - return gzipFile(tarPath, outputTarball) + return gzipFile(oldTarPath, outputTarball) } if err := os.MkdirAll(cacheDir, 0o700); err != nil { @@ -131,7 +140,7 @@ func packagePlainTarWithCache(ctx context.Context, repoPath, configPath string, if old != nil { mode = "rebuild-warm" } - return rebuildWarmSnapshot(repoPath, dirName, files, old, tarPath, outputTarball) + return rebuildWarmSnapshot(repoPath, dirName, files, old, cacheDir, oldTarPath, outputTarball) } // snapshotChanged reports whether the current file set differs from the manifest by @@ -150,52 +159,70 @@ func snapshotChanged(files []snapshotFile, old *snapshotManifest) bool { 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) { +// rebuildWarmSnapshot writes a fresh warm tar and its gzipped upload copy in one pass +// (copying unchanged members verbatim from the previous build's tar), then installs a +// manifest pointing at the new tar. +// +// Correctness under crashes and concurrent runs rests on two properties: each build +// names its tar uniquely (snapshot..tar) and never overwrites another's, and the +// manifest is renamed into place atomically only after the tar it names is fully +// written. So a manifest and the tar it indexes are always a consistent pair — there is +// no window where the manifest's byte offsets describe a tar with a different layout +// (which would silently corrupt a later verbatim reuse), and two concurrent rebuilds are +// last-writer-wins on the manifest rather than interleaving into one file. No lock needed. +func rebuildWarmSnapshot(repoPath, dirName string, files []snapshotFile, old *snapshotManifest, cacheDir, oldTarPath, outputTarball string) (err error) { gz, closeGz, err := newGzFile(outputTarball) if err != nil { return err } defer func() { err = firstErr(err, closeGz()) }() - newTarPath := tarPath + ".tmp" + buildID, err := randomID() + if err != nil { + return err + } + newTarName := snapshotTarPrefix + buildID + ".tar" + newTarPath := filepath.Join(cacheDir, newTarName) 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) + tarFile.Close() + os.Remove(newTarPath) // unreferenced on failure; don't leave an orphan } }() + // Reuse unchanged members from the previous build's tar. Its offsets come from old's + // manifest, which indexes exactly that (immutable) tar, so the ranges stay valid. var oldTar io.ReaderAt - if old != nil { - if f, e := os.Open(tarPath); e == nil { + if old != nil && oldTarPath != "" { + if f, e := os.Open(oldTarPath); e == nil { defer f.Close() oldTar = f } else { - old = nil // warm tar gone; rebuild every member from disk + old = nil // previous tar gone; rebuild every member from disk } } manifest, err := writeSnapshot(repoPath, dirName, files, old, oldTar, tarFile, gz) if err != nil { - tarFile.Close() return err } + manifest.TarName = newTarName 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) + // Install the manifest only now that its tar is durable; the write is atomic. + if err = saveSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName), manifest); err != nil { + return err } - return saveSnapshotManifest(filepath.Join(filepath.Dir(tarPath), snapshotCacheManifestName), manifest) + cleanupOldTars(cacheDir, newTarName) + return nil } // writeGzOnly packs files straight to a gzipped tarball without persisting a cache, @@ -336,8 +363,8 @@ func newGzFile(outputTarball string) (io.Writer, func() error, error) { 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. +// gzipFile gzips src to outputTarball via newGzFile (parallel gzip). 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 { @@ -365,17 +392,56 @@ func loadSnapshotManifest(path string) *snapshotManifest { return &m } +// saveSnapshotManifest writes the manifest atomically (unique temp + rename) so a reader +// or a concurrent run never sees a half-written manifest, and the pair with its tar swaps +// in as a unit. 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 { + tmp, err := os.CreateTemp(filepath.Dir(path), "manifest-*.json") + if err != nil { + return fmt.Errorf("failed to write snapshot manifest: %w", err) + } + tmpName := tmp.Name() + if _, err := tmp.Write(data); err != nil { + tmp.Close() + os.Remove(tmpName) + return fmt.Errorf("failed to write snapshot manifest: %w", err) + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) return fmt.Errorf("failed to write snapshot manifest: %w", err) } + if err := os.Rename(tmpName, path); err != nil { + os.Remove(tmpName) + return fmt.Errorf("failed to install snapshot manifest: %w", err) + } return nil } +// randomID returns a short random hex string used to name each warm tar uniquely. +func randomID() (string, error) { + var b [8]byte + if _, err := rand.Read(b[:]); err != nil { + return "", fmt.Errorf("failed to generate cache build id: %w", err) + } + return hex.EncodeToString(b[:]), nil +} + +// cleanupOldTars removes warm tars other than keep: superseded builds, and tars a +// concurrent rebuild orphaned. Best-effort — only keep is referenced by the manifest, +// and a wrongly removed tar costs at most a cold rebuild, never correctness. +func cleanupOldTars(cacheDir, keep string) { + matches, _ := filepath.Glob(filepath.Join(cacheDir, snapshotTarPrefix+"*.tar")) + for _, p := range matches { + if filepath.Base(p) != keep { + os.Remove(p) + } + } +} + func fileExists(path string) bool { _, err := os.Stat(path) return err == nil diff --git a/experimental/air/cmd/snapshot_cache_test.go b/experimental/air/cmd/snapshot_cache_test.go index 134f3418c0e..717eb62fdfb 100644 --- a/experimental/air/cmd/snapshot_cache_test.go +++ b/experimental/air/cmd/snapshot_cache_test.go @@ -79,6 +79,15 @@ func TestSnapshotChanged(t *testing.T) { assert.True(t, snapshotChanged(added, old), "addition detected") } +// warmTarPath returns the warm tar the cache's manifest currently points at. +func warmTarPath(t *testing.T, cacheDir string) string { + t.Helper() + m := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, m) + require.NotEmpty(t, m.TarName) + return filepath.Join(cacheDir, m.TarName) +} + func TestWarmSnapshotColdBuild(t *testing.T) { repo := t.TempDir() writeRepoFile(t, repo, "a.txt", "alpha") @@ -86,22 +95,23 @@ func TestWarmSnapshotColdBuild(t *testing.T) { 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)) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", 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) + // The warm tar and manifest are persisted for the next run, and the manifest names + // the tar it indexes. m := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) require.NotNil(t, m) assert.Equal(t, dirName, m.DirName) assert.Len(t, m.Entries, 2) + assert.NotEmpty(t, m.TarName) + assert.FileExists(t, filepath.Join(cacheDir, m.TarName)) } func TestWarmSnapshotReusesUnchangedFromCache(t *testing.T) { @@ -111,12 +121,12 @@ func TestWarmSnapshotReusesUnchangedFromCache(t *testing.T) { 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"))) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", filepath.Join(t.TempDir(), "cold.tar.gz"))) old := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) require.NotNil(t, old) + oldTarPath := filepath.Join(cacheDir, old.TarName) // 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 @@ -124,7 +134,7 @@ func TestWarmSnapshotReusesUnchangedFromCache(t *testing.T) { 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)) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, old, cacheDir, oldTarPath, out)) contents := extractTarball(t, out) assert.Equal(t, "keep", contents[dirName+"/keep.py"], "unchanged member copied from warm tar") @@ -138,11 +148,11 @@ func TestWarmSnapshotRebuildEditAddDelete(t *testing.T) { 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"))) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", filepath.Join(t.TempDir(), "cold.tar.gz"))) old := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) require.NotNil(t, old) + oldTarPath := filepath.Join(cacheDir, old.TarName) // Edit a.txt (changed), delete b.txt, add c.txt. writeRepoFile(t, repo, "a.txt", "alpha-v2") @@ -151,7 +161,7 @@ func TestWarmSnapshotRebuildEditAddDelete(t *testing.T) { 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)) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, newFiles, old, cacheDir, oldTarPath, out)) contents := extractTarball(t, out) assert.Equal(t, "alpha-v2", contents[dirName+"/a.txt"], "edited file updated") @@ -171,13 +181,42 @@ func TestGzipFileRoundTrip(t *testing.T) { 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"))) + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", 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)) + require.NoError(t, gzipFile(warmTarPath(t, cacheDir), out)) contents := extractTarball(t, out) assert.Equal(t, "alpha", contents[dirName+"/a.txt"]) } + +func TestWarmSnapshotRebuildRotatesTar(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() + + files := statFiles(t, repo, "a.txt", "b.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files, nil, cacheDir, "", filepath.Join(t.TempDir(), "1.tar.gz"))) + m1 := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, m1) + tar1 := filepath.Join(cacheDir, m1.TarName) + assert.FileExists(t, tar1) + + // Change a file and rebuild: a new uniquely named tar replaces the old one, and the + // manifest points at the survivor. This is what keeps the manifest/tar pair + // consistent instead of overwriting a fixed name in place. + writeRepoFile(t, repo, "a.txt", "alpha-2") + files2 := statFiles(t, repo, "a.txt", "b.txt") + require.NoError(t, rebuildWarmSnapshot(repo, dirName, files2, m1, cacheDir, tar1, filepath.Join(t.TempDir(), "2.tar.gz"))) + m2 := loadSnapshotManifest(filepath.Join(cacheDir, snapshotCacheManifestName)) + require.NotNil(t, m2) + assert.NotEqual(t, m1.TarName, m2.TarName, "each build gets a unique tar name") + assert.FileExists(t, filepath.Join(cacheDir, m2.TarName)) + assert.NoFileExists(t, tar1, "superseded warm tar is cleaned up") + + matches, _ := filepath.Glob(filepath.Join(cacheDir, snapshotTarPrefix+"*.tar")) + assert.Len(t, matches, 1, "exactly one warm tar remains") +} From d40316bc522562a1f5de85e40575135c61163232 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Tue, 15 Sep 2026 17:47:23 +0000 Subject: [PATCH 7/7] experimental/air: upload incremental code snapshots --- experimental/air/cmd/runsubmit_test.go | 4 +- experimental/air/cmd/snapshot_cache.go | 33 +- experimental/air/cmd/snapshot_cache_test.go | 39 +- experimental/air/cmd/snapshot_cachekey.go | 11 +- .../air/cmd/snapshot_cachekey_test.go | 4 +- experimental/air/cmd/snapshot_dabs.go | 15 +- experimental/air/cmd/snapshot_overlay.go | 560 ++++++++++++++++++ experimental/air/cmd/snapshot_overlay_test.go | 217 +++++++ experimental/air/cmd/snapshot_package.go | 30 +- experimental/air/cmd/snapshot_resolve.go | 4 +- 10 files changed, 885 insertions(+), 32 deletions(-) create mode 100644 experimental/air/cmd/snapshot_overlay.go create mode 100644 experimental/air/cmd/snapshot_overlay_test.go diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 839264d8673..9613d7db2d4 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -378,7 +378,7 @@ func TestSubmitWorkloadStagingErrorPreventsSubmit(t *testing.T) { cfg, err := loadRunConfig(cfgPath) 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.Error(t, err) assert.ErrorContains(t, err, "failed to create launch directory") assert.Zero(t, submitCalls.Load()) @@ -439,7 +439,7 @@ environment: cfg, err := loadRunConfig(cfgPath) 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) tasks, ok := got["tasks"].([]any) diff --git a/experimental/air/cmd/snapshot_cache.go b/experimental/air/cmd/snapshot_cache.go index e319c84be90..5f1f6fdfe9c 100644 --- a/experimental/air/cmd/snapshot_cache.go +++ b/experimental/air/cmd/snapshot_cache.go @@ -3,7 +3,7 @@ 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) +// tar of the tree on local disk plus a manifest of each member's metadata identity // 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 @@ -33,7 +33,7 @@ import ( const ( // snapshotCacheVersion invalidates on-disk caches when the layout below changes. - snapshotCacheVersion = "v1" + snapshotCacheVersion = "v2" snapshotCacheManifestName = "manifest.json" // snapshotTarPrefix begins each warm tar's filename; the rest is a per-build random // id (snapshot..tar). A build never overwrites another build's tar, so a manifest @@ -51,14 +51,16 @@ const ( // embed a trailer between members). var tarTrailer = make([]byte, 2*512) -// cacheEntry records a member's identity for change detection (size+mtime) and its +// cacheEntry records a member's identity for change detection 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"` + Size int64 `json:"size"` + ModTime int64 `json:"mtime_ns"` + Mode uint32 `json:"mode"` + LinkTarget string `json:"link_target,omitempty"` + Offset int64 `json:"offset"` + Length int64 `json:"length"` } // snapshotManifest is the on-disk index of a warm snapshot tar. TarName binds it to the @@ -144,7 +146,7 @@ func packagePlainTarWithCache(ctx context.Context, repoPath, configPath string, } // 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 +// any add, delete, or metadata modification. 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) { @@ -152,7 +154,7 @@ func snapshotChanged(files []snapshotFile, old *snapshotManifest) bool { } for _, f := range files { e, ok := old.Entries[filepath.ToSlash(f.rel)] - if !ok || e.Size != f.size || e.ModTime != f.modTime { + if !ok || e.Size != f.size || e.ModTime != f.modTime || e.Mode != f.mode || e.LinkTarget != f.linkTarget { return true } } @@ -239,7 +241,7 @@ func writeGzOnly(repoPath, dirName string, files []snapshotFile, outputTarball s // 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) +// the tarDst stream. When old+oldTar are set, an unchanged member (matching metadata) // 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 @@ -260,7 +262,7 @@ func writeSnapshot(repoPath, dirName string, files []snapshotFile, old *snapshot reused := false if old != nil && oldTar != nil { - if e, ok := old.Entries[rel]; ok && e.Size == f.size && e.ModTime == f.modTime { + if e, ok := old.Entries[rel]; ok && e.Size == f.size && e.ModTime == f.modTime && e.Mode == f.mode && e.LinkTarget == f.linkTarget { 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) } @@ -272,7 +274,14 @@ func writeSnapshot(repoPath, dirName string, files []snapshotFile, old *snapshot return snapshotManifest{}, err } } - entries[rel] = cacheEntry{Size: f.size, ModTime: f.modTime, Offset: start, Length: cw.n - start} + entries[rel] = cacheEntry{ + Size: f.size, + ModTime: f.modTime, + Mode: f.mode, + LinkTarget: f.linkTarget, + Offset: start, + Length: cw.n - start, + } } if _, err := cw.Write(tarTrailer); err != nil { return snapshotManifest{}, err diff --git a/experimental/air/cmd/snapshot_cache_test.go b/experimental/air/cmd/snapshot_cache_test.go index 717eb62fdfb..ae0686708a2 100644 --- a/experimental/air/cmd/snapshot_cache_test.go +++ b/experimental/air/cmd/snapshot_cache_test.go @@ -5,6 +5,7 @@ import ( "compress/gzip" "os" "path/filepath" + "slices" "testing" "github.com/stretchr/testify/assert" @@ -39,9 +40,21 @@ 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))) + fullPath := filepath.Join(repo, filepath.FromSlash(rel)) + info, err := os.Lstat(fullPath) require.NoError(t, err) - files = append(files, snapshotFile{rel: filepath.FromSlash(rel), size: info.Size(), modTime: info.ModTime().UnixNano()}) + linkTarget := "" + if info.Mode()&os.ModeSymlink != 0 { + linkTarget, err = os.Readlink(fullPath) + require.NoError(t, err) + } + files = append(files, snapshotFile{ + rel: filepath.FromSlash(rel), + size: info.Size(), + modTime: info.ModTime().UnixNano(), + mode: uint32(info.Mode()), + linkTarget: linkTarget, + }) } return files } @@ -79,6 +92,26 @@ func TestSnapshotChanged(t *testing.T) { assert.True(t, snapshotChanged(added, old), "addition detected") } +func TestSnapshotChangedDetectsModeAndSymlinkTarget(t *testing.T) { + old := &snapshotManifest{Entries: map[string]cacheEntry{ + "run.sh": {Size: 1, ModTime: 100, Mode: 0o644}, + "model": {Size: 6, ModTime: 200, Mode: uint32(os.ModeSymlink | 0o777), LinkTarget: "v1.bin"}, + }} + unchanged := []snapshotFile{ + {rel: "run.sh", size: 1, modTime: 100, mode: 0o644}, + {rel: "model", size: 6, modTime: 200, mode: uint32(os.ModeSymlink | 0o777), linkTarget: "v1.bin"}, + } + assert.False(t, snapshotChanged(unchanged, old)) + + modeChanged := slices.Clone(unchanged) + modeChanged[0].mode = 0o755 + assert.True(t, snapshotChanged(modeChanged, old)) + + linkChanged := slices.Clone(unchanged) + linkChanged[1].linkTarget = "v2.bin" + assert.True(t, snapshotChanged(linkChanged, old)) +} + // warmTarPath returns the warm tar the cache's manifest currently points at. func warmTarPath(t *testing.T, cacheDir string) string { t.Helper() @@ -129,7 +162,7 @@ func TestWarmSnapshotReusesUnchangedFromCache(t *testing.T) { oldTarPath := filepath.Join(cacheDir, old.TarName) // 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 + // cached metadata: 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"))) diff --git a/experimental/air/cmd/snapshot_cachekey.go b/experimental/air/cmd/snapshot_cachekey.go index 4bebf2b7643..e3b782303f6 100644 --- a/experimental/air/cmd/snapshot_cachekey.go +++ b/experimental/air/cmd/snapshot_cachekey.go @@ -20,13 +20,13 @@ const snapshotPackagingVersion = "v1" // with a git_archive key) and lets us invalidate it if the fingerprint scheme changes. // computePlainTarKey also folds in the shared snapshotPackagingVersion, so a // packaging-logic bump invalidates both modes' keys. -const plainTarKeyVersion = "plaintar-v1" +const plainTarKeyVersion = "plaintar-v2" // computePlainTarKey returns a content-addressed key for a working-tree snapshot: the -// SHA-256 over every file's path, size and mtime (sorted for stability). An unchanged +// SHA-256 over every file's path, size, mtime, mode, and symlink target (sorted for stability). An unchanged // tree yields the same key, so an already-uploaded tarball can be reused instead of -// re-packaged and re-uploaded. The fingerprint is size+mtime, not content — the same -// trade-off DABs file-sync makes — so an edit preserving both size and mtime is not seen. +// re-packaged and re-uploaded. Regular-file content detection is size+mtime, not content — +// the same trade-off DABs file-sync makes — so an edit preserving both is not seen. func computePlainTarKey(files []snapshotFile) string { sorted := slices.Clone(files) slices.SortFunc(sorted, func(a, b snapshotFile) int { @@ -35,7 +35,8 @@ func computePlainTarKey(files []snapshotFile) string { h := sha256.New() for _, f := range sorted { - fmt.Fprintf(h, "%s\x00%d\x00%d\n", filepath.ToSlash(f.rel), f.size, f.modTime) + fmt.Fprintf(h, "%s\x00%d\x00%d\x00%d\x00%s\n", + filepath.ToSlash(f.rel), f.size, f.modTime, f.mode, f.linkTarget) } fmt.Fprintf(h, "%s\x00%s", plainTarKeyVersion, snapshotPackagingVersion) return hex.EncodeToString(h.Sum(nil)) diff --git a/experimental/air/cmd/snapshot_cachekey_test.go b/experimental/air/cmd/snapshot_cachekey_test.go index 2607bc8dd89..755e78030c7 100644 --- a/experimental/air/cmd/snapshot_cachekey_test.go +++ b/experimental/air/cmd/snapshot_cachekey_test.go @@ -72,7 +72,7 @@ func TestComputeSnapshotCacheKeyProperties(t *testing.T) { } // TestComputePlainTarKeyProperties pins the working-tree fingerprint behavior: it is -// order-independent and reacts to any change in a file's path, size, or mtime. +// order-independent and reacts to any change in a file's metadata identity. func TestComputePlainTarKeyProperties(t *testing.T) { base := []snapshotFile{ {rel: "a.txt", size: 10, modTime: 100}, @@ -89,6 +89,8 @@ func TestComputePlainTarKeyProperties(t *testing.T) { assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 11, modTime: 100}, base[1]})) assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 10, modTime: 101}, base[1]})) assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "renamed.txt", size: 10, modTime: 100}, base[1]})) + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 10, modTime: 100, mode: 0o755}, base[1]})) + assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey([]snapshotFile{{rel: "a.txt", size: 10, modTime: 100, linkTarget: "target"}, base[1]})) // Adding or dropping a file changes the key. assert.NotEqual(t, computePlainTarKey(base), computePlainTarKey(base[:1])) diff --git a/experimental/air/cmd/snapshot_dabs.go b/experimental/air/cmd/snapshot_dabs.go index 8c54a87fc90..7763526f561 100644 --- a/experimental/air/cmd/snapshot_dabs.go +++ b/experimental/air/cmd/snapshot_dabs.go @@ -131,7 +131,7 @@ func uploadSnapshotSidecars(ctx context.Context, sidecarStore filer.Filer, sidec // for plain_tar, the working-tree file listing used to build both the key and the tarball // (nil for git_archive, which lists nothing locally). The name is _.tar.gz, // keyed on (commit, include_paths, root_path subtree) for git_archive and on the -// working-tree fingerprint (path+size+mtime) for plain_tar, so an identical input +// working-tree metadata fingerprint for plain_tar, so an identical input // reuses the same remote object (see the skip in uploadSnapshotViaDABs). func snapshotTarName(ctx context.Context, repoPath string, plan snapshotPlan) (string, []snapshotFile, error) { dirName := filepath.Base(repoPath) @@ -139,10 +139,16 @@ func snapshotTarName(ctx context.Context, repoPath string, plan snapshotPlan) (s key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths, plan.subtreePrefix) return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]), nil, nil } + start := time.Now() files, err := snapshotFiles(ctx, repoPath, plan.includePaths, plan.isGitRepo) if err != nil { return "", nil, err } + var totalBytes int64 + for _, file := range files { + totalBytes += file.size + } + log.Debugf(ctx, "air snapshot listing: files=%d uncompressed_bytes=%d duration=%s", len(files), totalBytes, time.Since(start)) return fmt.Sprintf("%s_%s.tar.gz", dirName, computePlainTarKey(files)[:16]), files, nil } @@ -166,7 +172,7 @@ func packageSnapshot(ctx context.Context, repoPath, configPath string, plan snap // already-resolved user's repo_snapshots directory or a configured UC Volume. // // The tarball name is content-addressed — by (commit, include_paths, root_path subtree) -// for git_archive and by the working-tree fingerprint (path+size+mtime) for plain_tar — +// for git_archive and by the working-tree metadata fingerprint for plain_tar — // so if the identical object is already uploaded we skip packaging and upload entirely // and reuse the remote path. func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath, configPath string, plan snapshotPlan, artifactPath string, noCache bool) (snapshotResult, error) { @@ -219,6 +225,11 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r if diags.HasError() { return snapshotResult{}, diags.Error() } + if result, handled, err := maybeUploadSnapshotOverlay( + ctx, f, uploadPath, repoPath, configPath, plan, files, tmp, artifactPath, noCache, + ); handled { + return result, err + } exists, err := snapshotExists(ctx, f, tarName) if err != nil { return snapshotResult{}, err diff --git a/experimental/air/cmd/snapshot_overlay.go b/experimental/air/cmd/snapshot_overlay.go new file mode 100644 index 00000000000..873aa43ceae --- /dev/null +++ b/experimental/air/cmd/snapshot_overlay.go @@ -0,0 +1,560 @@ +package aircmd + +import ( + "archive/tar" + "context" + "crypto/sha256" + "encoding/hex" + "encoding/json" + "errors" + "fmt" + "io" + "io/fs" + "os" + "path" + "path/filepath" + "slices" + "strings" + "time" + + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/log" + "github.com/klauspost/pgzip" +) + +const ( + snapshotOverlayVersion = "v1" + snapshotOverlayDescriptorName = ".databricks/air-source-overlay-v1/descriptor" + snapshotOverlayDeletionsRoot = ".databricks/air-source-overlay-v1/deletions" + + snapshotAnchorMaxAge = 30 * 24 * time.Hour + snapshotOverlayMaxChanged = 1024 + snapshotOverlayMaxBytes int64 = 128 << 20 + snapshotOverlayMaxRatio = 0.25 +) + +type snapshotOverlayEntry struct { + Size int64 `json:"size"` + ModTime int64 `json:"mtime_ns"` + Mode uint32 `json:"mode"` + LinkTarget string `json:"link_target,omitempty"` +} + +type snapshotAnchorState struct { + Version string `json:"version"` + Branch string `json:"branch"` + CreatedAtUnix int64 `json:"created_at_unix"` + AnchorName string `json:"anchor_name"` + AnchorFingerprint string `json:"anchor_fingerprint"` + Component string `json:"component"` + Entries map[string]snapshotOverlayEntry `json:"entries"` + Results map[string]string `json:"results"` +} + +type snapshotOverlayDescriptor struct { + FormatVersion int `json:"format_version"` + Anchor string `json:"anchor"` + AnchorFingerprint string `json:"anchor_fingerprint"` + ResultFingerprint string `json:"result_fingerprint"` + Component string `json:"component"` +} + +type snapshotOverlayDiff struct { + Changed []snapshotFile + Deleted []string + ChangedBytes int64 + TotalBytes int64 +} + +func snapshotOverlayManifest(files []snapshotFile) map[string]snapshotOverlayEntry { + entries := make(map[string]snapshotOverlayEntry, len(files)) + for _, file := range files { + entries[filepath.ToSlash(file.rel)] = snapshotOverlayEntry{ + Size: file.size, + ModTime: file.modTime, + Mode: file.mode, + LinkTarget: file.linkTarget, + } + } + return entries +} + +func snapshotOverlayFingerprint(files []snapshotFile) string { + return computePlainTarKey(files) +} + +func computeSnapshotOverlayDiff(files []snapshotFile, anchor map[string]snapshotOverlayEntry) snapshotOverlayDiff { + current := make(map[string]struct{}, len(files)) + diff := snapshotOverlayDiff{} + for _, file := range files { + rel := filepath.ToSlash(file.rel) + current[rel] = struct{}{} + diff.TotalBytes += file.size + entry, ok := anchor[rel] + if !ok || entry.Size != file.size || entry.ModTime != file.modTime || + entry.Mode != file.mode || entry.LinkTarget != file.linkTarget { + diff.Changed = append(diff.Changed, file) + diff.ChangedBytes += file.size + } + } + for rel := range anchor { + if _, ok := current[rel]; !ok { + diff.Deleted = append(diff.Deleted, rel) + } + } + slices.SortFunc(diff.Changed, func(a, b snapshotFile) int { + return strings.Compare(filepath.ToSlash(a.rel), filepath.ToSlash(b.rel)) + }) + slices.Sort(diff.Deleted) + return diff +} + +func snapshotOverlayAnchorReason(diff snapshotOverlayDiff) string { + changedPaths := len(diff.Changed) + len(diff.Deleted) + if changedPaths > snapshotOverlayMaxChanged { + return "changed-path-count" + } + if diff.ChangedBytes > snapshotOverlayMaxBytes { + return "changed-bytes" + } + if diff.TotalBytes > 0 && float64(diff.ChangedBytes)/float64(diff.TotalBytes) > snapshotOverlayMaxRatio { + return "changed-byte-ratio" + } + return "" +} + +func snapshotOverlayStatePath(repoPath, configPath string, includePaths []string, branch string) (string, error) { + absRepo, err := filepath.Abs(repoPath) + if err != nil { + return "", err + } + absConfig, err := filepath.Abs(configPath) + if err != nil { + return "", err + } + branchKey := sha256.Sum256([]byte(snapshotOverlayVersion + "\x00" + branch)) + return filepath.Join( + snapshotCacheDir(absRepo, absConfig, includePaths), + "overlay-anchor-"+hex.EncodeToString(branchKey[:8])+".json", + ), nil +} + +func loadSnapshotAnchorState(statePath, branch, component string) *snapshotAnchorState { + file, err := os.Open(statePath) + if err != nil { + return nil + } + defer file.Close() + decoder := json.NewDecoder(io.LimitReader(file, 64<<20)) + decoder.DisallowUnknownFields() + var state snapshotAnchorState + if err := decoder.Decode(&state); err != nil { + return nil + } + var extra any + if err := decoder.Decode(&extra); err != io.EOF || state.Version != snapshotOverlayVersion || + state.Branch != branch || state.Component != component || state.CreatedAtUnix <= 0 || + state.CreatedAtUnix > time.Now().Add(time.Hour).Unix() || + state.AnchorName == "" || state.AnchorFingerprint == "" || state.Entries == nil || state.Results == nil { + return nil + } + if !validSnapshotObjectNameForKind(state.AnchorName, "anchor") || !validSnapshotFingerprint(state.AnchorFingerprint) { + return nil + } + anchorFiles := make([]snapshotFile, 0, len(state.Entries)) + for rel, entry := range state.Entries { + if validateSnapshotRelativePath(rel) != nil { + return nil + } + mode := os.FileMode(entry.Mode) + if entry.Size < 0 || (!mode.IsRegular() && mode&os.ModeSymlink == 0) || + (mode&os.ModeSymlink == 0 && entry.LinkTarget != "") { + return nil + } + anchorFiles = append(anchorFiles, snapshotFile{ + rel: filepath.FromSlash(rel), + size: entry.Size, + modTime: entry.ModTime, + mode: entry.Mode, + linkTarget: entry.LinkTarget, + }) + } + if snapshotOverlayFingerprint(anchorFiles) != state.AnchorFingerprint || state.Results[state.AnchorFingerprint] != state.AnchorName { + return nil + } + for fingerprint, name := range state.Results { + if !validSnapshotFingerprint(fingerprint) || !validSnapshotObjectName(name) { + return nil + } + } + return &state +} + +func validSnapshotObjectName(name string) bool { + return validSnapshotObjectNameForKind(name, "anchor") || validSnapshotObjectNameForKind(name, "overlay") +} + +func validSnapshotObjectNameForKind(name, kind string) bool { + if len(name) > 255 || path.Base(name) != name || strings.ContainsAny(name, "/\\") || !strings.HasSuffix(name, ".tar.gz") { + return false + } + prefix := "air_" + kind + "_" + snapshotOverlayVersion + "_" + stem := strings.TrimSuffix(name, ".tar.gz") + return strings.HasPrefix(stem, prefix) && validSnapshotFingerprint(strings.TrimPrefix(stem, prefix)) +} + +func validSnapshotFingerprint(fingerprint string) bool { + if len(fingerprint) != sha256.Size*2 { + return false + } + _, err := hex.DecodeString(fingerprint) + return err == nil +} + +func saveSnapshotAnchorState(statePath string, state snapshotAnchorState) error { + data, err := json.Marshal(state) + if err != nil { + return err + } + if err := os.MkdirAll(filepath.Dir(statePath), 0o700); err != nil { + return fmt.Errorf("failed to create overlay state directory: %w", err) + } + tmp, err := os.CreateTemp(filepath.Dir(statePath), "overlay-anchor-*.json") + if err != nil { + return fmt.Errorf("failed to write overlay state: %w", err) + } + tmpName := tmp.Name() + if _, err := tmp.Write(data); err != nil { + tmp.Close() + os.Remove(tmpName) + return fmt.Errorf("failed to write overlay state: %w", err) + } + if err := tmp.Close(); err != nil { + os.Remove(tmpName) + return fmt.Errorf("failed to write overlay state: %w", err) + } + if err := os.Rename(tmpName, statePath); err != nil { + os.Remove(tmpName) + return fmt.Errorf("failed to install overlay state: %w", err) + } + return nil +} + +// maybeUploadSnapshotOverlay handles the WSFS incremental path. A false +// handled result leaves the caller on the existing full-snapshot path. +func maybeUploadSnapshotOverlay( + ctx context.Context, + store filer.Filer, + uploadPath string, + repoPath string, + configPath string, + plan snapshotPlan, + files []snapshotFile, + tmp string, + artifactPath string, + noCache bool, +) (result snapshotResult, handled bool, err error) { + isVolumePath := artifactPath == "/Volumes" || strings.HasPrefix(artifactPath, "/Volumes/") + if noCache || isVolumePath || plan.mode != modePlainTar { + return snapshotResult{}, false, nil + } + + var totalBytes int64 + for _, file := range files { + totalBytes += file.size + } + if totalBytes < snapshotCacheMinBytes || filepath.Base(repoPath) == ".databricks" { + log.Debugf(ctx, "air snapshot overlay: decision=full-existing reason=small-or-reserved files=%d uncompressed_bytes=%d", len(files), totalBytes) + return snapshotResult{}, false, nil + } + + branch := "" + if plan.isGitRepo { + branch = newGitRepo(repoPath).currentBranch(ctx) + } + component := filepath.Base(repoPath) + statePath, err := snapshotOverlayStatePath(repoPath, configPath, plan.includePaths, branch) + if err != nil { + return snapshotResult{}, true, err + } + fingerprint := snapshotOverlayFingerprint(files) + state := loadSnapshotAnchorState(statePath, branch, component) + anchorReason := "" + if state == nil { + anchorReason = "missing-or-invalid-state" + } else if time.Since(time.Unix(state.CreatedAtUnix, 0)) >= snapshotAnchorMaxAge { + anchorReason = "anchor-age" + } else { + exists, statErr := snapshotExists(ctx, store, state.AnchorName) + if statErr != nil { + return snapshotResult{}, true, statErr + } + if !exists { + anchorReason = "missing-remote-anchor" + } + } + + if anchorReason == "" { + if name, ok := state.Results[fingerprint]; ok { + exists, statErr := snapshotExists(ctx, store, name) + if statErr != nil { + return snapshotResult{}, true, statErr + } + if exists { + log.Debugf(ctx, "air snapshot overlay: decision=reuse branch=%q object=%s", branch, name) + return snapshotResult{CodeSourcePath: path.Join(uploadPath, name)}, true, nil + } + delete(state.Results, fingerprint) + } + + diff := computeSnapshotOverlayDiff(files, state.Entries) + anchorReason = snapshotOverlayAnchorReason(diff) + changedPaths := len(diff.Changed) + len(diff.Deleted) + ratio := 0.0 + if diff.TotalBytes > 0 { + ratio = float64(diff.ChangedBytes) / float64(diff.TotalBytes) + } + log.Debugf(ctx, "air snapshot overlay: decision-input branch=%q changed_paths=%d changed_bytes=%d total_bytes=%d changed_ratio=%.4f reanchor_reason=%q", + branch, changedPaths, diff.ChangedBytes, diff.TotalBytes, ratio, anchorReason) + if anchorReason == "" { + archivePath := filepath.Join(tmp, "overlay.tar.gz") + packageStart := time.Now() + if err := createSnapshotOverlay(repoPath, component, state, fingerprint, diff, archivePath); err != nil { + return snapshotResult{}, true, err + } + afterFiles, err := snapshotFiles(ctx, repoPath, plan.includePaths, plan.isGitRepo) + if err != nil { + return snapshotResult{}, true, err + } + if snapshotOverlayFingerprint(afterFiles) != fingerprint { + return snapshotResult{}, true, errors.New("code source changed while creating its overlay snapshot; retry submission") + } + packageDuration := time.Since(packageStart) + archiveInfo, err := os.Stat(archivePath) + if err != nil { + return snapshotResult{}, true, err + } + publishStart := time.Now() + name, err := publishImmutableSnapshot(ctx, store, "air_overlay_"+snapshotOverlayVersion, archivePath) + if err != nil { + return snapshotResult{}, true, err + } + log.Debugf(ctx, "air snapshot overlay: decision=overlay package=%s publish=%s compressed_bytes=%d object=%s", + packageDuration, time.Since(publishStart), archiveInfo.Size(), name) + state.Results[fingerprint] = name + if err := saveSnapshotAnchorState(statePath, *state); err != nil { + return snapshotResult{}, true, err + } + return snapshotResult{CodeSourcePath: path.Join(uploadPath, name)}, true, nil + } + } + + log.Debugf(ctx, "air snapshot overlay: decision=new-anchor branch=%q reason=%s files=%d uncompressed_bytes=%d", branch, anchorReason, len(files), totalBytes) + archivePath := filepath.Join(tmp, "anchor.tar.gz") + packageStart := time.Now() + if err := packageSnapshot(ctx, repoPath, configPath, plan, files, archivePath, false); err != nil { + return snapshotResult{}, true, err + } + afterFiles, err := snapshotFiles(ctx, repoPath, plan.includePaths, plan.isGitRepo) + if err != nil { + return snapshotResult{}, true, err + } + if snapshotOverlayFingerprint(afterFiles) != fingerprint { + return snapshotResult{}, true, errors.New("code source changed while creating its anchor snapshot; retry submission") + } + packageDuration := time.Since(packageStart) + archiveInfo, err := os.Stat(archivePath) + if err != nil { + return snapshotResult{}, true, err + } + publishStart := time.Now() + name, err := publishImmutableSnapshot(ctx, store, "air_anchor_"+snapshotOverlayVersion, archivePath) + if err != nil { + return snapshotResult{}, true, err + } + log.Debugf(ctx, "air snapshot overlay: decision=anchor package=%s publish=%s compressed_bytes=%d object=%s", + packageDuration, time.Since(publishStart), archiveInfo.Size(), name) + newState := snapshotAnchorState{ + Version: snapshotOverlayVersion, + Branch: branch, + CreatedAtUnix: time.Now().Unix(), + AnchorName: name, + AnchorFingerprint: fingerprint, + Component: component, + Entries: snapshotOverlayManifest(files), + Results: map[string]string{fingerprint: name}, + } + if err := saveSnapshotAnchorState(statePath, newState); err != nil { + return snapshotResult{}, true, err + } + return snapshotResult{CodeSourcePath: path.Join(uploadPath, name)}, true, nil +} + +func createSnapshotOverlay( + repoPath string, + component string, + state *snapshotAnchorState, + resultFingerprint string, + diff snapshotOverlayDiff, + outputPath string, +) (err error) { + output, err := os.Create(outputPath) + if err != nil { + return fmt.Errorf("failed to create overlay: %w", err) + } + gz, err := pgzip.NewWriterLevel(output, pgzip.DefaultCompression) + if err != nil { + output.Close() + return err + } + tw := tar.NewWriter(gz) + defer func() { + err = firstErr(err, tw.Close(), gz.Close(), output.Close()) + }() + + descriptor, err := json.Marshal(snapshotOverlayDescriptor{ + FormatVersion: 1, + Anchor: state.AnchorName, + AnchorFingerprint: state.AnchorFingerprint, + ResultFingerprint: resultFingerprint, + Component: component, + }) + if err != nil { + return err + } + if err := tw.WriteHeader(&tar.Header{ + Name: snapshotOverlayDescriptorName, + Mode: 0o600, + Size: int64(len(descriptor)), + }); err != nil { + return fmt.Errorf("failed to write overlay descriptor: %w", err) + } + if _, err := tw.Write(descriptor); err != nil { + return fmt.Errorf("failed to write overlay descriptor: %w", err) + } + + for _, rel := range diff.Deleted { + if err := validateSnapshotRelativePath(rel); err != nil { + return err + } + if err := tw.WriteHeader(&tar.Header{ + Name: path.Join(snapshotOverlayDeletionsRoot, rel), + Mode: 0o600, + }); err != nil { + return fmt.Errorf("failed to write deletion marker for %q: %w", rel, err) + } + } + for _, file := range diff.Changed { + if err := writeSnapshotOverlayMember(tw, repoPath, component, file); err != nil { + return err + } + } + return nil +} + +func validateSnapshotRelativePath(rel string) error { + clean := path.Clean(filepath.ToSlash(rel)) + if clean == "." || clean == ".." || strings.HasPrefix(clean, "../") || path.IsAbs(clean) || clean != filepath.ToSlash(rel) { + return fmt.Errorf("invalid snapshot path %q", rel) + } + return nil +} + +func writeSnapshotOverlayMember(tw *tar.Writer, repoPath, component string, file snapshotFile) error { + rel := filepath.ToSlash(file.rel) + if err := validateSnapshotRelativePath(rel); err != nil { + return err + } + fullPath := filepath.Join(repoPath, file.rel) + info, linkTarget, err := inspectSnapshotOverlayFile(fullPath) + if err != nil { + return fmt.Errorf("failed to inspect changed snapshot file %q: %w", rel, err) + } + if !snapshotFileMatches(file, info, linkTarget) { + return fmt.Errorf("snapshot file changed while packaging: %q", rel) + } + if !info.Mode().IsRegular() && info.Mode()&os.ModeSymlink == 0 { + return nil + } + header, err := tar.FileInfoHeader(info, linkTarget) + if err != nil { + return fmt.Errorf("failed to build overlay header for %q: %w", rel, err) + } + header.Name = path.Join(component, rel) + if err := tw.WriteHeader(header); err != nil { + return fmt.Errorf("failed to write overlay header for %q: %w", rel, err) + } + if info.Mode().IsRegular() { + input, err := os.Open(fullPath) + if err != nil { + return fmt.Errorf("failed to open changed snapshot file %q: %w", rel, err) + } + _, copyErr := io.Copy(tw, input) + closeErr := input.Close() + if copyErr != nil { + return fmt.Errorf("failed to write changed snapshot file %q: %w", rel, copyErr) + } + if closeErr != nil { + return fmt.Errorf("failed to close changed snapshot file %q: %w", rel, closeErr) + } + } + after, afterLink, err := inspectSnapshotOverlayFile(fullPath) + if err != nil || !snapshotFileMatches(file, after, afterLink) { + return fmt.Errorf("snapshot file changed while packaging: %q", rel) + } + return nil +} + +func inspectSnapshotOverlayFile(name string) (os.FileInfo, string, error) { + info, err := os.Lstat(name) + if err != nil { + return nil, "", err + } + linkTarget := "" + if info.Mode()&os.ModeSymlink != 0 { + linkTarget, err = os.Readlink(name) + } + return info, linkTarget, err +} + +func snapshotFileMatches(file snapshotFile, info os.FileInfo, linkTarget string) bool { + return info != nil && file.size == info.Size() && file.modTime == info.ModTime().UnixNano() && + file.mode == uint32(info.Mode()) && file.linkTarget == linkTarget +} + +func publishImmutableSnapshot(ctx context.Context, store filer.Filer, kind, archivePath string) (string, error) { + input, err := os.Open(archivePath) + if err != nil { + return "", err + } + hash := sha256.New() + if _, err := io.Copy(hash, input); err != nil { + input.Close() + return "", fmt.Errorf("failed to hash snapshot archive: %w", err) + } + if _, err := input.Seek(0, io.SeekStart); err != nil { + input.Close() + return "", err + } + name := kind + "_" + hex.EncodeToString(hash.Sum(nil)) + ".tar.gz" + err = store.Write(ctx, name, input, filer.CreateParentDirectories) + closeErr := input.Close() + if err != nil && !errors.Is(err, fs.ErrExist) { + return "", fmt.Errorf("failed to publish immutable snapshot %s: %w", name, err) + } + if closeErr != nil { + return "", closeErr + } + localInfo, err := os.Stat(archivePath) + if err != nil { + return "", err + } + remoteInfo, err := store.Stat(ctx, name) + if err != nil { + return "", fmt.Errorf("failed to verify immutable snapshot %s: %w", name, err) + } + if remoteInfo.Size() != localInfo.Size() { + return "", fmt.Errorf("immutable snapshot %s has unexpected size: got %d, want %d", name, remoteInfo.Size(), localInfo.Size()) + } + return name, nil +} diff --git a/experimental/air/cmd/snapshot_overlay_test.go b/experimental/air/cmd/snapshot_overlay_test.go new file mode 100644 index 00000000000..3118e836553 --- /dev/null +++ b/experimental/air/cmd/snapshot_overlay_test.go @@ -0,0 +1,217 @@ +package aircmd + +import ( + "archive/tar" + "compress/gzip" + "encoding/json" + "io" + "os" + "path/filepath" + "testing" + "time" + + "github.com/databricks/cli/libs/filer" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func overlayTestFile(t *testing.T, repo, rel string) snapshotFile { + t.Helper() + info, err := os.Lstat(filepath.Join(repo, filepath.FromSlash(rel))) + require.NoError(t, err) + linkTarget := "" + if info.Mode()&os.ModeSymlink != 0 { + linkTarget, err = os.Readlink(filepath.Join(repo, filepath.FromSlash(rel))) + require.NoError(t, err) + } + return snapshotFile{ + rel: filepath.FromSlash(rel), + size: info.Size(), + modTime: info.ModTime().UnixNano(), + mode: uint32(info.Mode()), + linkTarget: linkTarget, + } +} + +func TestComputeSnapshotOverlayDiff(t *testing.T) { + anchor := map[string]snapshotOverlayEntry{ + "same": {Size: 1, ModTime: 1, Mode: 0o644}, + "mode": {Size: 2, ModTime: 2, Mode: 0o644}, + "link": {Size: 3, ModTime: 3, Mode: uint32(os.ModeSymlink | 0o777), LinkTarget: "old"}, + "deleted": {Size: 10, ModTime: 4, Mode: 0o644}, + } + files := []snapshotFile{ + {rel: "same", size: 1, modTime: 1, mode: 0o644}, + {rel: "mode", size: 2, modTime: 2, mode: 0o755}, + {rel: "link", size: 3, modTime: 3, mode: uint32(os.ModeSymlink | 0o777), linkTarget: "new"}, + {rel: "added", size: 4, modTime: 5, mode: 0o600}, + } + + diff := computeSnapshotOverlayDiff(files, anchor) + + require.Len(t, diff.Changed, 3) + assert.Equal(t, []string{"added", "link", "mode"}, []string{ + filepath.ToSlash(diff.Changed[0].rel), + filepath.ToSlash(diff.Changed[1].rel), + filepath.ToSlash(diff.Changed[2].rel), + }) + assert.Equal(t, []string{"deleted"}, diff.Deleted) + assert.Equal(t, int64(9), diff.ChangedBytes) + assert.Equal(t, int64(10), diff.TotalBytes) +} + +func TestSnapshotOverlayAnchorReason(t *testing.T) { + assert.Empty(t, snapshotOverlayAnchorReason(snapshotOverlayDiff{ChangedBytes: 1, TotalBytes: 100})) + assert.Equal(t, "changed-path-count", snapshotOverlayAnchorReason(snapshotOverlayDiff{ + Changed: make([]snapshotFile, snapshotOverlayMaxChanged+1), + })) + assert.Equal(t, "changed-bytes", snapshotOverlayAnchorReason(snapshotOverlayDiff{ + ChangedBytes: snapshotOverlayMaxBytes + 1, + TotalBytes: snapshotOverlayMaxBytes * 10, + })) + assert.Equal(t, "changed-byte-ratio", snapshotOverlayAnchorReason(snapshotOverlayDiff{ + ChangedBytes: 26, + TotalBytes: 100, + })) +} + +func TestCreateSnapshotOverlay(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "changed.sh", "#!/bin/sh\necho changed\n") + require.NoError(t, os.Chmod(filepath.Join(repo, "changed.sh"), 0o755)) + writeRepoFile(t, repo, "added.txt", "added") + require.NoError(t, os.Symlink("added.txt", filepath.Join(repo, "current"))) + files := []snapshotFile{ + overlayTestFile(t, repo, "changed.sh"), + overlayTestFile(t, repo, "added.txt"), + overlayTestFile(t, repo, "current"), + } + state := &snapshotAnchorState{ + AnchorName: "air_anchor_v1_aabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccdd.tar.gz", + AnchorFingerprint: "aabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccddaabbccdd", + } + out := filepath.Join(t.TempDir(), "overlay.tar.gz") + require.NoError(t, createSnapshotOverlay(repo, filepath.Base(repo), state, "b"+state.AnchorFingerprint[1:], snapshotOverlayDiff{ + Changed: files, + Deleted: []string{"deleted.txt"}, + }, out)) + + archive, err := os.Open(out) + require.NoError(t, err) + defer archive.Close() + gz, err := gzip.NewReader(archive) + require.NoError(t, err) + defer gz.Close() + reader := tar.NewReader(gz) + + first, err := reader.Next() + require.NoError(t, err) + assert.Equal(t, snapshotOverlayDescriptorName, first.Name) + descriptorBytes, err := io.ReadAll(reader) + require.NoError(t, err) + var descriptor snapshotOverlayDescriptor + require.NoError(t, json.Unmarshal(descriptorBytes, &descriptor)) + assert.Equal(t, state.AnchorName, descriptor.Anchor) + assert.Equal(t, filepath.Base(repo), descriptor.Component) + + entries := map[string]*tar.Header{} + for { + header, err := reader.Next() + if err == io.EOF { + break + } + require.NoError(t, err) + copy := *header + entries[header.Name] = © + } + assert.Contains(t, entries, snapshotOverlayDeletionsRoot+"/deleted.txt") + assert.Equal(t, int64(0), entries[snapshotOverlayDeletionsRoot+"/deleted.txt"].Size) + assert.Equal(t, int64(0o755), entries[filepath.Base(repo)+"/changed.sh"].Mode) + assert.Equal(t, byte(tar.TypeSymlink), entries[filepath.Base(repo)+"/current"].Typeflag) + assert.Equal(t, "added.txt", entries[filepath.Base(repo)+"/current"].Linkname) +} + +func TestSnapshotAnchorStateIsBranchLocalAndStrict(t *testing.T) { + repo := t.TempDir() + config := filepath.Join(repo, "run.yaml") + mainPath, err := snapshotOverlayStatePath(repo, config, []string{"src"}, "main") + require.NoError(t, err) + featurePath, err := snapshotOverlayStatePath(repo, config, []string{"src"}, "feature") + require.NoError(t, err) + assert.NotEqual(t, mainPath, featurePath) + + entries := map[string]snapshotOverlayEntry{"src/train.py": {Size: 1}} + anchorFingerprint := snapshotOverlayFingerprint([]snapshotFile{{rel: filepath.FromSlash("src/train.py"), size: 1}}) + anchorName := "air_anchor_v1_aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa.tar.gz" + state := snapshotAnchorState{ + Version: snapshotOverlayVersion, + Branch: "main", + CreatedAtUnix: time.Now().Unix(), + AnchorName: anchorName, + AnchorFingerprint: anchorFingerprint, + Component: filepath.Base(repo), + Entries: entries, + Results: map[string]string{anchorFingerprint: anchorName}, + } + require.NoError(t, saveSnapshotAnchorState(mainPath, state)) + assert.NotNil(t, loadSnapshotAnchorState(mainPath, "main", filepath.Base(repo))) + assert.Nil(t, loadSnapshotAnchorState(mainPath, "other", filepath.Base(repo))) + + require.NoError(t, os.WriteFile(mainPath, append([]byte(`{"unknown":true}`), '\n'), 0o600)) + assert.Nil(t, loadSnapshotAnchorState(mainPath, "main", filepath.Base(repo))) +} + +func TestMaybeUploadSnapshotOverlayLifecycle(t *testing.T) { + repo := t.TempDir() + largePath := filepath.Join(repo, "large.bin") + large, err := os.Create(largePath) + require.NoError(t, err) + require.NoError(t, large.Truncate(snapshotCacheMinBytes+1)) + require.NoError(t, large.Close()) + writeRepoFile(t, repo, "train.py", "print('one')") + configPath := filepath.Join(repo, "run.yaml") + require.NoError(t, os.WriteFile(configPath, []byte("test"), 0o600)) + + remoteDir := t.TempDir() + store, err := filer.NewLocalClient(remoteDir) + require.NoError(t, err) + plan := snapshotPlan{mode: modePlainTar, isGitRepo: false} + files, err := snapshotFiles(t.Context(), repo, nil, false) + require.NoError(t, err) + statePath, err := snapshotOverlayStatePath(repo, configPath, nil, "") + require.NoError(t, err) + t.Cleanup(func() { _ = os.RemoveAll(filepath.Dir(statePath)) }) + + anchor, handled, err := maybeUploadSnapshotOverlay( + t.Context(), store, "/Workspace/test/.internal", repo, configPath, plan, files, + t.TempDir(), "", false, + ) + require.NoError(t, err) + require.True(t, handled) + assert.Contains(t, filepath.Base(anchor.CodeSourcePath), "air_anchor_v1_") + + oldInfo, err := os.Stat(filepath.Join(repo, "train.py")) + require.NoError(t, err) + writeRepoFile(t, repo, "train.py", "print('two')") + require.NoError(t, os.Chtimes( + filepath.Join(repo, "train.py"), oldInfo.ModTime().Add(time.Second), oldInfo.ModTime().Add(time.Second), + )) + files, err = snapshotFiles(t.Context(), repo, nil, false) + require.NoError(t, err) + overlay, handled, err := maybeUploadSnapshotOverlay( + t.Context(), store, "/Workspace/test/.internal", repo, configPath, plan, files, + t.TempDir(), "", false, + ) + require.NoError(t, err) + require.True(t, handled) + assert.Contains(t, filepath.Base(overlay.CodeSourcePath), "air_overlay_v1_") + assert.NotEqual(t, anchor.CodeSourcePath, overlay.CodeSourcePath) + + reused, handled, err := maybeUploadSnapshotOverlay( + t.Context(), store, "/Workspace/test/.internal", repo, configPath, plan, files, + t.TempDir(), "", false, + ) + require.NoError(t, err) + require.True(t, handled) + assert.Equal(t, overlay.CodeSourcePath, reused.CodeSourcePath) +} diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index cce694d5ddb..e13417011af 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -91,11 +91,13 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, fil } // snapshotFile is a file selected for the snapshot: its repo-relative path (native -// separators) plus the size and mtime used by content addressing and the warm cache. +// separators) plus metadata used by content addressing, overlays, and the warm cache. type snapshotFile struct { - rel string - size int64 - modTime int64 // Unix nanoseconds + rel string + size int64 + modTime int64 // Unix nanoseconds + mode uint32 + linkTarget string } func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, isGitRepo bool) ([]snapshotFile, error) { @@ -139,7 +141,25 @@ func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, if err != nil { return nil, fmt.Errorf("failed to inspect snapshot path %q: %w", name, err) } - files = append(files, snapshotFile{rel: filepath.FromSlash(name), size: info.Size(), modTime: info.ModTime().UnixNano()}) + // tarpack omits directories and other special files. Filter them from the + // manifest too, so an overlay describes exactly the full archive's contents. + if !info.Mode().IsRegular() && info.Mode()&os.ModeSymlink == 0 { + continue + } + linkTarget := "" + if info.Mode()&os.ModeSymlink != 0 { + linkTarget, err = os.Readlink(filepath.Join(repoPath, filepath.FromSlash(name))) + if err != nil { + return nil, fmt.Errorf("failed to inspect snapshot symlink %q: %w", name, err) + } + } + files = append(files, snapshotFile{ + rel: filepath.FromSlash(name), + size: info.Size(), + modTime: info.ModTime().UnixNano(), + mode: uint32(info.Mode()), + linkTarget: linkTarget, + }) } return files, nil } diff --git a/experimental/air/cmd/snapshot_resolve.go b/experimental/air/cmd/snapshot_resolve.go index 0fcd3e82657..df465fdb505 100644 --- a/experimental/air/cmd/snapshot_resolve.go +++ b/experimental/air/cmd/snapshot_resolve.go @@ -20,8 +20,8 @@ const ( // root_path subtree). modeGitArchive snapshotMode = iota // modePlainTar packages the working tree (including uncommitted changes) via - // `tar`. Content-addressed by the working-tree fingerprint (path+size+mtime), so an - // unchanged tree reuses the uploaded tarball; a same-size, same-mtime edit is missed. + // `tar`. Content-addressed by the working-tree metadata fingerprint, so an + // unchanged tree reuses the uploaded tarball; a same-size, same-mtime regular-file edit is missed. modePlainTar )