diff --git a/acceptance/experimental/air/run-submit/databricks.yml b/acceptance/experimental/air/run-submit/databricks.yml new file mode 100644 index 00000000000..4a1c612d600 --- /dev/null +++ b/acceptance/experimental/air/run-submit/databricks.yml @@ -0,0 +1,2 @@ +bundle: + name: air-run-submit diff --git a/acceptance/experimental/air/run-submit/out.test.toml b/acceptance/experimental/air/run-submit/out.test.toml new file mode 100644 index 00000000000..d6187dcb046 --- /dev/null +++ b/acceptance/experimental/air/run-submit/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt new file mode 100644 index 00000000000..8d9de387bb1 --- /dev/null +++ b/acceptance/experimental/air/run-submit/output.txt @@ -0,0 +1,47 @@ + +=== submit with a git code_source +>>> [CLI] experimental air run -f run.yaml +Submitted run 555 +View at: [DATABRICKS_URL]/jobs/runs/555 + +=== the ai_runtime_task carries the snapshot + provenance paths +>>> print_requests.py //api/2.2/jobs/runs/submit +{ + "method": "POST", + "path": "/api/2.2/jobs/runs/submit", + "body": { + "run_name": "submit-smoke", + "tasks": [ + { + "task_key": "submit-smoke", + "run_if": "ALL_SUCCESS", + "ai_runtime_task": { + "experiment": "submit-smoke", + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/command.sh", + "compute": { + "accelerator_type": "GPU_1xH100", + "accelerator_count": 1 + } + } + ], + "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/[SNAPSHOT_TARBALL]", + "git_state_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/git_state.json" + }, + "environment_key": "default", + "max_retries": 3, + "retry_on_timeout": true + } + ], + "environments": [ + { + "environment_key": "default", + "spec": { + "environment_version": "4" + } + } + ], + "idempotency_token": "[UUID]" + } +} diff --git a/acceptance/experimental/air/run-submit/run.yaml.tmpl b/acceptance/experimental/air/run-submit/run.yaml.tmpl new file mode 100644 index 00000000000..3fdbf48eb85 --- /dev/null +++ b/acceptance/experimental/air/run-submit/run.yaml.tmpl @@ -0,0 +1,11 @@ +experiment_name: submit-smoke +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +code_source: + type: snapshot + snapshot: + root_path: . + git: + commit: COMMIT_SHA diff --git a/acceptance/experimental/air/run-submit/script b/acceptance/experimental/air/run-submit/script new file mode 100644 index 00000000000..b43e481d882 --- /dev/null +++ b/acceptance/experimental/air/run-submit/script @@ -0,0 +1,18 @@ +# Pinned commit dates keep the resolved HEAD SHA — and thus the snapshot cache key +# baked into the tarball name — stable across runs. +export GIT_AUTHOR_DATE="2020-01-01T00:00:00Z" +export GIT_COMMITTER_DATE="2020-01-01T00:00:00Z" +git-repo-init + +# Pin the config to the committed HEAD. git.commit tolerates a dirty working tree +# (the pinned revision is what gets archived), which matters here because the +# acceptance harness writes output.txt into the working dir as the script runs. +sed "s/COMMIT_SHA/$(git rev-parse HEAD)/" run.yaml.tmpl > run.yaml + +title "submit with a git code_source" +trace $CLI experimental air run -f run.yaml + +title "the ai_runtime_task carries the snapshot + provenance paths" +trace print_requests.py //api/2.2/jobs/runs/submit + +rm -fr .git diff --git a/acceptance/experimental/air/run-submit/test.toml b/acceptance/experimental/air/run-submit/test.toml new file mode 100644 index 00000000000..fcdf8fd1242 --- /dev/null +++ b/acceptance/experimental/air/run-submit/test.toml @@ -0,0 +1,35 @@ +# A real (non-dry-run) submit that packages a git code_source, uploads the +# tarball + provenance sidecars, and POSTs runs/submit. No bundle deploy, so no +# engine matrix. +RecordRequests = true + +# run.yaml is generated from run.yaml.tmpl at test time (commit SHA templated in); +# it isn't a committed input to diff. +Ignore = ["run.yaml"] + +[EnvMatrix] +DATABRICKS_BUNDLE_ENGINE = [] + +# The SDK probes host reachability with a HEAD request; stub it for determinism. +[[Server]] +Pattern = "HEAD /" +Response.Body = '' + +[[Server]] +Pattern = "POST /api/2.2/jobs/runs/submit" +Response.Body = ''' +{"run_id": 555} +''' + +# The snapshot tarball is named _.tar.gz, where is the +# test's temp-dir basename and the cache key derives from the pinned commit SHA. +# Both are stable given the pinned commit dates in the script, but the temp-dir +# basename varies per run, so collapse the whole tarball filename to a stable token. +[[Repls]] +Old = '[a-zA-Z0-9._-]+_[0-9a-f]{16}\.tar\.gz' +New = '[SNAPSHOT_TARBALL]' + +# The per-run launch directory ends in _<16 hex>; the random suffix varies. +[[Repls]] +Old = 'submit-smoke_[0-9a-f]{16}' +New = 'submit-smoke_[RUN_ID]' diff --git a/acceptance/experimental/air/run/git-remote.yaml b/acceptance/experimental/air/run/git-remote.yaml new file mode 100644 index 00000000000..0161fe14972 --- /dev/null +++ b/acceptance/experimental/air/run/git-remote.yaml @@ -0,0 +1,12 @@ +experiment_name: smoke-test +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +code_source: + type: snapshot + snapshot: + root_path: . + git: + branch: main + remote: origin diff --git a/acceptance/experimental/air/run/output.txt b/acceptance/experimental/air/run/output.txt index 180886290ba..a753eabd198 100644 --- a/acceptance/experimental/air/run/output.txt +++ b/acceptance/experimental/air/run/output.txt @@ -26,6 +26,16 @@ Error: --watch is not yet supported Exit code: 1 +=== code_source config passes validation +>>> [CLI] experimental air run -f with-code-source.yaml --dry-run +Dry run: configuration for "smoke-test" is valid; not submitting. + +=== git.remote is rejected +>>> [CLI] experimental air run -f git-remote.yaml --dry-run +Error: git.remote is no longer supported: the snapshot archives your local copy, so a branch resolves to its local HEAD. To deploy a specific committed revision, use git.commit + +Exit code: 1 + === invalid config is rejected >>> [CLI] experimental air run -f invalid.yaml --dry-run Error: invalid experiment_name "bad.name": only alphanumeric characters, hyphens (-), and underscores (_) are allowed diff --git a/acceptance/experimental/air/run/script b/acceptance/experimental/air/run/script index 806bd321e6d..312b2f6fecf 100644 --- a/acceptance/experimental/air/run/script +++ b/acceptance/experimental/air/run/script @@ -10,6 +10,12 @@ errcode trace $CLI experimental air run -f valid.yaml --dry-run --override a=b title "watch not yet supported" errcode trace $CLI experimental air run -f valid.yaml --dry-run --watch +title "code_source config passes validation" +trace $CLI experimental air run -f with-code-source.yaml --dry-run + +title "git.remote is rejected" +errcode trace $CLI experimental air run -f git-remote.yaml --dry-run + title "invalid config is rejected" errcode trace $CLI experimental air run -f invalid.yaml --dry-run diff --git a/acceptance/experimental/air/run/with-code-source.yaml b/acceptance/experimental/air/run/with-code-source.yaml new file mode 100644 index 00000000000..86a32c138cb --- /dev/null +++ b/acceptance/experimental/air/run/with-code-source.yaml @@ -0,0 +1,13 @@ +experiment_name: smoke-test +command: python train.py +compute: + accelerator_type: GPU_1xH100 + num_accelerators: 1 +code_source: + type: snapshot + snapshot: + root_path: . + git: + branch: main + include_paths: + - src diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go index 5a0b6f27e1e..09437f50a5b 100644 --- a/experimental/air/cmd/runconfig.go +++ b/experimental/air/cmd/runconfig.go @@ -27,8 +27,8 @@ import ( // Jobs API task_key rejects. var taskKeyRe = regexp.MustCompile(`^[A-Za-z0-9_-]+$`) -// gitRefRe guards branch/remote names against command injection. Only safe ref -// characters are allowed. +// gitRefRe guards the branch name against command injection (it flows into git +// exec args). Only safe ref characters are allowed. var gitRefRe = regexp.MustCompile(`^[\w./-]+$`) // runConfig is the top-level run YAML schema: experiment_name + compute / @@ -372,13 +372,12 @@ func (g *gitRef) validate() error { if g.Branch != nil && !gitRefRe.MatchString(*g.Branch) { return fmt.Errorf("invalid git.branch format %q: only alphanumeric characters, hyphens, dots, slashes, and underscores are allowed", *g.Branch) } - if g.Remote.isString { - if g.Remote.name == "" { - return errors.New("git.remote string cannot be empty; use 'true' to auto-detect") - } - if !gitRefRe.MatchString(g.Remote.name) { - return fmt.Errorf("invalid git.remote name %q: only alphanumeric characters, hyphens, dots, slashes, and underscores are allowed", g.Remote.name) - } + + // The remote-fetch path (fetching a branch's remote HEAD) is deprecated: the + // snapshot archives the local copy only. A truthy git.remote (a name or `true`) + // is rejected; `remote: false` is the default (local HEAD) and stays valid. + if g.Remote.truthy() { + return errors.New("git.remote is no longer supported: the snapshot archives your local copy, so a branch resolves to its local HEAD. To deploy a specific committed revision, use git.commit") } if g.Branch == nil && g.Commit == nil { @@ -387,9 +386,6 @@ func (g *gitRef) validate() error { if g.Branch != nil && g.Commit != nil { return errors.New("git: 'branch' and 'commit' are mutually exclusive — specify only one") } - if g.Remote.truthy() && g.Branch == nil { - return errors.New("git.remote requires git.branch (only valid with branch refs)") - } return nil } diff --git a/experimental/air/cmd/runconfig_test.go b/experimental/air/cmd/runconfig_test.go index 1dd5ce1ea39..26b54127265 100644 --- a/experimental/air/cmd/runconfig_test.go +++ b/experimental/air/cmd/runconfig_test.go @@ -64,7 +64,6 @@ code_source: remote_volume: /Volumes/main/default/code git: branch: main - remote: origin include_paths: - src - configs/train.yaml @@ -92,8 +91,6 @@ permissions: require.NotNil(t, cfg.CodeSource.Snapshot.Git) require.NotNil(t, cfg.CodeSource.Snapshot.Git.Branch) assert.Equal(t, "main", *cfg.CodeSource.Snapshot.Git.Branch) - assert.True(t, cfg.CodeSource.Snapshot.Git.Remote.isString) - assert.Equal(t, "origin", cfg.CodeSource.Snapshot.Git.Remote.name) assert.Len(t, cfg.Permissions, 2) } @@ -111,8 +108,8 @@ environment: assert.Equal(t, "requirements.yaml", cfg.Environment.Dependencies.path) }) - t.Run("git remote as bool true", func(t *testing.T) { - cfg, err := loadRunConfig(writeConfig(t, minimalConfig+` + t.Run("git remote as bool true is rejected", func(t *testing.T) { + _, err := loadRunConfig(writeConfig(t, minimalConfig+` code_source: type: snapshot snapshot: @@ -121,11 +118,8 @@ code_source: branch: main remote: true `)) - require.NoError(t, err) - r := cfg.CodeSource.Snapshot.Git.Remote - assert.False(t, r.isString) - assert.True(t, r.enabled) - assert.True(t, r.truthy()) + require.Error(t, err) + assert.Contains(t, err.Error(), "git.remote is no longer supported") }) t.Run("git remote defaults to false when unset", func(t *testing.T) { @@ -334,12 +328,12 @@ func TestGitRefValidate(t *testing.T) { }{ {"branch only ok", gitRef{Branch: str("main")}, ""}, {"commit only ok", gitRef{Commit: str("abc123")}, ""}, - {"branch with remote ok", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: true}}, ""}, + {"remote false is ok", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: false}}, ""}, {"neither branch nor commit", gitRef{}, "must specify either 'branch' or 'commit'"}, {"both branch and commit", gitRef{Branch: str("main"), Commit: str("abc")}, "mutually exclusive"}, - {"remote without branch", gitRef{Commit: str("abc"), Remote: gitRemote{set: true, isString: true, name: "origin"}}, "requires git.branch"}, + {"remote true rejected", gitRef{Branch: str("main"), Remote: gitRemote{set: true, enabled: true}}, "git.remote is no longer supported"}, + {"remote name rejected", gitRef{Branch: str("main"), Remote: gitRemote{set: true, isString: true, name: "origin"}}, "git.remote is no longer supported"}, {"bad branch chars", gitRef{Branch: str("bad branch")}, "invalid git.branch"}, - {"empty remote string", gitRef{Branch: str("main"), Remote: gitRemote{set: true, isString: true, name: ""}}, "cannot be empty"}, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 49f51a8aa66..343c606146c 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -52,6 +52,13 @@ type aiRuntimeTask struct { Deployments []aiRuntimeDeployment `json:"deployments"` MlflowRun string `json:"mlflow_run,omitempty"` MlflowExperimentDirectory string `json:"mlflow_experiment_directory,omitempty"` + // CodeSourcePath points at the uploaded snapshot tarball; the remote + // entry_script extracts it to /databricks/code_source/. GitStatePath and + // GitDiffPath are the optional provenance sidecars (see snapshot.go). All are + // omitempty: a workspace-only run (no code_source) carries none of them. + CodeSourcePath string `json:"code_source_path,omitempty"` + GitStatePath string `json:"git_state_path,omitempty"` + GitDiffPath string `json:"git_diff_path,omitempty"` } // environmentSpec carries the bare runtime channel ("4", "5", ...). @@ -106,8 +113,9 @@ func dlRuntimeImage(ctx context.Context, runtimeVersion string) string { } // buildSubmitPayload assembles the runs/submit payload. commandPath is the -// workspace path of the uploaded command.sh; dlImage is the runtime channel. -func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string) jobsSubmitRun { +// workspace path of the uploaded command.sh; dlImage is the runtime channel; snap +// carries the code snapshot paths (zero value when the run has no code_source). +func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapshotResult) jobsSubmitRun { task := aiRuntimeTask{ Experiment: cfg.ExperimentName, Deployments: []aiRuntimeDeployment{{ @@ -117,6 +125,9 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string) jobsSubmitR AcceleratorCount: cfg.Compute.NumAccelerators, }, }}, + CodeSourcePath: snap.CodeSourcePath, + GitStatePath: snap.GitStatePath, + GitDiffPath: snap.GitDiffPath, } if cfg.MLflowRunName != nil { task.MlflowRun = *cfg.MLflowRunName @@ -195,14 +206,11 @@ func submitToken(flag string, cfg *runConfig) (string, error) { // upload the launch artifacts, assemble the Jobs payload, and submit it. It // returns the new run_id and its dashboard URL. func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) { - // Resolving usage_policy_name to a budget policy id and packaging a - // code_source snapshot are not ported yet; reject rather than silently drop. + // Resolving usage_policy_name to a budget policy id is not ported yet; reject + // rather than silently drop. if cfg.UsagePolicyName != nil { return 0, "", errors.New("usage_policy_name is not yet supported") } - if cfg.CodeSource != nil { - return 0, "", errors.New("code_source is not yet supported") - } // Resolve the idempotency token first so a bad key fails before any upload. token, err := submitToken(idempotencyKey, cfg) @@ -240,8 +248,18 @@ 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. + var snap snapshotResult + if cfg.CodeSource != nil { + snap, err = snapshotCodeSource(ctx, w, cfg.CodeSource.Snapshot, configPath, base, funcDir) + if err != nil { + return 0, "", err + } + } + runtimeVersion, _ := cfg.runtimeVersion() - payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion)) + payload := buildSubmitPayload(cfg, path.Join(funcDir, commandScriptName), dlRuntimeImage(ctx, runtimeVersion), snap) payload.IdempotencyToken = token jc, err := newJobsSubmitClient(w) diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index d57b14f80a6..35cd6dd8526 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -2,6 +2,7 @@ package aircmd import ( "encoding/json" + "path/filepath" "strings" "testing" @@ -36,7 +37,7 @@ func TestBuildSubmitPayload(t *testing.T) { MLflowExperimentDirectory: new("/Workspace/Users/me/exp"), } - p := buildSubmitPayload(cfg, "/d/command.sh", "5") + p := buildSubmitPayload(cfg, "/d/command.sh", "5", snapshotResult{}) assert.Equal(t, "exp", p.RunName) assert.Equal(t, 1800, p.TimeoutSeconds) @@ -69,7 +70,7 @@ func TestBuildSubmitPayloadDefaultRetries(t *testing.T) { Command: new("x"), Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, } - task := buildSubmitPayload(cfg, "/d/command.sh", "4").Tasks[0] + task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] assert.Equal(t, defaultMaxRetries, task.MaxRetries) assert.True(t, task.RetryOnTimeout) } @@ -84,7 +85,7 @@ func TestBuildSubmitPayloadNoRetries(t *testing.T) { Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, MaxRetries: new(0), } - task := buildSubmitPayload(cfg, "/d/command.sh", "4").Tasks[0] + task := buildSubmitPayload(cfg, "/d/command.sh", "4", snapshotResult{}).Tasks[0] assert.Equal(t, 0, task.MaxRetries) assert.False(t, task.RetryOnTimeout) @@ -170,6 +171,50 @@ func TestSubmitWorkload(t *testing.T) { assert.Equal(t, aiRuntimeCompute{AcceleratorType: "GPU_1xH100", 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. +func TestSubmitWorkloadWithCodeSource(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + testserver.AddDefaultHandlers(server) + + var got jobsSubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return submitRunResponse{RunID: 555} + }) + 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) + 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") + require.NoError(t, err) + + at := got.Tasks[0].AiRuntimeTask + // The tarball path is under the user's repo_snapshots dir; the clean repo also + // attaches a git_state sidecar (no diff). + assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/"+filepath.Base(repo)+"/") + assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) + assert.True(t, strings.HasSuffix(at.GitStatePath, "/"+gitStateName), at.GitStatePath) + assert.Empty(t, at.GitDiffPath) +} + func TestSubmitWorkloadGuards(t *testing.T) { w := newFakeWorkspaceClient(t) cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) @@ -182,11 +227,4 @@ func TestSubmitWorkloadGuards(t *testing.T) { _, _, err := submitWorkload(t.Context(), w, &cfg, cfgPath, "") require.ErrorContains(t, err, "usage_policy_name is not yet supported") }) - - t.Run("code_source rejected", func(t *testing.T) { - cfg := *base - cfg.CodeSource = &codeSourceConfig{Type: "snapshot"} - _, _, err := submitWorkload(t.Context(), w, &cfg, cfgPath, "") - require.ErrorContains(t, err, "code_source is not yet supported") - }) } diff --git a/experimental/air/cmd/snapshot.go b/experimental/air/cmd/snapshot.go new file mode 100644 index 00000000000..aba67f6109c --- /dev/null +++ b/experimental/air/cmd/snapshot.go @@ -0,0 +1,279 @@ +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). +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. +func resolveRootPath(ctx context.Context, rawPath, configDir string) (string, error) { + expanded := os.ExpandEnv(rawPath) + if home, err := env.UserHomeDir(ctx); err == nil { + if expanded == "~" { + expanded = home + } else if rest, ok := strings.CutPrefix(expanded, "~/"); ok { + expanded = filepath.Join(home, rest) + } + } + + var resolved string + switch { + case strings.HasPrefix(expanded, "project_root/"): + resolved = filepath.Join(configDir, strings.TrimPrefix(expanded, "project_root/")) + case filepath.IsAbs(expanded): + resolved = expanded + default: + 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) + } + resolved = abs + + info, err := os.Stat(resolved) + if err != nil { + return "", fmt.Errorf("root_path does not exist: %s", resolved) + } + if !info.IsDir() { + return "", fmt.Errorf("root_path must be a directory: %s", resolved) + } + 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 new file mode 100644 index 00000000000..44c58ee903b --- /dev/null +++ b/experimental/air/cmd/snapshot_cachekey.go @@ -0,0 +1,34 @@ +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 new file mode 100644 index 00000000000..5743217c003 --- /dev/null +++ b/experimental/air/cmd/snapshot_cachekey_test.go @@ -0,0 +1,62 @@ +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_git.go b/experimental/air/cmd/snapshot_git.go new file mode 100644 index 00000000000..616b3049f74 --- /dev/null +++ b/experimental/air/cmd/snapshot_git.go @@ -0,0 +1,309 @@ +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 new file mode 100644 index 00000000000..cf1a821b1d6 --- /dev/null +++ b/experimental/air/cmd/snapshot_git_test.go @@ -0,0 +1,296 @@ +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 new file mode 100644 index 00000000000..672366086c9 --- /dev/null +++ b/experimental/air/cmd/snapshot_package.go @@ -0,0 +1,132 @@ +package aircmd + +import ( + "bytes" + "context" + "fmt" + "os" + "os/exec" + "path/filepath" + "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 +} + +// createPlainTarball writes a gzipped tar of repoPath's working tree 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) error { + dirName := filepath.Base(repoPath) + parent := filepath.Dir(repoPath) + + args := []string{"-czf", outputTarball} + + // Exclude macOS AppleDouble files: they sort before the real top-level dir and + // hijack a remote `head -1` parse. No-op on Linux. + args = append(args, "--exclude=._*") + + // Never ship .git — provenance flows via the git_state.json sidecar. + args = append(args, "--exclude=.git") + + // Honor .gitignore if present. + gitignorePath := filepath.Join(repoPath, ".gitignore") + if patterns, err := parseGitignore(gitignorePath); err == nil { + for _, p := range patterns { + if strings.Contains(p, "/") { + // Anchor path-relative patterns to the archive root so they don't + // match identically-named paths in subdirectories. + args = append(args, "--exclude="+dirName+"/"+strings.TrimPrefix(p, "/")) + } else { + args = append(args, "--exclude="+p) + } + } + } + + // Archive from the parent so the directory name is preserved; with include_paths, + // prefix each so entries nest under it (matching git archive --prefix). + args = append(args, "-C", parent) + if len(includePaths) > 0 { + for _, p := range includePaths { + args = append(args, dirName+"/"+p) + } + } else { + args = append(args, dirName) + } + + cmd := exec.CommandContext(ctx, "tar", args...) + var stderr bytes.Buffer + cmd.Stderr = &stderr + if err := cmd.Run(); err != nil { + if msg := strings.TrimSpace(stderr.String()); msg != "" { + return fmt.Errorf("failed to create plain tarball: %w: %s", err, msg) + } + return fmt.Errorf("failed to create plain tarball: %w", err) + } + return nil +} + +// parseGitignore reads a .gitignore and returns tar --exclude patterns. It mirrors +// the Python CLI's lossy normalization so plain-tar snapshots exclude the same set: +// +// - comments (#…) and blank lines are skipped; +// - negation patterns (!…) are unsupported by tar --exclude and skipped; +// - a trailing "/" (directory marker) is stripped; +// - "**" is not a path-separator-agnostic wildcard in tar, so "**/foo" → "foo" +// and "foo/**" → "foo"; a mid-path "**" has no tar equivalent and is skipped. +// +// A missing file returns (nil, error); callers treat any error as "no patterns". +func parseGitignore(path string) ([]string, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + + var patterns []string + for raw := range strings.SplitSeq(string(data), "\n") { + line := strings.TrimRight(raw, " \t\r") + if line == "" || strings.HasPrefix(line, "#") { + continue + } + if strings.HasPrefix(line, "!") { + continue + } + line = strings.TrimRight(line, "/") + if strings.Contains(line, "**") { + switch { + case strings.HasPrefix(line, "**/"): + line = line[len("**/"):] + case strings.HasSuffix(line, "/**"): + line = line[:len(line)-len("/**")] + default: + continue + } + } + patterns = append(patterns, line) + } + return patterns, nil +} diff --git a/experimental/air/cmd/snapshot_package_test.go b/experimental/air/cmd/snapshot_package_test.go new file mode 100644 index 00000000000..d895d59b98e --- /dev/null +++ b/experimental/air/cmd/snapshot_package_test.go @@ -0,0 +1,163 @@ +package aircmd + +import ( + "archive/tar" + "compress/gzip" + "os" + "path/filepath" + "slices" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// tarballEntries returns the sorted list of entry names in a .tar.gz. +func tarballEntries(t *testing.T, path 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() + + var names []string + tr := tar.NewReader(gz) + for { + hdr, err := tr.Next() + if err != nil { + break + } + names = append(names, hdr.Name) + } + slices.Sort(names) + 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") +} + +func TestCreatePlainTarball(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + 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") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + require.NoError(t, createPlainTarball(ctx, repo, out, nil)) + + dirName := filepath.Base(repo) + entries := tarballEntries(t, out) + assert.Contains(t, entries, dirName+"/a.txt") + assert.Contains(t, entries, dirName+"/dirty.txt") + // .git is never shipped. + for _, e := range entries { + assert.NotContains(t, e, "/.git/") + } +} + +func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "keep.txt", "1") + writeRepoFile(t, repo, "junk.log", "noise") + writeRepoFile(t, repo, ".gitignore", "*.log\n") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + require.NoError(t, createPlainTarball(ctx, repo, out, nil)) + + dirName := filepath.Base(repo) + entries := tarballEntries(t, out) + assert.Contains(t, entries, dirName+"/keep.txt") + assert.NotContains(t, entries, dirName+"/junk.log") +} + +func TestCreatePlainTarball_IncludePaths(t *testing.T) { + ctx := t.Context() + repo := newTestRepo(t) + writeRepoFile(t, repo, "a.txt", "1") + writeRepoFile(t, repo, "src/model.py", "print()") + + out := filepath.Join(t.TempDir(), "snap.tar.gz") + require.NoError(t, createPlainTarball(ctx, repo, out, []string{"src"})) + + dirName := filepath.Base(repo) + entries := tarballEntries(t, out) + assert.Contains(t, entries, dirName+"/src/model.py") + assert.NotContains(t, entries, dirName+"/a.txt") +} + +func TestParseGitignore(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, ".gitignore") + content := "# comment\n" + + "\n" + + "*.log\n" + + "!keep.log\n" + // negation: skipped + "build/\n" + // trailing slash stripped + "**/node_modules\n" + // **/foo -> foo + "dist/**\n" + // foo/** -> foo + "a/**/b\n" + // mid ** : skipped + "src/config\n" // path-relative kept as-is + require.NoError(t, os.WriteFile(path, []byte(content), 0o600)) + + patterns, err := parseGitignore(path) + require.NoError(t, err) + assert.Equal(t, []string{ + "*.log", + "build", + "node_modules", + "dist", + "src/config", + }, patterns) +} + +func TestParseGitignore_Missing(t *testing.T) { + _, err := parseGitignore(filepath.Join(t.TempDir(), "nope")) + require.Error(t, err) +} diff --git a/experimental/air/cmd/snapshot_resolve.go b/experimental/air/cmd/snapshot_resolve.go new file mode 100644 index 00000000000..1146b641b92 --- /dev/null +++ b/experimental/air/cmd/snapshot_resolve.go @@ -0,0 +1,120 @@ +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 new file mode 100644 index 00000000000..c8c946f8394 --- /dev/null +++ b/experimental/air/cmd/snapshot_resolve_test.go @@ -0,0 +1,114 @@ +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 new file mode 100644 index 00000000000..d94fe005fc9 --- /dev/null +++ b/experimental/air/cmd/snapshot_test.go @@ -0,0 +1,155 @@ +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...) +} diff --git a/experimental/air/cmd/testdata/cache_keys.json b/experimental/air/cmd/testdata/cache_keys.json new file mode 100644 index 00000000000..06673499109 --- /dev/null +++ b/experimental/air/cmd/testdata/cache_keys.json @@ -0,0 +1,77 @@ +[ + { + "name": "commit_no_paths", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": null, + "cache_key": "8d7bc445ac83dfe432a353718ae8b1ae40eb00c860352662e6df96b7fdf862a5" + }, + { + "name": "commit_empty_paths_none", + "commit_sha": "0000000000000000000000000000000000000000", + "include_paths": null, + "cache_key": "a07e0fa4d38f3c2d08c3ff6bbff0158cea01796aaaf5368b0d0e247cebd27c80" + }, + { + "name": "single_path", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "research" + ], + "cache_key": "97bd562ad591af23b9a05651d88c01abbbcc24632e9a8f2534230a61492b9e0a" + }, + { + "name": "multi_path_sorted", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "a", + "b", + "c" + ], + "cache_key": "17d021577e3615e9a603a2f701f918d74ad16218efc4b2950ee20f6a18b1a174" + }, + { + "name": "multi_path_unsorted", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "c", + "a", + "b" + ], + "cache_key": "17d021577e3615e9a603a2f701f918d74ad16218efc4b2950ee20f6a18b1a174" + }, + { + "name": "path_whitespace", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + " research ", + " data " + ], + "cache_key": "c6a6fa1d866ac90a5f371610ee08efa8aeb061a43b8a4acae30765db21f4155b" + }, + { + "name": "nested_paths", + "commit_sha": "deadbeefdeadbeefdeadbeefdeadbeefdeadbeef", + "include_paths": [ + "src/models", + "src/data", + "README.md" + ], + "cache_key": "2b18452b1019186a8e5b34af540a0545f96f24bc61771f58fa71fc218cb0275c" + }, + { + "name": "short_sha", + "commit_sha": "a3492b8", + "include_paths": null, + "cache_key": "16be4d7d83ecd01b8d3d274dafb07b3bccc72cf7f7299578e634cd01e1825345" + }, + { + "name": "dup_paths", + "commit_sha": "a3492b801c0ffee00000000000000000000dead", + "include_paths": [ + "x", + "x", + "y" + ], + "cache_key": "407a5d8c3a0e9438d5abf653c4d3a43928c300b5630e2f83f435661b9fbd599e" + } +]