diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 8c0be55260..5d4857e7cc 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -163,13 +163,12 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run return 0, "", err } - // Package and upload the code snapshot, if any. The resulting paths ride on the - // ai_runtime_task; a run with no code_source leaves them empty. Snapshot is the - // only code_source type; guard against a nil block so snapshotCodeSource never - // dereferences a missing snapshot. + // Package and upload the code snapshot, if any, via DABs' artifact-upload + // plumbing; the remote code_source_path rides the ai_runtime_task. A run with no + // code_source leaves it empty. Snapshot is the only code_source type. var snap snapshotResult if cfg.CodeSource != nil && cfg.CodeSource.Snapshot != nil { - snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir) + snap, err = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath) if err != nil { return 0, "", err } diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index fd5103599d..e5e4dace41 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -6,6 +6,7 @@ import ( "strings" "testing" + "github.com/databricks/cli/libs/cmdio" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/jobs" @@ -156,8 +157,8 @@ func TestSubmitWorkload(t *testing.T) { assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu1xH100, AcceleratorCount: 1}, d.Compute) } -// TestSubmitWorkloadWithCodeSource exercises the snapshot path end to end: a -// git-pinned code_source is packaged, uploaded, and its paths attached to the task. +// A working-tree code_source is packaged into a tarball, uploaded via DABs' artifact +// plumbing, and its remote code_source_path attached to the submitted task. func TestSubmitWorkloadWithCodeSource(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) @@ -172,32 +173,29 @@ func TestSubmitWorkloadWithCodeSource(t *testing.T) { w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) require.NoError(t, err) - // A git repo committed at HEAD, referenced by commit so packaging is git_archive. - repo := newTestRepo(t) + // A plain working-tree directory: packaging is plain-tar. + repo := filepath.Join(t.TempDir(), "src") writeRepoFile(t, repo, "train.py", "print()") - sha := commitAll(t, repo, "init") cfg := minimalConfig + ` code_source: type: snapshot snapshot: root_path: ` + repo + ` - git: - commit: ` + sha + ` ` cfgPath := writeConfigFile(t, "run.yaml", cfg) loaded, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(t.Context(), w, loaded, cfgPath, "idem") + // 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") require.NoError(t, err) at := got.Tasks[0].AiRuntimeTask - // The tarball path is under the user's repo_snapshots dir. git_state_path / - // git_diff_path are not asserted: the typed jobs.AiRuntimeTask has no such fields - // (see the TEMP note in buildSubmitPayload), so they aren't sent. The git_state - // sidecar file is still uploaded next to the tarball — covered by TestRunSnapshot. - assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/"+filepath.Base(repo)+"/") + // The tarball is uploaded to the artifact .internal dir and code_source_path + // rewritten to it. + assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/.internal/") assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) } diff --git a/experimental/air/cmd/snapshot.go b/experimental/air/cmd/snapshot.go index aba67f6109..3f9c388c00 100644 --- a/experimental/air/cmd/snapshot.go +++ b/experimental/air/cmd/snapshot.go @@ -1,58 +1,25 @@ package aircmd import ( - "bytes" "context" - "errors" "fmt" - "io/fs" "os" - "path" "path/filepath" "strings" - "time" "github.com/databricks/cli/libs/env" - "github.com/databricks/cli/libs/filer" - "github.com/databricks/cli/libs/log" - "github.com/databricks/databricks-sdk-go" ) -// Snapshot orchestrator: resolve → package+upload → sidecars, uploading via -// libs/filer. The Python CLI did this inline; here it's split into steps. - -// snapshotResult holds the paths wired into the submit payload: the uploaded -// tarball and the optional provenance sidecars (empty when not produced). +// snapshotResult holds the code_source_path wired into the submit payload: the +// uploaded code archive's remote path. type snapshotResult struct { CodeSourcePath string - GitStatePath string - GitDiffPath string -} - -// repoSnapshotsSubdir is the per-user workspace location for cached tarballs, under -// the user's home. Volume uploads use remote_volume directly instead. -const repoSnapshotsSubdir = ".air/repo_snapshots" - -// snapshotCodeSource packages and uploads the code_source snapshot, returning the -// paths to attach to the ai_runtime_task. userDir is the user's workspace home; -// funcDir is the run's launch directory (where sidecars land). -func snapshotCodeSource(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath, userDir, funcDir string) (snapshotResult, error) { - repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) - if err != nil { - return snapshotResult{}, err - } - - up, err := newSnapshotUploader(w, snap, userDir, funcDir, filepath.Base(repoPath)) - if err != nil { - return snapshotResult{}, err - } - return runSnapshot(ctx, up, repoPath, snap) } -// resolveRootPath resolves a snapshot root_path the way the Python normalize layer -// does: expand environment variables and ~, strip a leading "project_root/" (meaning -// "relative to the YAML file"), and resolve the rest against the config's directory. -// It then confirms the path exists and is a directory. +// resolveRootPath resolves a code_source snapshot root_path: expand environment +// variables and ~, strip a leading "project_root/" (meaning "relative to the YAML +// file"), and resolve the rest against the config's directory. It then confirms the +// path exists and is a directory. func resolveRootPath(ctx context.Context, rawPath, configDir string) (string, error) { expanded := os.ExpandEnv(rawPath) if home, err := env.UserHomeDir(ctx); err == nil { @@ -73,8 +40,6 @@ func resolveRootPath(ctx context.Context, rawPath, configDir string) (string, er resolved = filepath.Join(configDir, expanded) } - // Resolve to an absolute path so the directory name (used for the tarball name - // and archive prefix) is a real basename, not "." or a trailing relative segment. abs, err := filepath.Abs(resolved) if err != nil { return "", fmt.Errorf("failed to resolve root_path %s: %w", resolved, err) @@ -90,190 +55,3 @@ func resolveRootPath(ctx context.Context, rawPath, configDir string) (string, er } return resolved, nil } - -// snapshotUploader splits the snapshot's two destinations: the tarball goes to a -// cache location (the user's repo_snapshots dir or a Volume), sidecars to the run's -// funcDir. tarBase/sidecarBase are the absolute roots, for reporting final paths. -type snapshotUploader struct { - tarStore filer.Filer - sidecarStore filer.Filer - tarBase string - sidecarBase string -} - -// runSnapshot resolves the packaging plan, uploads the tarball, then uploads the -// provenance sidecars. repoPath is the resolved root_path. -func runSnapshot(ctx context.Context, up snapshotUploader, repoPath string, snap *snapshotSourceConfig) (snapshotResult, error) { - git := newGitRepo(repoPath) - plan, err := resolveSnapshotPlan(ctx, git, snap.Git, snap.IncludePaths) - if err != nil { - return snapshotResult{}, err - } - - dirName := filepath.Base(repoPath) - - tarName, err := uploadTarball(ctx, up, git, plan, repoPath, dirName) - if err != nil { - return snapshotResult{}, err - } - - result := snapshotResult{CodeSourcePath: path.Join(up.tarBase, tarName)} - - // Provenance sidecars are best-effort: a git/upload hiccup here must not fail an - // otherwise-valid submission. Non-git roots have no provenance to record. - if plan.isGitRepo { - result.GitStatePath, result.GitDiffPath = uploadSidecars(ctx, up, git, plan) - } - return result, nil -} - -// uploadTarball packages the snapshot and uploads it, returning the tarball's name -// within the tar store. For git_archive it checks the cache first and skips -// packaging+upload on a hit. It writes the tarball to a temp file that is always -// cleaned up. -func uploadTarball(ctx context.Context, up snapshotUploader, git gitRepo, plan snapshotPlan, repoPath, dirName string) (string, error) { - // git_archive is cacheable by (commit, include_paths); a hit means the identical - // tarball is already uploaded, so packaging and upload are skipped entirely. - if plan.mode == modeGitArchive { - cacheKey := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths) - tarName := fmt.Sprintf("%s_%s.tar.gz", dirName, cacheKey[:16]) - if exists, err := fileExists(ctx, up.tarStore, tarName); err != nil { - return "", err - } else if exists { - log.Debugf(ctx, "snapshot cache hit for %s at %s", shortSHA(plan.commitSHA), path.Join(up.tarBase, tarName)) - return tarName, nil - } - if err := packageAndUpload(ctx, up, tarName, func(out string) error { - return createGitArchiveSnapshot(ctx, git, plan.commitSHA, out, dirName, plan.includePaths) - }); err != nil { - return "", err - } - return tarName, nil - } - - // plain_tar is not cacheable (working-tree content isn't pinned to a SHA), so it - // is timestamp-named to avoid clobbering a concurrent submission. - tarName := fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) - if err := packageAndUpload(ctx, up, tarName, func(out string) error { - return createPlainTarball(ctx, repoPath, out, plan.includePaths) - }); err != nil { - return "", err - } - return tarName, nil -} - -// packageAndUpload writes the tarball via pkg into a temp file, then uploads it to -// tarName in the tar store. The temp file is always removed. -func packageAndUpload(ctx context.Context, up snapshotUploader, tarName string, pkg func(outputPath string) error) error { - tmp, err := os.CreateTemp("", "air-snapshot-*.tar.gz") - if err != nil { - return fmt.Errorf("failed to create temp tarball: %w", err) - } - tmpPath := tmp.Name() - tmp.Close() - defer os.Remove(tmpPath) - - if err := pkg(tmpPath); err != nil { - return err - } - - f, err := os.Open(tmpPath) - if err != nil { - return fmt.Errorf("failed to open tarball: %w", err) - } - defer f.Close() - - if err := up.tarStore.Write(ctx, tarName, f, filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { - return fmt.Errorf("failed to upload snapshot to %s: %w", path.Join(up.tarBase, tarName), err) - } - return nil -} - -// uploadSidecars builds and uploads the git_state.json and optional git_diff.patch -// provenance sidecars into the run's funcDir. It is best-effort: any failure logs a -// warning and returns whatever paths did upload (possibly none), never an error. -func uploadSidecars(ctx context.Context, up snapshotUploader, git gitRepo, plan snapshotPlan) (statePath, diffPath string) { - mode := packagingModePlainTar - pinnedTip := "" - if plan.mode == modeGitArchive { - mode = packagingModeGitArchive - pinnedTip = plan.commitSHA - } - - sidecar, err := buildGitStateSidecar(ctx, git, mode, pinnedTip, time.Now()) - if err != nil { - log.Warnf(ctx, "skipping git provenance sidecar: %v", err) - return "", "" - } - - // Capture the dirty diff first so its status/path land in git_state.json. - if sidecar.Dirty { - status, diff := captureDirtyDiff(ctx, git, dirtyDiffSizeCapBytes, dirtyDiffTimeout) - sidecar.DiffStatus = status - if status == diffStatusCaptured { - if err := up.sidecarStore.Write(ctx, gitDiffName, bytes.NewReader(diff), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { - log.Warnf(ctx, "failed to upload git diff sidecar: %v", err) - sidecar.DiffStatus = diffStatusClean - } else { - diffPath = path.Join(up.sidecarBase, gitDiffName) - sidecar.DiffPath = &diffPath - } - } - } - - data, err := sidecar.marshal() - if err != nil { - log.Warnf(ctx, "failed to encode git state sidecar: %v", err) - return "", diffPath - } - if err := up.sidecarStore.Write(ctx, gitStateName, bytes.NewReader(data), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { - log.Warnf(ctx, "failed to upload git state sidecar: %v", err) - return "", diffPath - } - return path.Join(up.sidecarBase, gitStateName), diffPath -} - -// gitStateName and gitDiffName are the sidecar basenames read by the backend. -const ( - gitStateName = "git_state.json" - gitDiffName = "git_diff.patch" -) - -// fileExists reports whether name exists in the store, treating fs.ErrNotExist as -// "no". Any other error propagates. -func fileExists(ctx context.Context, store filer.Filer, name string) (bool, error) { - _, err := store.Stat(ctx, name) - if err == nil { - return true, nil - } - if errors.Is(err, fs.ErrNotExist) { - return false, nil - } - return false, fmt.Errorf("failed to check snapshot cache: %w", err) -} - -// newSnapshotUploader builds the uploader for a submission. The tarball store is a -// Volume (when remote_volume is set) or the user's repo_snapshots workspace dir; -// sidecars always go to the run's funcDir in the workspace. -func newSnapshotUploader(w *databricks.WorkspaceClient, snap *snapshotSourceConfig, userDir, funcDir, dirName string) (snapshotUploader, error) { - sidecarStore, err := filer.NewWorkspaceFilesClient(w, funcDir) - if err != nil { - return snapshotUploader{}, err - } - - if snap.RemoteVolume != nil { - tarBase := strings.TrimRight(*snap.RemoteVolume, "/") - tarStore, err := filer.NewFilesClient(w, tarBase) - if err != nil { - return snapshotUploader{}, err - } - return snapshotUploader{tarStore: tarStore, sidecarStore: sidecarStore, tarBase: tarBase, sidecarBase: funcDir}, nil - } - - tarBase := path.Join(userDir, repoSnapshotsSubdir, dirName) - tarStore, err := filer.NewWorkspaceFilesClient(w, tarBase) - if err != nil { - return snapshotUploader{}, err - } - return snapshotUploader{tarStore: tarStore, sidecarStore: sidecarStore, tarBase: tarBase, sidecarBase: funcDir}, nil -} diff --git a/experimental/air/cmd/snapshot_cachekey.go b/experimental/air/cmd/snapshot_cachekey.go deleted file mode 100644 index 44c58ee903..0000000000 --- a/experimental/air/cmd/snapshot_cachekey.go +++ /dev/null @@ -1,34 +0,0 @@ -package aircmd - -// This file packages a local code directory into a tarball, uploads it to the -// workspace (or a Volume), and records git provenance sidecars for cache -// invalidation — the Go port of the Python CLI's code_source snapshot path. - -import ( - "crypto/sha256" - "encoding/hex" - "slices" - "strings" -) - -// snapshotPackagingVersion is bumped when packaging logic changes in a way that invalidates existing caches -const snapshotPackagingVersion = "v1" - -// 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. -func computeSnapshotCacheKey(commitSHA string, includePaths []string) string { - var normalizedPaths string - if len(includePaths) > 0 { - trimmed := make([]string, len(includePaths)) - for i, p := range includePaths { - trimmed[i] = strings.TrimSpace(p) - } - slices.Sort(trimmed) - normalizedPaths = strings.Join(trimmed, "\n") - } - - keyMaterial := commitSHA + "\n" + normalizedPaths + "\n" + snapshotPackagingVersion - sum := sha256.Sum256([]byte(keyMaterial)) - return hex.EncodeToString(sum[:]) -} diff --git a/experimental/air/cmd/snapshot_cachekey_test.go b/experimental/air/cmd/snapshot_cachekey_test.go deleted file mode 100644 index 5743217c00..0000000000 --- a/experimental/air/cmd/snapshot_cachekey_test.go +++ /dev/null @@ -1,62 +0,0 @@ -package aircmd - -import ( - "encoding/json" - "os" - "path/filepath" - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -type goldenCase struct { - Name string `json:"name"` - CommitSHA string `json:"commit_sha"` - IncludePaths []string `json:"include_paths"` - CacheKey string `json:"cache_key"` -} - -// TestComputeSnapshotCacheKeyGolden asserts byte-for-byte parity with golden -// fixtures across the local-only matrix (commit + include_paths permutations). -func TestComputeSnapshotCacheKeyGolden(t *testing.T) { - data, err := os.ReadFile(filepath.Join("testdata", "cache_keys.json")) - require.NoError(t, err) - - var cases []goldenCase - require.NoError(t, json.Unmarshal(data, &cases)) - require.NotEmpty(t, cases) - - for _, tc := range cases { - t.Run(tc.Name, func(t *testing.T) { - assert.Equal(t, tc.CacheKey, computeSnapshotCacheKey(tc.CommitSHA, tc.IncludePaths)) - }) - } -} - -// TestComputeSnapshotCacheKeyProperties pins the normalization behavior the golden cases -// encode, so a regression is legible without decoding hashes. -func TestComputeSnapshotCacheKeyProperties(t *testing.T) { - sha := "a3492b801c0ffee00000000000000000000dead" - - // Order-independent: sorting means unsorted input yields the sorted key. - assert.Equal(t, - computeSnapshotCacheKey(sha, []string{"a", "b", "c"}), - computeSnapshotCacheKey(sha, []string{"c", "a", "b"}), - ) - - // nil and empty include_paths are equivalent (both contribute an empty line). - assert.Equal(t, computeSnapshotCacheKey(sha, nil), computeSnapshotCacheKey(sha, []string{})) - - // Paths are trimmed before hashing. - assert.Equal(t, - computeSnapshotCacheKey(sha, []string{"research", "data"}), - computeSnapshotCacheKey(sha, []string{" research ", " data "}), - ) - - // Duplicates are NOT collapsed — they are sorted and kept, matching Python. - assert.NotEqual(t, computeSnapshotCacheKey(sha, []string{"x", "y"}), computeSnapshotCacheKey(sha, []string{"x", "x", "y"})) - - // The version constant participates: a different version is a different key. - assert.NotEqual(t, snapshotPackagingVersion, "") -} diff --git a/experimental/air/cmd/snapshot_dabs.go b/experimental/air/cmd/snapshot_dabs.go new file mode 100644 index 0000000000..afe63fd6ba --- /dev/null +++ b/experimental/air/cmd/snapshot_dabs.go @@ -0,0 +1,142 @@ +package aircmd + +import ( + "context" + "errors" + "fmt" + "os" + "path" + "path/filepath" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/bundle/libraries" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/vfs" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" +) + +// snapshotViaDABsUpload packages the code_source into a tarball and uploads it using +// DABs' artifact-upload plumbing (the same path a bundle uses for a file-valued +// code_source_path), returning the remote path to attach to the ai_runtime_task. +// +// The packaging + upload logic is CLI-owned (this file, OWNERS = us); it only reuses +// DABs' libraries.ReplaceWithRemotePath + libraries.Upload as the uploader so we do +// not reimplement workspace/volume upload. A minimal in-memory bundle carries the +// local tarball path as code_source_path; ReplaceWithRemotePath rewrites it to the +// artifact .internal path and Upload pushes the bytes. +func snapshotViaDABsUpload(ctx context.Context, w *databricks.WorkspaceClient, snap *snapshotSourceConfig, configPath string) (snapshotResult, error) { + if snap.RemoteVolume != nil { + return snapshotResult{}, errors.New("code_source.snapshot.remote_volume is not yet supported") + } + if snap.Git != nil { + return snapshotResult{}, errors.New("code_source.snapshot.git is not yet supported on this path") + } + + repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) + if err != nil { + return snapshotResult{}, err + } + + tarball, cleanup, err := buildSnapshotTarball(ctx, repoPath, snap.IncludePaths) + if err != nil { + return snapshotResult{}, err + } + defer cleanup() + + return uploadSnapshotViaDABs(ctx, w, tarball) +} + +// buildSnapshotTarball writes a plain-tar snapshot of repoPath to a temp file named +// .tar.gz (the basename becomes the uploaded filename), returning the path +// and a cleanup func. +func buildSnapshotTarball(ctx context.Context, repoPath string, includePaths []string) (string, func(), error) { + noop := func() {} + tmp, err := os.MkdirTemp("", "air-snapshot-*") + if err != nil { + return "", noop, err + } + cleanup := func() { _ = os.RemoveAll(tmp) } + + tarball := filepath.Join(tmp, filepath.Base(repoPath)+".tar.gz") + if err := createPlainTarball(ctx, repoPath, tarball, includePaths); err != nil { + cleanup() + return "", noop, err + } + return tarball, cleanup, nil +} + +// uploadSnapshotViaDABs uploads a local tarball through DABs' artifact-upload +// machinery and returns its remote code_source_path. It builds a minimal bundle whose +// only artifact is the tarball (as a file-valued code_source_path), rewrites the field +// to the remote .internal path, and uploads the bytes. +func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, tarball string) (snapshotResult, error) { + base, err := userWorkspaceDir(ctx, w) + if err != nil { + return snapshotResult{}, err + } + // Upload under the user's repo_snapshots dir (not the default bundle artifact_path, + // which a deploy would clean up); ReplaceWithRemotePath appends /.internal. + artifactPath := path.Join(base, ".air", "repo_snapshots") + + b := &bundle.Bundle{ + BundleRootPath: filepath.Dir(tarball), + BundleRoot: vfs.MustNew(filepath.Dir(tarball)), + SyncRootPath: filepath.Dir(tarball), + SyncRoot: vfs.MustNew(filepath.Dir(tarball)), + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Workspace: config.Workspace{ArtifactPath: artifactPath}, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "air": { + JobSettings: jobs.JobSettings{ + Tasks: []jobs.Task{{ + TaskKey: "air", + // Relative to SyncRootPath (the tarball's dir); collectLocalLibraries + // joins it back and uploads the file. + AiRuntimeTask: &jobs.AiRuntimeTask{CodeSourcePath: filepath.Base(tarball)}, + }}, + }, + }, + }, + }, + }, + } + b.SetWorkpaceClient(w) + if err := b.Config.Mutate(func(v dyn.Value) (dyn.Value, error) { return v, nil }); err != nil { + return snapshotResult{}, err + } + + // Rewrite code_source_path to its remote .internal path (returns the upload set), + // then upload the bytes. + libs, diags := libraries.ReplaceWithRemotePath(ctx, b) + if diags.HasError() { + return snapshotResult{}, diags.Error() + } + if diags := bundle.Apply(ctx, b, libraries.Upload(libs)); diags.HasError() { + return snapshotResult{}, diags.Error() + } + + remote, err := readCodeSourcePath(b) + if err != nil { + return snapshotResult{}, err + } + return snapshotResult{CodeSourcePath: remote}, nil +} + +// readCodeSourcePath returns the (rewritten) code_source_path from the bundle config. +func readCodeSourcePath(b *bundle.Bundle) (string, error) { + v, err := dyn.GetByPath(b.Config.Value(), + dyn.MustPathFromString("resources.jobs.air.tasks[0].ai_runtime_task.code_source_path")) + if err != nil { + return "", fmt.Errorf("code snapshot was not packaged: %w", err) + } + s, ok := v.AsString() + if !ok { + return "", fmt.Errorf("unexpected code_source_path value %v", v.AsAny()) + } + return s, nil +} diff --git a/experimental/air/cmd/snapshot_git.go b/experimental/air/cmd/snapshot_git.go deleted file mode 100644 index 616b3049f7..0000000000 --- a/experimental/air/cmd/snapshot_git.go +++ /dev/null @@ -1,309 +0,0 @@ -package aircmd - -import ( - "bytes" - "context" - "encoding/json" - "errors" - "fmt" - "os/exec" - "strings" - "time" -) - -// Local, no-network git introspection and the git-state provenance sidecar, ported -// from the Python CLI's cli/utils/git_state.py. The remote-fetch helpers -// (fetch_branch_sha, remote detection, partial clone) are deliberately not ported: -// the snapshot archives the local copy only, so a ref must resolve to a local commit. - -// gitRepo runs git subcommands scoped to one repository via `git -C`. Arguments are -// passed as a slice, never a shell string, so branch/commit values can't inject. -type gitRepo struct { - path string -} - -func newGitRepo(path string) gitRepo { - return gitRepo{path: path} -} - -// run executes `git ` and returns stdout; a non-zero exit wraps stderr. -func (g gitRepo) run(ctx context.Context, args ...string) (string, error) { - out, err := g.runBytes(ctx, args...) - return string(out), err -} - -// runBytes is run returning raw stdout bytes, for the dirty-diff capture which needs -// exact bytes and a size measurement. -func (g gitRepo) runBytes(ctx context.Context, args ...string) ([]byte, error) { - full := append([]string{"-C", g.path}, args...) - cmd := exec.CommandContext(ctx, "git", full...) - - var stdout, stderr bytes.Buffer - cmd.Stdout = &stdout - cmd.Stderr = &stderr - if err := cmd.Run(); err != nil { - msg := strings.TrimSpace(stderr.String()) - if msg == "" { - return nil, fmt.Errorf("git %s: %w", strings.Join(args, " "), err) - } - return nil, fmt.Errorf("git %s: %w: %s", strings.Join(args, " "), err, msg) - } - return stdout.Bytes(), nil -} - -// isRepository reports whether the path is inside a git work tree. Using -// `rev-parse --is-inside-work-tree` (not a .git lookup) means a subdirectory of a -// repo counts — the common case when root_path is a subfolder of a monorepo. -func (g gitRepo) isRepository(ctx context.Context) bool { - out, err := g.run(ctx, "rev-parse", "--is-inside-work-tree") - if err != nil { - return false - } - return strings.TrimSpace(out) == "true" -} - -// headSHA returns the current HEAD commit SHA. -func (g gitRepo) headSHA(ctx context.Context) (string, error) { - out, err := g.run(ctx, "rev-parse", "HEAD") - if err != nil { - return "", err - } - return strings.TrimSpace(out), nil -} - -// hasUncommittedChanges reports whether there are staged or unstaged changes under -// the repo subtree. The `-- .` pathspec scopes the check so a subfolder snapshot -// considers only changes that could land in it. -func (g gitRepo) hasUncommittedChanges(ctx context.Context) (bool, error) { - out, err := g.run(ctx, "status", "--porcelain", "--", ".") - if err != nil { - return false, err - } - return strings.TrimSpace(out) != "", nil -} - -// hasUncommittedChangesInPaths reports whether there are uncommitted changes within -// the include paths (empty includePaths yields false). -// -// The pathspecs limit `git status` to those subtrees (it is O(working tree), slow on -// a large monorepo) and already scope the output to what could land in the snapshot. -// Unlike the Python source we don't re-parse the entries to filter by name: git -// reports a rename as `R \x00`, so a name-based re-filter keys off the old -// path and could miss a rename into an include path. The only caller needs the bool. -func (g gitRepo) hasUncommittedChangesInPaths(ctx context.Context, includePaths []string) (bool, error) { - var pathspecs []string - for _, p := range includePaths { - if s := strings.TrimRight(p, "/"); s != "" { - pathspecs = append(pathspecs, s) - } - } - if len(pathspecs) == 0 { - return false, nil - } - - args := append([]string{"status", "--porcelain", "--"}, pathspecs...) - out, err := g.run(ctx, args...) - if err != nil { - return false, err - } - return strings.TrimSpace(out) != "", nil -} - -// resolveLocalBranchSHA resolves a branch to its local-HEAD commit. No remote is -// contacted; the branch must exist locally. -func (g gitRepo) resolveLocalBranchSHA(ctx context.Context, branch string) (string, error) { - out, err := g.run(ctx, "rev-parse", "refs/heads/"+branch) - if err != nil { - return "", fmt.Errorf("failed to resolve local branch %q; ensure the branch exists locally and root_path is correct: %w", branch, err) - } - return strings.TrimSpace(out), nil -} - -// commitExistsLocally reports whether commitSHA is in the local object store, without -// triggering a promisor/lazy fetch. -func (g gitRepo) commitExistsLocally(ctx context.Context, commitSHA string) bool { - _, err := g.run(ctx, "cat-file", "-e", commitSHA) - return err == nil -} - -// currentBranch returns the branch name, or "" for a detached HEAD or on error. -func (g gitRepo) currentBranch(ctx context.Context) string { - out, err := g.run(ctx, "rev-parse", "--abbrev-ref", "HEAD") - if err != nil { - return "" - } - branch := strings.TrimSpace(out) - if branch == "HEAD" { - return "" - } - return branch -} - -// remoteURL returns the URL of the named remote, or "" if it has none. -func (g gitRepo) remoteURL(ctx context.Context, remoteName string) string { - out, err := g.run(ctx, "remote", "get-url", remoteName) - if err != nil { - return "" - } - return strings.TrimSpace(out) -} - -// mergeBaseWithUpstream resolves the merge-base of HEAD and a likely upstream ref, -// trying /HEAD, /main, then /master. It reads only local remote-tracking -// refs (no fetch), returning "" if none resolve. -func (g gitRepo) mergeBaseWithUpstream(ctx context.Context, remoteName string) string { - for _, ref := range []string{remoteName + "/HEAD", remoteName + "/main", remoteName + "/master"} { - out, err := g.run(ctx, "merge-base", "HEAD", ref) - if err != nil { - continue - } - if base := strings.TrimSpace(out); base != "" { - return base - } - } - return "" -} - -// validateIncludePathsExist checks that every include path exists at commitSHA. -// `git ls-tree` (without -d, so both blobs and trees count) reports an entry when the -// path exists; empty output means missing. -func (g gitRepo) validateIncludePathsExist(ctx context.Context, commitSHA string, includePaths []string) error { - var missing []string - for _, p := range includePaths { - out, err := g.run(ctx, "ls-tree", commitSHA, p) - if err != nil { - return err - } - if strings.TrimSpace(out) == "" { - missing = append(missing, p) - } - } - if len(missing) > 0 { - return fmt.Errorf("include_paths do not exist at commit %s: %s", shortSHA(commitSHA), strings.Join(missing, ", ")) - } - return nil -} - -// shortSHA abbreviates a commit SHA to 8 chars for log/error messages, tolerating -// user-supplied abbreviations shorter than that. -func shortSHA(sha string) string { - return sha[:min(len(sha), 8)] -} - -// --- git-state provenance sidecar (git_state.json + git_diff.patch) --- -// -// The backend reads git_state.json next to the tarball to tag the MLflow run with -// base/tip/dirty provenance, and logs git_diff.patch when the tree was dirty. -// Producing the sidecar is best-effort: callers warn and continue, never fail submit. - -// snapshotStateSchemaVersion is the git_state.json schema version. Bump only in -// coordination with the backend reader. -const snapshotStateSchemaVersion = 1 - -// defaultRemoteName is the remote consulted for merge-base and repo URL (local refs -// only — the remote-fetch path is gone). -const defaultRemoteName = "origin" - -// dirtyDiffSizeCapBytes caps the git_diff.patch sidecar; a larger diff records -// size_exceeded and is skipped to keep the upload small. -const dirtyDiffSizeCapBytes = 1024 * 1024 - -// dirtyDiffTimeout bounds `git diff HEAD` so provenance never delays submission. -const dirtyDiffTimeout = 5 * time.Second - -// packaging_mode values: how the uploaded tarball was produced. -const ( - packagingModeGitArchive = "git_archive" - packagingModePlainTar = "plain_tar" -) - -// diff_status values recorded in the sidecar. -const ( - diffStatusClean = "clean" - diffStatusCaptured = "captured" - diffStatusSizeExceeded = "size_exceeded" - diffStatusTimeout = "timeout" -) - -// gitStateSidecar is the git_state.json record. Field names and the null-for-absent -// encoding match the Python source, so nullable fields are *string (absent → null). -type gitStateSidecar struct { - SchemaVersion int `json:"schema_version"` - PackagingMode string `json:"packaging_mode"` - BaseCommit *string `json:"base_commit"` - TipCommit *string `json:"tip_commit"` - Branch *string `json:"branch"` - RepoURL *string `json:"repo_url"` - Dirty bool `json:"dirty"` - DiffStatus string `json:"diff_status"` - DiffPath *string `json:"diff_path"` - GeneratedAtUTC string `json:"generated_at_utc"` -} - -// nilIfEmpty maps "" to nil so an absent value serializes as JSON null. -func nilIfEmpty(s string) *string { - if s == "" { - return nil - } - return &s -} - -// buildGitStateSidecar gathers git provenance. pinnedTip overrides the HEAD-derived -// tip for git_archive (the tarball reflects that commit, not HEAD); pass "" for -// plain_tar. Metadata is best-effort — unavailable fields become null. -func buildGitStateSidecar(ctx context.Context, git gitRepo, packagingMode, pinnedTip string, now time.Time) (gitStateSidecar, error) { - tip := pinnedTip - if tip == "" { - head, err := git.headSHA(ctx) - if err != nil { - return gitStateSidecar{}, err - } - tip = head - } - - dirty, err := git.hasUncommittedChanges(ctx) - if err != nil { - return gitStateSidecar{}, err - } - - return gitStateSidecar{ - SchemaVersion: snapshotStateSchemaVersion, - PackagingMode: packagingMode, - BaseCommit: nilIfEmpty(git.mergeBaseWithUpstream(ctx, defaultRemoteName)), - TipCommit: nilIfEmpty(tip), - Branch: nilIfEmpty(git.currentBranch(ctx)), - RepoURL: nilIfEmpty(git.remoteURL(ctx, defaultRemoteName)), - Dirty: dirty, - DiffStatus: diffStatusClean, - DiffPath: nil, - GeneratedAtUTC: now.UTC().Format("2006-01-02T15:04:05.000000") + "Z", - }, nil -} - -// marshal renders the sidecar as indented JSON (matching Python's json.dump indent=2). -func (s gitStateSidecar) marshal() ([]byte, error) { - return json.MarshalIndent(s, "", " ") -} - -// captureDirtyDiff runs `git diff HEAD` over the repo subtree, returning a diff_status -// and the diff bytes (non-nil only when captured): clean (no changes or diff failed), -// captured (under the cap), size_exceeded, or timeout. -func captureDirtyDiff(ctx context.Context, git gitRepo, sizeCapBytes int, timeout time.Duration) (string, []byte) { - diffCtx, cancel := context.WithTimeout(ctx, timeout) - defer cancel() - - out, err := git.runBytes(diffCtx, "diff", "HEAD", "--", ".") - if err != nil { - if errors.Is(diffCtx.Err(), context.DeadlineExceeded) { - return diffStatusTimeout, nil - } - return diffStatusClean, nil - } - if len(out) == 0 { - return diffStatusClean, nil - } - if len(out) > sizeCapBytes { - return diffStatusSizeExceeded, nil - } - return diffStatusCaptured, out -} diff --git a/experimental/air/cmd/snapshot_git_test.go b/experimental/air/cmd/snapshot_git_test.go deleted file mode 100644 index cf1a821b1d..0000000000 --- a/experimental/air/cmd/snapshot_git_test.go +++ /dev/null @@ -1,296 +0,0 @@ -package aircmd - -import ( - "encoding/json" - "os" - "os/exec" - "path/filepath" - "testing" - "time" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// newTestRepo initializes a git repo in a temp dir with a deterministic identity -// and returns its path. Tests build up real commits/branches/dirty states on top, -// mirroring the Python git_state tests (which drive real repos, not a fake). -func newTestRepo(t *testing.T) string { - t.Helper() - dir := t.TempDir() - runGit(t, dir, "init", "-q", "-b", "main") - // Deterministic identity so commits succeed in a bare CI environment. - runGit(t, dir, "config", "user.email", "test@example.test") - runGit(t, dir, "config", "user.name", "Test") - return dir -} - -// runGit runs a git command in dir and fails the test on error. -func runGit(t *testing.T, dir string, args ...string) { - t.Helper() - cmd := exec.Command("git", append([]string{"-C", dir}, args...)...) - out, err := cmd.CombinedOutput() - require.NoError(t, err, "git %v: %s", args, out) -} - -// writeRepoFile writes a file at a repo-relative path, creating parent dirs. -func writeRepoFile(t *testing.T, repo, rel, content string) { - t.Helper() - full := filepath.Join(repo, rel) - require.NoError(t, os.MkdirAll(filepath.Dir(full), 0o755)) - require.NoError(t, os.WriteFile(full, []byte(content), 0o600)) -} - -// commitAll stages everything and commits, returning the new HEAD SHA. -func commitAll(t *testing.T, repo, msg string) string { - t.Helper() - runGit(t, repo, "add", "-A") - runGit(t, repo, "commit", "-q", "-m", msg) - sha, err := newGitRepo(repo).headSHA(t.Context()) - require.NoError(t, err) - return sha -} - -func TestGitRepo_IsRepository(t *testing.T) { - ctx := t.Context() - - repo := newTestRepo(t) - assert.True(t, newGitRepo(repo).isRepository(ctx)) - - // A subdirectory of a repo is still inside the work tree. - writeRepoFile(t, repo, "sub/x.txt", "hi") - assert.True(t, newGitRepo(filepath.Join(repo, "sub")).isRepository(ctx)) - - // A plain temp dir with no repo is not. - assert.False(t, newGitRepo(t.TempDir()).isRepository(ctx)) -} - -func TestGitRepo_HeadSHA(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - sha := commitAll(t, repo, "init") - - got, err := newGitRepo(repo).headSHA(ctx) - require.NoError(t, err) - assert.Equal(t, sha, got) - assert.Len(t, got, 40) -} - -func TestGitRepo_HasUncommittedChanges(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - - // Clean tree. - dirty, err := newGitRepo(repo).hasUncommittedChanges(ctx) - require.NoError(t, err) - assert.False(t, dirty) - - // Unstaged modification. - writeRepoFile(t, repo, "a.txt", "2") - dirty, err = newGitRepo(repo).hasUncommittedChanges(ctx) - require.NoError(t, err) - assert.True(t, dirty) -} - -func TestGitRepo_HasUncommittedChangesInPaths(t *testing.T) { - ctx := t.Context() - - repo := newTestRepo(t) - writeRepoFile(t, repo, "src/model.py", "1") - writeRepoFile(t, repo, "other/x.py", "1") - commitAll(t, repo, "init") - g := newGitRepo(repo) - - // No paths: no changes, and git is never consulted. - dirty, err := g.hasUncommittedChangesInPaths(ctx, nil) - require.NoError(t, err) - assert.False(t, dirty) - - // A change outside the included paths is ignored. - writeRepoFile(t, repo, "other/x.py", "2") - dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"src", "configs"}) - require.NoError(t, err) - assert.False(t, dirty) - - // A change inside an included path is reported. - writeRepoFile(t, repo, "src/model.py", "2") - dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"src", "configs"}) - require.NoError(t, err) - assert.True(t, dirty) - - // Trailing slashes on include paths are trimmed for the pathspec. - dirty, err = g.hasUncommittedChangesInPaths(ctx, []string{"other/"}) - require.NoError(t, err) - assert.True(t, dirty) -} - -func TestGitRepo_HasUncommittedChangesInPaths_Rename(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "src/old.py", "content") - commitAll(t, repo, "init") - - // A rename within an included path counts as a change, however git classifies - // it (rename vs delete+add); we only assert the boolean. - runGit(t, repo, "mv", "src/old.py", "src/new.py") - dirty, err := newGitRepo(repo).hasUncommittedChangesInPaths(ctx, []string{"src"}) - require.NoError(t, err) - assert.True(t, dirty) -} - -func TestGitRepo_ResolveLocalBranchSHA(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - mainSHA := commitAll(t, repo, "init") - - // A second branch at its own commit. - runGit(t, repo, "checkout", "-q", "-b", "feature") - writeRepoFile(t, repo, "b.txt", "2") - featSHA := commitAll(t, repo, "feature work") - g := newGitRepo(repo) - - got, err := g.resolveLocalBranchSHA(ctx, "main") - require.NoError(t, err) - assert.Equal(t, mainSHA, got) - - got, err = g.resolveLocalBranchSHA(ctx, "feature") - require.NoError(t, err) - assert.Equal(t, featSHA, got) - - // A branch that does not exist locally errors (no remote is contacted). - _, err = g.resolveLocalBranchSHA(ctx, "nope") - require.Error(t, err) - assert.Contains(t, err.Error(), "resolve local branch") -} - -func TestGitRepo_CommitExistsLocally(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - sha := commitAll(t, repo, "init") - g := newGitRepo(repo) - - assert.True(t, g.commitExistsLocally(ctx, sha)) - assert.False(t, g.commitExistsLocally(ctx, "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef")) -} - -func TestGitRepo_ValidateIncludePathsExist(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "src/model.py", "1") - writeRepoFile(t, repo, "configs/train.yaml", "x") - writeRepoFile(t, repo, "train.py", "print()") - sha := commitAll(t, repo, "init") - g := newGitRepo(repo) - - // Both directory and file include_paths are accepted (ls-tree without -d). - require.NoError(t, g.validateIncludePathsExist(ctx, sha, []string{"src", "configs", "train.py"})) - - err := g.validateIncludePathsExist(ctx, sha, []string{"src", "missing"}) - require.Error(t, err) - assert.Contains(t, err.Error(), "missing") - assert.Contains(t, err.Error(), sha[:8]) -} - -func TestBuildGitStateSidecar_PlainTarClean(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - head := commitAll(t, repo, "init") - - sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) - require.NoError(t, err) - assert.Equal(t, snapshotStateSchemaVersion, sc.SchemaVersion) - assert.Equal(t, packagingModePlainTar, sc.PackagingMode) - require.NotNil(t, sc.TipCommit) - assert.Equal(t, head, *sc.TipCommit) - assert.False(t, sc.Dirty) - assert.Equal(t, diffStatusClean, sc.DiffStatus) - assert.Nil(t, sc.DiffPath) - // No remote in a bare test repo → base_commit and repo_url are null. - assert.Nil(t, sc.BaseCommit) - assert.Nil(t, sc.RepoURL) - require.NotNil(t, sc.Branch) - assert.Equal(t, "main", *sc.Branch) - assert.Equal(t, "2026-07-10T12:00:00.000000Z", sc.GeneratedAtUTC) -} - -func TestBuildGitStateSidecar_GitArchivePinsTip(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - first := commitAll(t, repo, "init") - // Advance HEAD; the pinned tip must win over HEAD. - writeRepoFile(t, repo, "b.txt", "2") - commitAll(t, repo, "second") - - sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModeGitArchive, first, fixedNow) - require.NoError(t, err) - require.NotNil(t, sc.TipCommit) - assert.Equal(t, first, *sc.TipCommit) - assert.Equal(t, packagingModeGitArchive, sc.PackagingMode) -} - -func TestBuildGitStateSidecar_Dirty(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - writeRepoFile(t, repo, "a.txt", "2") // uncommitted - - sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) - require.NoError(t, err) - assert.True(t, sc.Dirty) -} - -func TestGitStateSidecar_MarshalNullsAbsentFields(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - - sc, err := buildGitStateSidecar(ctx, newGitRepo(repo), packagingModePlainTar, "", fixedNow) - require.NoError(t, err) - data, err := sc.marshal() - require.NoError(t, err) - - // Absent fields serialize as JSON null (not "" or omitted), matching Python. - var raw map[string]any - require.NoError(t, json.Unmarshal(data, &raw)) - require.Contains(t, raw, "base_commit") - assert.Nil(t, raw["base_commit"]) - require.Contains(t, raw, "repo_url") - assert.Nil(t, raw["repo_url"]) - require.Contains(t, raw, "diff_path") - assert.Nil(t, raw["diff_path"]) - assert.EqualValues(t, 1, raw["schema_version"]) -} - -func TestCaptureDirtyDiff(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "one\n") - commitAll(t, repo, "init") - - // Clean tree → no diff. - status, diff := captureDirtyDiff(ctx, newGitRepo(repo), dirtyDiffSizeCapBytes, dirtyDiffTimeout) - assert.Equal(t, diffStatusClean, status) - assert.Nil(t, diff) - - // Dirty tree → captured, and the diff mentions the changed file. - writeRepoFile(t, repo, "a.txt", "two\n") - status, diff = captureDirtyDiff(ctx, newGitRepo(repo), dirtyDiffSizeCapBytes, dirtyDiffTimeout) - assert.Equal(t, diffStatusCaptured, status) - assert.Contains(t, string(diff), "a.txt") - - // A tiny size cap forces size_exceeded and drops the bytes. - status, diff = captureDirtyDiff(ctx, newGitRepo(repo), 1, dirtyDiffTimeout) - assert.Equal(t, diffStatusSizeExceeded, status) - assert.Nil(t, diff) -} - -var fixedNow = time.Date(2026, 7, 10, 12, 0, 0, 0, time.UTC) diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index 672366086c..38fdb76c0b 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -10,30 +10,10 @@ import ( "strings" ) -// Tar builders ported from cli/utils/snapshot.py. Both shell out (git archive / tar) -// for parity and to reuse git's/tar's symlink, gitignore, and AppleDouble handling. -// The tarball's top-level dir name is load-bearing — the remote entry_script extracts -// to /databricks/code_source/ — so the --prefix / `-C parent dir` forms preserve it. - -// createGitArchiveSnapshot writes a gzipped tar of commitSHA to outputTarball via -// `git archive`, with every entry prefixed by directoryName/. When includePaths is -// set, only those paths are archived. -func createGitArchiveSnapshot(ctx context.Context, git gitRepo, commitSHA, outputTarball, directoryName string, includePaths []string) error { - // Single git invocation writes the gzipped tar with the desired prefix; no - // extract/repack. Provenance lives in the git_state.json sidecar, not here. - args := []string{ - "archive", - "--format=tar.gz", - "--prefix=" + directoryName + "/", - "-o", outputTarball, - commitSHA, - } - args = append(args, includePaths...) - if _, err := git.run(ctx, args...); err != nil { - return fmt.Errorf("failed to create git archive: %w", err) - } - return nil -} +// Tar builder ported from cli/utils/snapshot.py. It shells out to `tar` for parity +// and to reuse tar's symlink, gitignore, and AppleDouble handling. The tarball's +// top-level dir name is load-bearing — the remote entry_script extracts to +// /databricks/code_source/ — so the `-C parent dir` form preserves it. // createPlainTarball writes a gzipped tar of repoPath's working tree to // outputTarball via `tar`. The archive preserves repoPath's directory name as the diff --git a/experimental/air/cmd/snapshot_package_test.go b/experimental/air/cmd/snapshot_package_test.go index d895d59b98..a3cdaf72b0 100644 --- a/experimental/air/cmd/snapshot_package_test.go +++ b/experimental/air/cmd/snapshot_package_test.go @@ -6,7 +6,6 @@ import ( "os" "path/filepath" "slices" - "strings" "testing" "github.com/stretchr/testify/assert" @@ -37,56 +36,22 @@ func tarballEntries(t *testing.T, path string) []string { return names } -func TestCreateGitArchiveSnapshot(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - writeRepoFile(t, repo, "src/model.py", "print()") - sha := commitAll(t, repo, "init") - - out := filepath.Join(t.TempDir(), "snap.tar.gz") - dirName := filepath.Base(repo) - require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, nil)) - - entries := tarballEntries(t, out) - // Every real entry is prefixed with the directory name; the tracked files are - // present. git archive also emits a `pax_global_header` pseudo-entry carrying - // the commit SHA — it has no prefix and tar ignores it on extraction. - assert.Contains(t, entries, dirName+"/a.txt") - assert.Contains(t, entries, dirName+"/src/model.py") - for _, e := range entries { - if e == "pax_global_header" { - continue - } - assert.True(t, strings.HasPrefix(e, dirName+"/"), "entry %q lacks prefix", e) - } -} - -func TestCreateGitArchiveSnapshot_IncludePaths(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - writeRepoFile(t, repo, "src/model.py", "print()") - sha := commitAll(t, repo, "init") - - out := filepath.Join(t.TempDir(), "snap.tar.gz") - dirName := filepath.Base(repo) - require.NoError(t, createGitArchiveSnapshot(ctx, newGitRepo(repo), sha, out, dirName, []string{"src"})) - - entries := tarballEntries(t, out) - assert.Contains(t, entries, dirName+"/src/model.py") - // a.txt is outside the include path, so it must not appear. - assert.NotContains(t, entries, dirName+"/a.txt") +// writeRepoFile writes a file at a dir-relative path, creating parent dirs. +func writeRepoFile(t *testing.T, dir, rel, content string) { + t.Helper() + full := filepath.Join(dir, rel) + require.NoError(t, os.MkdirAll(filepath.Dir(full), 0o755)) + require.NoError(t, os.WriteFile(full, []byte(content), 0o600)) } func TestCreatePlainTarball(t *testing.T) { ctx := t.Context() - repo := newTestRepo(t) + repo := t.TempDir() writeRepoFile(t, repo, "a.txt", "1") writeRepoFile(t, repo, "src/model.py", "print()") - commitAll(t, repo, "init") - // Uncommitted file must be included in a plain tar. writeRepoFile(t, repo, "dirty.txt", "wip") + // A .git dir must never be shipped. + writeRepoFile(t, repo, ".git/config", "x") out := filepath.Join(t.TempDir(), "snap.tar.gz") require.NoError(t, createPlainTarball(ctx, repo, out, nil)) @@ -103,7 +68,7 @@ func TestCreatePlainTarball(t *testing.T) { func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { ctx := t.Context() - repo := newTestRepo(t) + repo := t.TempDir() writeRepoFile(t, repo, "keep.txt", "1") writeRepoFile(t, repo, "junk.log", "noise") writeRepoFile(t, repo, ".gitignore", "*.log\n") @@ -119,7 +84,7 @@ func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { func TestCreatePlainTarball_IncludePaths(t *testing.T) { ctx := t.Context() - repo := newTestRepo(t) + repo := t.TempDir() writeRepoFile(t, repo, "a.txt", "1") writeRepoFile(t, repo, "src/model.py", "print()") diff --git a/experimental/air/cmd/snapshot_resolve.go b/experimental/air/cmd/snapshot_resolve.go deleted file mode 100644 index 1146b641b9..0000000000 --- a/experimental/air/cmd/snapshot_resolve.go +++ /dev/null @@ -1,120 +0,0 @@ -package aircmd - -import ( - "context" - "errors" - "fmt" -) - -// This file ports the mode/ref resolution from the Python CLI's cli_entrypoint -// snapshot block (the if/elif at lines ~1541–1722), local-only. The remote-fetch -// branches are dropped: a git ref must resolve to a commit already present -// locally (git.remote is rejected at validation — see gitRef.validate). - -// snapshotMode is how the snapshot tarball is produced. -type snapshotMode int - -const ( - // modeGitArchive packages a pinned commit via `git archive`. The commit is - // 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. - modePlainTar -) - -// snapshotPlan is the outcome of resolving how to package a snapshot: the mode, -// the commit SHA to archive (git_archive only; empty for plain_tar), and whether -// the working tree under the snapshot root has uncommitted changes. -type snapshotPlan struct { - mode snapshotMode - commitSHA string - hasUncommit bool - isGitRepo bool - includePaths []string -} - -// resolveSnapshotPlan decides how to package the snapshot (local-only): -// - git.commit → pin the SHA (must exist locally) → git_archive. -// - git.branch → the branch's local HEAD SHA → git_archive. -// - no ref / non-git dir → the working tree → plain_tar (no caching). -// -// The dirty check runs at most once (git status is O(working tree)) and is threaded -// into the plan. Dirty + git.branch is an error: the committed HEAD wouldn't include -// the uncommitted changes. -func resolveSnapshotPlan(ctx context.Context, git gitRepo, ref *gitRef, includePaths []string) (snapshotPlan, error) { - plan := snapshotPlan{includePaths: includePaths} - plan.isGitRepo = git.isRepository(ctx) - - // Detect uncommitted changes once. When include_paths is set, only changes - // under those paths can land in the snapshot, so scope the check to them — - // both more correct and cheaper than scanning the whole repo. - if plan.isGitRepo { - var err error - if len(includePaths) > 0 { - plan.hasUncommit, err = git.hasUncommittedChangesInPaths(ctx, includePaths) - } else { - plan.hasUncommit, err = git.hasUncommittedChanges(ctx) - } - if err != nil { - return snapshotPlan{}, err - } - } - - // Non-git directory: plain tar, no ref allowed. gitRef.validate already rejects - // git.* on a non-git dir at load time, but guard here too since this function - // is the single decision point. - if !plan.isGitRepo { - if ref != nil { - return snapshotPlan{}, fmt.Errorf("git.* is set but %s is not a git repository", git.path) - } - plan.mode = modePlainTar - return plan, nil - } - - // git repo, no ref: package the working tree as plain tar (uncommitted changes - // included). Provenance is captured separately via the git_state sidecar. - if ref == nil { - plan.mode = modePlainTar - return plan, nil - } - - switch { - case ref.Commit != nil: - // git.commit pins a committed SHA; local uncommitted changes are irrelevant - // and won't be included. The commit must exist locally — no remote fetch. - commit := *ref.Commit - if !git.commitExistsLocally(ctx, commit) { - return snapshotPlan{}, fmt.Errorf("commit %q does not exist locally; fetch it (e.g. `git fetch`) before submitting — the snapshot archives your local copy and does not fetch from a remote", commit) - } - plan.mode = modeGitArchive - plan.commitSHA = commit - - case ref.Branch != nil: - // git.branch deploys the branch's local HEAD. A dirty tree here is an error: - // the committed HEAD wouldn't include the uncommitted changes. - if plan.hasUncommit { - return snapshotPlan{}, fmt.Errorf("uncommitted changes under %s would not be included: git.branch deploys the committed HEAD of %q. Commit your changes, or use git.commit to pin a specific revision", git.path, *ref.Branch) - } - sha, err := git.resolveLocalBranchSHA(ctx, *ref.Branch) - if err != nil { - return snapshotPlan{}, err - } - plan.mode = modeGitArchive - plan.commitSHA = sha - - default: - // gitRef.validate guarantees exactly one of branch/commit is set. - return snapshotPlan{}, errors.New("git: must specify either 'branch' or 'commit'") - } - - // For git_archive with include_paths, verify each path exists at the resolved - // commit so a typo fails fast rather than producing an empty subtree. - if len(includePaths) > 0 { - if err := git.validateIncludePathsExist(ctx, plan.commitSHA, includePaths); err != nil { - return snapshotPlan{}, err - } - } - - return plan, nil -} diff --git a/experimental/air/cmd/snapshot_resolve_test.go b/experimental/air/cmd/snapshot_resolve_test.go deleted file mode 100644 index c8c946f839..0000000000 --- a/experimental/air/cmd/snapshot_resolve_test.go +++ /dev/null @@ -1,114 +0,0 @@ -package aircmd - -import ( - "testing" - - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestResolveSnapshotPlan_Commit(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - sha := commitAll(t, repo, "init") - - plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, nil) - require.NoError(t, err) - assert.Equal(t, modeGitArchive, plan.mode) - assert.Equal(t, sha, plan.commitSHA) - assert.True(t, plan.isGitRepo) - - // A commit pin is valid even with a dirty tree: local changes are irrelevant. - writeRepoFile(t, repo, "a.txt", "2") - plan, err = resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, nil) - require.NoError(t, err) - assert.Equal(t, modeGitArchive, plan.mode) - assert.True(t, plan.hasUncommit) -} - -func TestResolveSnapshotPlan_CommitNotLocal(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - - absent := "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef" - _, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(absent)}, nil) - require.Error(t, err) - assert.Contains(t, err.Error(), "does not exist locally") -} - -func TestResolveSnapshotPlan_BranchLocalHead(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - mainSHA := commitAll(t, repo, "init") - - plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Branch: new("main")}, nil) - require.NoError(t, err) - assert.Equal(t, modeGitArchive, plan.mode) - assert.Equal(t, mainSHA, plan.commitSHA) -} - -func TestResolveSnapshotPlan_BranchDirtyIsError(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - writeRepoFile(t, repo, "a.txt", "2") // uncommitted - - _, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Branch: new("main")}, nil) - require.Error(t, err) - assert.Contains(t, err.Error(), "uncommitted changes") - assert.Contains(t, err.Error(), "git.commit") -} - -func TestResolveSnapshotPlan_NoRefPlainTar(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "a.txt", "1") - commitAll(t, repo, "init") - writeRepoFile(t, repo, "a.txt", "2") // dirty tree is fine for plain tar - - plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), nil, nil) - require.NoError(t, err) - assert.Equal(t, modePlainTar, plan.mode) - assert.Empty(t, plan.commitSHA) - assert.True(t, plan.isGitRepo) - assert.True(t, plan.hasUncommit) -} - -func TestResolveSnapshotPlan_NonGitDir(t *testing.T) { - ctx := t.Context() - dir := t.TempDir() - - plan, err := resolveSnapshotPlan(ctx, newGitRepo(dir), nil, nil) - require.NoError(t, err) - assert.Equal(t, modePlainTar, plan.mode) - assert.False(t, plan.isGitRepo) - - // A git ref on a non-git directory is an error. - _, err = resolveSnapshotPlan(ctx, newGitRepo(dir), &gitRef{Branch: new("main")}, nil) - require.Error(t, err) - assert.Contains(t, err.Error(), "not a git repository") -} - -func TestResolveSnapshotPlan_IncludePaths(t *testing.T) { - ctx := t.Context() - repo := newTestRepo(t) - writeRepoFile(t, repo, "src/model.py", "1") - writeRepoFile(t, repo, "configs/train.yaml", "x") - sha := commitAll(t, repo, "init") - - // All include paths exist at the commit. - plan, err := resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, []string{"src", "configs"}) - require.NoError(t, err) - assert.Equal(t, modeGitArchive, plan.mode) - assert.Equal(t, []string{"src", "configs"}, plan.includePaths) - - // A missing include path fails fast. - _, err = resolveSnapshotPlan(ctx, newGitRepo(repo), &gitRef{Commit: new(sha)}, []string{"src", "missing"}) - require.Error(t, err) - assert.Contains(t, err.Error(), "missing") -} diff --git a/experimental/air/cmd/snapshot_test.go b/experimental/air/cmd/snapshot_test.go deleted file mode 100644 index d94fe005fc..0000000000 --- a/experimental/air/cmd/snapshot_test.go +++ /dev/null @@ -1,155 +0,0 @@ -package aircmd - -import ( - "context" - "io" - "os" - "path" - "path/filepath" - "testing" - - "github.com/databricks/cli/libs/filer" - "github.com/databricks/cli/libs/testserver" - "github.com/databricks/databricks-sdk-go" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestResolveRootPath(t *testing.T) { - ctx := t.Context() - dir := t.TempDir() - require.NoError(t, os.MkdirAll(filepath.Join(dir, "proj"), 0o755)) - - // root_path "." resolves against configDir to an absolute path whose basename is - // the real directory name — not "." (which would name the tarball ._.tar.gz, - // colliding with the AppleDouble exclude pattern the remote strips). - got, err := resolveRootPath(ctx, ".", filepath.Join(dir, "proj")) - require.NoError(t, err) - assert.True(t, filepath.IsAbs(got)) - assert.Equal(t, "proj", filepath.Base(got)) - - // A relative subpath resolves against configDir and keeps its own basename. - require.NoError(t, os.MkdirAll(filepath.Join(dir, "proj", "sub"), 0o755)) - got, err = resolveRootPath(ctx, "sub", filepath.Join(dir, "proj")) - require.NoError(t, err) - assert.Equal(t, "sub", filepath.Base(got)) - - // A non-existent path errors. - _, err = resolveRootPath(ctx, "missing", dir) - require.Error(t, err) -} - -// newSnapshotTestClient returns a workspace client backed by the in-process fake, -// which models workspace get-status / import-file with real state. -func newSnapshotTestClient(t *testing.T) *databricks.WorkspaceClient { - t.Helper() - server := testserver.New(t) - t.Cleanup(server.Close) - testserver.AddDefaultHandlers(server) - w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) - require.NoError(t, err) - return w -} - -// testUploader builds a snapshotUploader whose tar store and sidecar store both live -// under distinct workspace roots on the fake server. -func testUploader(t *testing.T, w *databricks.WorkspaceClient, tarBase, sidecarBase string) snapshotUploader { - t.Helper() - tarStore, err := filer.NewWorkspaceFilesClient(w, tarBase) - require.NoError(t, err) - sidecarStore, err := filer.NewWorkspaceFilesClient(w, sidecarBase) - require.NoError(t, err) - return snapshotUploader{tarStore: tarStore, sidecarStore: sidecarStore, tarBase: tarBase, sidecarBase: sidecarBase} -} - -func TestRunSnapshot_GitArchive(t *testing.T) { - ctx := t.Context() - w := newSnapshotTestClient(t) - repo := newTestRepo(t) - writeRepoFile(t, repo, "train.py", "print()") - sha := commitAll(t, repo, "init") - - up := testUploader(t, w, "/Workspace/Users/me/.air/repo_snapshots/repo", "/Workspace/Users/me/.air/cli_launch/exp/run") - res, err := runSnapshot(ctx, up, repo, &snapshotSourceConfig{RootPath: repo, Git: &gitRef{Commit: &sha}}) - require.NoError(t, err) - - // Tarball is cache-key-named under the tar base, prefixed with the repo dir name - // (the temp dir's basename); a clean git repo yields a git_state sidecar, no diff. - cacheKey := computeSnapshotCacheKey(sha, nil) - wantName := filepath.Base(repo) + "_" + cacheKey[:16] + ".tar.gz" - assert.Equal(t, path.Join(up.tarBase, wantName), res.CodeSourcePath) - assert.Equal(t, path.Join(up.sidecarBase, gitStateName), res.GitStatePath) - assert.Empty(t, res.GitDiffPath) -} - -func TestRunSnapshot_CacheHitSkipsUpload(t *testing.T) { - ctx := t.Context() - w := newSnapshotTestClient(t) - repo := newTestRepo(t) - writeRepoFile(t, repo, "train.py", "print()") - sha := commitAll(t, repo, "init") - - up := testUploader(t, w, "/Workspace/Users/me/.air/repo_snapshots/repo", "/Workspace/Users/me/.air/cli_launch/exp/run") - snap := &snapshotSourceConfig{RootPath: repo, Git: &gitRef{Commit: &sha}} - - // First submission uploads the tarball. - res1, err := runSnapshot(ctx, up, repo, snap) - require.NoError(t, err) - - // Count uploads to the tarball path on a fresh uploader: the second run should - // see the cached tarball via Stat and not re-upload it. - writes := &countingFiler{Filer: up.tarStore} - up2 := up - up2.tarStore = writes - res2, err := runSnapshot(ctx, up2, repo, snap) - require.NoError(t, err) - - assert.Equal(t, res1.CodeSourcePath, res2.CodeSourcePath) - assert.Zero(t, writes.writes, "cache hit must not re-upload the tarball") -} - -func TestRunSnapshot_PlainTarDirty(t *testing.T) { - ctx := t.Context() - w := newSnapshotTestClient(t) - repo := newTestRepo(t) - writeRepoFile(t, repo, "train.py", "print()") - commitAll(t, repo, "init") - writeRepoFile(t, repo, "train.py", "print('wip')") // dirty, no git ref - - up := testUploader(t, w, "/Workspace/Users/me/.air/repo_snapshots/repo", "/Workspace/Users/me/.air/cli_launch/exp/run") - res, err := runSnapshot(ctx, up, repo, &snapshotSourceConfig{RootPath: repo}) - require.NoError(t, err) - - // Plain tar is timestamp-named (not cache-key-named); a dirty tree captures both - // the state and the diff sidecar. - assert.Contains(t, res.CodeSourcePath, path.Join(up.tarBase, filepath.Base(repo)+"_")) - assert.Equal(t, path.Join(up.sidecarBase, gitStateName), res.GitStatePath) - assert.Equal(t, path.Join(up.sidecarBase, gitDiffName), res.GitDiffPath) -} - -func TestRunSnapshot_NonGitDir(t *testing.T) { - ctx := t.Context() - w := newSnapshotTestClient(t) - dir := t.TempDir() - writeRepoFile(t, dir, "train.py", "print()") - - up := testUploader(t, w, "/Workspace/Users/me/.air/repo_snapshots/proj", "/Workspace/Users/me/.air/cli_launch/exp/run") - res, err := runSnapshot(ctx, up, dir, &snapshotSourceConfig{RootPath: dir}) - require.NoError(t, err) - - // Non-git dir: plain tar, and no provenance sidecars. - assert.NotEmpty(t, res.CodeSourcePath) - assert.Empty(t, res.GitStatePath) - assert.Empty(t, res.GitDiffPath) -} - -// countingFiler wraps a Filer to count Write calls, for asserting cache-hit skips. -type countingFiler struct { - filer.Filer - writes int -} - -func (c *countingFiler) Write(ctx context.Context, name string, reader io.Reader, mode ...filer.WriteMode) error { - c.writes++ - return c.Filer.Write(ctx, name, reader, mode...) -}