diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 45560157bc2..7a619db25f1 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -391,15 +391,29 @@ func testSidecarStore(t *testing.T, w *databricks.WorkspaceClient) (filer.Filer, return f, base } -// 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} }) + // 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/") { + snapshotUploads++ + } + 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"}) @@ -420,14 +434,24 @@ 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, sidecarStore, sidecarBase) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, 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, 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, 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.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 44c58ee903b..6ad5a527119 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,31 @@ 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. +// 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 +// 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.Fprintf(h, "%s\x00%s", plainTarKeyVersion, snapshotPackagingVersion) + 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). // Changing any input yields a different entry. diff --git a/experimental/air/cmd/snapshot_cachekey_test.go b/experimental/air/cmd/snapshot_cachekey_test.go index 5743217c003..6737ce87eed 100644 --- a/experimental/air/cmd/snapshot_cachekey_test.go +++ b/experimental/air/cmd/snapshot_cachekey_test.go @@ -60,3 +60,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 ed809b09e01..1c0aed1be8e 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) — so -// an identical commit 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 { +// 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) 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) - return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]) + 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) +// 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, dirName, plan.includePaths) + return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, filepath.Base(repoPath), plan.includePaths) } - return createPlainTarball(ctx, repoPath, tarball, plan.includePaths, plan.isGitRepo) + 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. When remoteVolume is set the // tarball goes to that UC Volume; otherwise to the user's repo_snapshots dir. // -// git_archive snapshots are cacheable: the tarball name is content-addressed by -// (commit, include_paths), so if the identical object is already uploaded we skip -// packaging and upload entirely and just reuse the remote path. +// The tarball name is content-addressed — by (commit, include_paths) 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, remoteVolume 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. @@ -179,7 +186,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, @@ -210,33 +220,31 @@ func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, r return snapshotResult{}, err } - // git_archive is cacheable by (commit, include_paths): 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() { + // 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 } @@ -256,8 +264,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_package.go b/experimental/air/cmd/snapshot_package.go index c22f7dd61af..3f5cd474f74 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -34,12 +34,12 @@ 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 via `tar`. 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 { +// top-level 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) // Absolute so it resolves correctly regardless of tar's working dir (set below). parent, err := filepath.Abs(filepath.Dir(repoPath)) @@ -58,18 +58,13 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc } outName := filepath.Base(outputTarball) - files, err := snapshotFiles(ctx, repoPath, includePaths, isGitRepo) - if err != nil { - return err - } args := []string{"-czf", outName, "-C", parent, "--null", "--no-recursion", "-T", "-"} - cmd := exec.CommandContext(ctx, "tar", args...) // Run tar in the output directory so the bare -f basename lands there. cmd.Dir = outDirAbs var stdin bytes.Buffer - for _, file := range files { - stdin.WriteString(filepath.ToSlash(filepath.Join(dirName, file))) + for _, f := range files { + stdin.WriteString(filepath.ToSlash(filepath.Join(dirName, f.rel))) stdin.WriteByte(0) } cmd.Stdin = &stdin @@ -84,7 +79,15 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, inc return nil } -func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, isGitRepo bool) ([]string, error) { +// snapshotFile is a file selected for the snapshot: its repo-relative path (native +// separators) 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-") @@ -108,7 +111,7 @@ func snapshotFiles(ctx context.Context, repoPath string, includePaths []string, return nil, fmt.Errorf("failed to evaluate git ignore rules: %w", err) } - var files []string + var files []snapshotFile for raw := range bytes.SplitSeq(output, []byte{0}) { if len(raw) == 0 { continue @@ -118,14 +121,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 41503fd683c..ca7d7da664a 100644 --- a/experimental/air/cmd/snapshot_package_test.go +++ b/experimental/air/cmd/snapshot_package_test.go @@ -87,7 +87,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) @@ -107,7 +109,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) @@ -122,7 +126,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) @@ -142,7 +148,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) @@ -162,7 +170,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) diff --git a/experimental/air/cmd/snapshot_resolve.go b/experimental/air/cmd/snapshot_resolve.go index 1146b641b92..cacab72694f 100644 --- a/experimental/air/cmd/snapshot_resolve.go +++ b/experimental/air/cmd/snapshot_resolve.go @@ -19,7 +19,8 @@ const ( // deterministic, so the tarball is cacheable by (commit, include_paths). 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 )