diff --git a/acceptance/experimental/air/run-submit/output.txt b/acceptance/experimental/air/run-submit/output.txt index 8d52eed1dab..333b13f294f 100644 --- a/acceptance/experimental/air/run-submit/output.txt +++ b/acceptance/experimental/air/run-submit/output.txt @@ -1,6 +1,7 @@ === submit with a git code_source >>> [CLI] experimental air run -f run.yaml +Uploading [SNAPSHOT_TARBALL]... Submitted run 555 View at: [DATABRICKS_URL]/jobs/runs/555 @@ -23,7 +24,7 @@ View at: [DATABRICKS_URL]/jobs/runs/555 "tasks": [ { "ai_runtime_task": { - "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/001/[SNAPSHOT_TARBALL]", + "code_source_path": "/Workspace/Users/[USERNAME]/.air/repo_snapshots/.internal/[SNAPSHOT_TARBALL]", "deployments": [ { "command_path": "/Workspace/Users/[USERNAME]/.air/cli_launch/submit-smoke/submit-smoke_[RUN_ID]/command.sh", diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 8c0be55260d..d29bea427ba 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -58,15 +58,6 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string, snap snapsh }, }}, CodeSourcePath: snap.CodeSourcePath, - // TEMP: git_state_path / git_diff_path are intentionally NOT sent. The typed - // jobs.AiRuntimeTask (and its source proto, ai_runtime_task.proto) has no such - // fields, so the typed SDK path cannot carry them. This is safe today because - // nothing in the backend consumes those fields — the AI Runtime task proto - // never declared them, so even the Python CLI's raw-JSON values were dropped - // on deserialization. The git_state.json / git_diff.patch sidecars are still - // uploaded next to the tarball (see snapshot.go) for human inspection. - // If the backend later adds these fields to the proto, regenerate the SDK and - // wire snap.GitStatePath / snap.GitDiffPath back in here. } if cfg.MLflowRunName != nil { task.MlflowRun = *cfg.MLflowRunName @@ -163,13 +154,13 @@ 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) + // Sidecars land in the run's launch dir (funcDir) via fc, next to command.sh. + snap, err = snapshotViaDABsUpload(ctx, w, cfg.CodeSource.Snapshot, configPath, fc, funcDir) if err != nil { return 0, "", err } diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index fd5103599df..f903ba76f2c 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -2,10 +2,14 @@ package aircmd import ( "encoding/json" + "io" + "path" "path/filepath" "strings" "testing" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/testserver" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/service/jobs" @@ -156,8 +160,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,6 +176,47 @@ func TestSubmitWorkloadWithCodeSource(t *testing.T) { w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) require.NoError(t, err) + // A plain working-tree directory: packaging is plain-tar. + repo := filepath.Join(t.TempDir(), "src") + writeRepoFile(t, repo, "train.py", "print()") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + // 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 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) +} + +// A git-pinned code_source is git-archived at the commit, uploaded via DABs' artifact +// plumbing, and its remote code_source_path attached to the submitted task. +func TestSubmitWorkloadWithGitPinnedCodeSource(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + 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()") @@ -189,15 +234,209 @@ code_source: loaded, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(t.Context(), w, loaded, cfgPath, "idem") + ctx := cmdio.MockDiscard(t.Context()) + _, _, err = submitWorkload(ctx, w, loaded, cfgPath, "idem") + require.NoError(t, err) + + at := got.Tasks[0].AiRuntimeTask + assert.Contains(t, at.CodeSourcePath, "/.air/repo_snapshots/.internal/") + assert.True(t, strings.HasSuffix(at.CodeSourcePath, ".tar.gz"), at.CodeSourcePath) +} + +// testSidecarStore builds a workspace filer + base path standing in for the run's +// launch dir, where snapshotViaDABsUpload writes git provenance sidecars. +func testSidecarStore(t *testing.T, w *databricks.WorkspaceClient) (filer.Filer, string) { + t.Helper() + base := "/Workspace/Users/tester@databricks.com/.air/cli_launch/test" + f, err := filer.NewWorkspaceFilesClient(w, base) + require.NoError(t, err) + return f, base +} + +// A plain-tar (working-tree) snapshot is uploaded under a unique, timestamped name so +// two concurrent submissions of the same root_path don't clobber each other's upload. +func TestSubmitWorkloadPlainTarNameIsUnique(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + // A plain working-tree directory named "src": the old code named the tarball + // after the dir alone (src.tar.gz), so any two submissions collided. + repo := filepath.Join(t.TempDir(), "src") + writeRepoFile(t, repo, "train.py", "print()") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + // The uploaded name carries a discriminator (timestamp), not the bare dir name. + ctx := cmdio.MockDiscard(t.Context()) + sidecarStore, sidecarBase := testSidecarStore(t, w) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + base := path.Base(snap.CodeSourcePath) + assert.NotEqual(t, "src.tar.gz", base, "plain-tar name must be unique, not the bare dir name") + assert.Regexp(t, `^src_\d{8}_\d{6}\.tar\.gz$`, base) +} + +// A git_archive snapshot is content-addressed by (commit, include_paths): submitting +// the same commit twice reuses the already-uploaded tarball and skips the second +// upload (cache hit), while resolving to the identical remote path. +func TestSubmitWorkloadGitArchiveCaching(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 555} + }) + // Track which snapshot tarballs get uploaded, preserving fake-workspace + // persistence so the second submit's cache-existence Stat sees the first upload. + // Dedupe by path: the DABs uploader mkdirs-and-retries the import on a missing + // parent dir, so one logical upload can hit this route more than once. + uploaded := map[string]bool{} + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { + p := req.Vars["path"] + if strings.Contains(p, "/.air/repo_snapshots/") { + uploaded[p] = true + } + return req.Workspace.WorkspaceFilesImportFile(p, req.Body, req.URL.Query().Get("overwrite") == "true") + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + 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) + + ctx := cmdio.MockDiscard(t.Context()) + sidecarStore, sidecarBase := testSidecarStore(t, w) + first, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + second, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + + // Same pinned commit → identical content-addressed remote path, uploaded once + // (the second submit is a cache hit and moves no bytes). + assert.Equal(t, first.CodeSourcePath, second.CodeSourcePath) + assert.Len(t, uploaded, 1, "git_archive cache hit should skip the second upload") +} + +// A git code_source also uploads git provenance sidecars (git_state.json, and +// git_diff.patch when the tree is dirty) next to the run's launch dir, so the +// submitted commit + working-tree diff are inspectable. +func TestSubmitWorkloadUploadsGitSidecars(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 555} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + // Commit, then dirty the tree so both git_state.json and git_diff.patch are produced. + repo := newTestRepo(t) + writeRepoFile(t, repo, "train.py", "print()") + commitAll(t, repo, "init") + writeRepoFile(t, repo, "train.py", "print('dirty')") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + ctx := cmdio.MockDiscard(t.Context()) + sidecarStore, sidecarBase := testSidecarStore(t, w) + snap, err := snapshotViaDABsUpload(ctx, w, loaded.CodeSource.Snapshot, cfgPath, sidecarStore, sidecarBase) + require.NoError(t, err) + + // Both sidecars are reported under the launch dir and actually exist there. + assert.Equal(t, path.Join(sidecarBase, gitStateName), snap.GitStatePath) + assert.Equal(t, path.Join(sidecarBase, gitDiffName), snap.GitDiffPath) + + r, err := sidecarStore.Read(ctx, gitStateName) + require.NoError(t, err) + stateBytes, err := io.ReadAll(r) + require.NoError(t, err) + var state map[string]any + require.NoError(t, json.Unmarshal(stateBytes, &state)) + assert.Equal(t, "plain_tar", state["packaging_mode"]) + assert.Equal(t, true, state["dirty"]) + assert.Equal(t, "captured", state["diff_status"]) +} + +// remote_volume uploads the snapshot to a UC Volume: DABs' artifact uploader handles +// /Volumes destinations natively, so code_source_path lands under the Volume path. +func TestSubmitWorkloadWithRemoteVolumeCodeSource(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + return jobs.SubmitRunResponse{RunId: 555} + }) + // Stub the UC Volume file write: the fake server's default handler 404s when the + // parent dir is absent (no auto-mkdir), so accept the PUT to exercise the Volume + // upload route. This asserts we route to /api/2.0/fs/files/Volumes/... at all. + server.Handle("PUT", "/api/2.0/fs/files/Volumes/{path...}", func(req testserver.Request) any { + return testserver.Response{StatusCode: 204} + }) + testserver.AddDefaultHandlers(server) + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + repo := filepath.Join(t.TempDir(), "src") + writeRepoFile(t, repo, "train.py", "print()") + + cfg := minimalConfig + ` +code_source: + type: snapshot + snapshot: + root_path: ` + repo + ` + remote_volume: /Volumes/main/default/code +` + cfgPath := writeConfigFile(t, "run.yaml", cfg) + loaded, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + 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)+"/") + assert.Contains(t, at.CodeSourcePath, "/Volumes/main/default/code/.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 59041f5986d..603b7e6c0f8 100644 --- a/experimental/air/cmd/snapshot.go +++ b/experimental/air/cmd/snapshot.go @@ -1,58 +1,28 @@ 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) plus the remote paths of the best-effort git +// provenance sidecars (empty when not a git repo or upload failed). 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(ctx, 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 +43,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 +58,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(ctx context.Context, 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(ctx, 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_dabs.go b/experimental/air/cmd/snapshot_dabs.go new file mode 100644 index 00000000000..9becc5e91a2 --- /dev/null +++ b/experimental/air/cmd/snapshot_dabs.go @@ -0,0 +1,278 @@ +package aircmd + +import ( + "bytes" + "context" + "errors" + "fmt" + "io/fs" + "os" + "path" + "path/filepath" + "time" + + "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/filer" + "github.com/databricks/cli/libs/log" + "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, sidecarStore filer.Filer, sidecarBase string) (snapshotResult, error) { + repoPath, err := resolveRootPath(ctx, snap.RootPath, filepath.Dir(configPath)) + if err != nil { + return snapshotResult{}, err + } + + // Resolve how to package before touching the tarball: git_archive (pinned commit, + // cacheable) vs plain_tar (working tree, not cacheable). + plan, err := resolveSnapshotPlan(ctx, newGitRepo(repoPath), snap.Git, snap.IncludePaths) + if err != nil { + return snapshotResult{}, err + } + + // remote_volume, when set, is a UC Volume path; DABs' artifact uploader handles + // /Volumes destinations natively (GetFilerForLibraries → filerForVolume). + remoteVolume := "" + if snap.RemoteVolume != nil { + remoteVolume = *snap.RemoteVolume + } + result, err := uploadSnapshotViaDABs(ctx, w, repoPath, plan, remoteVolume) + if err != nil { + return snapshotResult{}, err + } + + // Upload git provenance sidecars (git_state.json / git_diff.patch) next to the + // run's launch dir so the submitted commit + working-tree diff are inspectable. + // Best-effort and git-only: any failure logs and leaves the paths empty rather + // than failing an otherwise-valid submission. + // + // The sidecars are deliberately NOT bundled into the code tarball. The git_archive + // tarball is content-addressed and cached by (commit, include_paths), so a second + // run at the same commit reuses it; but the sidecars vary per run (git_state's + // timestamp, and git_diff captures the working tree at submit time). Folding them + // in would force a distinct tarball per run (defeating the cache) or serve a prior + // run's stale provenance on a cache hit. They also live in the per-run launch dir, + // not the shared artifact dir, so they don't accumulate. Keep them out of the tar. + if plan.isGitRepo { + result.GitStatePath, result.GitDiffPath = uploadSnapshotSidecars(ctx, sidecarStore, sidecarBase, newGitRepo(repoPath), plan) + } + return result, nil +} + +// uploadSnapshotSidecars writes the git_state.json provenance record — and, when the +// working tree is dirty, a captured git_diff.patch — into the run's launch dir via +// sidecarStore (rooted at sidecarBase, used only to report absolute paths). It is +// best-effort: every failure logs a warning and yields an empty path, never an error, +// so provenance capture cannot fail a submission. +func uploadSnapshotSidecars(ctx context.Context, sidecarStore filer.Filer, sidecarBase string, 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 := 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(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 := 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(sidecarBase, gitStateName), diffPath +} + +// snapshotTarballName is the uploaded filename for the snapshot. It is deterministic +// for git_archive — _.tar.gz keyed on (commit, include_paths) — so +// an identical commit reuses the same remote object (see the cache check below). For +// plain_tar it is timestamped so concurrent submissions of the same directory don't +// clobber each other's upload (working-tree content isn't pinned to a SHA, so it +// can't be content-addressed). +func snapshotTarballName(plan snapshotPlan, dirName string) string { + if plan.mode == modeGitArchive { + key := computeSnapshotCacheKey(plan.commitSHA, plan.includePaths) + return fmt.Sprintf("%s_%s.tar.gz", dirName, key[:16]) + } + return fmt.Sprintf("%s_%s.tar.gz", dirName, time.Now().UTC().Format("20060102_150405")) +} + +// packageSnapshot writes the snapshot to tarball per the resolved plan: `git archive` +// of the pinned commit for git_archive, else a plain tar of the working tree. +func packageSnapshot(ctx context.Context, repoPath string, plan snapshotPlan, tarball string) error { + dirName := filepath.Base(repoPath) + if plan.mode == modeGitArchive { + return createGitArchiveSnapshot(ctx, newGitRepo(repoPath), plan.commitSHA, tarball, dirName, plan.includePaths) + } + return createPlainTarball(ctx, repoPath, tarball, plan.includePaths) +} + +// uploadSnapshotViaDABs uploads the snapshot 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. When remoteVolume is set the +// tarball goes to that UC Volume; otherwise to the user's repo_snapshots dir. +// +// git_archive snapshots are cacheable: the tarball name is content-addressed by +// (commit, include_paths), so if the identical object is already uploaded we skip +// packaging and upload entirely and just reuse the remote path. +func uploadSnapshotViaDABs(ctx context.Context, w *databricks.WorkspaceClient, repoPath string, plan snapshotPlan, remoteVolume string) (snapshotResult, error) { + // artifactPath is where DABs uploads the tarball; GetFilerForLibraries routes to + // a Workspace or Volume filer based on its prefix, then appends /.internal. + artifactPath := remoteVolume + if artifactPath == "" { + base, err := userWorkspaceDir(ctx, w) + if err != nil { + return snapshotResult{}, err + } + // The user's repo_snapshots dir (not the default bundle artifact_path, which a + // deploy would clean up). + artifactPath = path.Join(base, ".air", "repo_snapshots") + } + + tmp, err := os.MkdirTemp("", "air-snapshot-*") + if err != nil { + return snapshotResult{}, err + } + defer os.RemoveAll(tmp) + + tarName := snapshotTarballName(plan, filepath.Base(repoPath)) + + b := &bundle.Bundle{ + BundleRootPath: tmp, + BundleRoot: vfs.MustNew(tmp), + SyncRootPath: tmp, + SyncRoot: vfs.MustNew(tmp), + 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 temp dir); collectLocalLibraries + // joins it back and uploads the file. + AiRuntimeTask: &jobs.AiRuntimeTask{CodeSourcePath: tarName}, + }}, + }, + }, + }, + }, + }, + } + b.SetWorkpaceClient(w) + if err := b.Config.Mutate(func(v dyn.Value) (dyn.Value, error) { return v, nil }); err != nil { + return snapshotResult{}, err + } + + // git_archive is cacheable by (commit, include_paths): if the identical tarball is + // already uploaded, skip packaging + upload and reuse it. Only the config-path + // rewrite (ReplaceWithRemotePath) runs — no bytes move. + if plan.mode == modeGitArchive { + f, uploadPath, diags := libraries.GetFilerForLibraries(ctx, b) + if diags.HasError() { + return snapshotResult{}, diags.Error() + } + exists, err := snapshotExists(ctx, f, tarName) + if err != nil { + return snapshotResult{}, err + } + if exists { + if _, diags := libraries.ReplaceWithRemotePath(ctx, b); diags.HasError() { + return snapshotResult{}, diags.Error() + } + remote, err := readCodeSourcePath(b) + if err != nil { + return snapshotResult{}, err + } + log.Debugf(ctx, "snapshot cache hit for %s at %s", shortSHA(plan.commitSHA), path.Join(uploadPath, tarName)) + return snapshotResult{CodeSourcePath: remote}, nil + } + } + + // Cache miss (or plain_tar): package the tarball locally, then upload the bytes. + if err := packageSnapshot(ctx, repoPath, plan, filepath.Join(tmp, tarName)); err != nil { + return snapshotResult{}, err + } + + 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 +} + +// snapshotExists reports whether name already exists in the artifact store, used to +// short-circuit a cacheable git_archive upload. A not-found is a clean miss (false, +// nil); any other error is surfaced. +func snapshotExists(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) +} + +// 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 index 616b3049f74..70dd34f4da9 100644 --- a/experimental/air/cmd/snapshot_git.go +++ b/experimental/air/cmd/snapshot_git.go @@ -200,6 +200,13 @@ func shortSHA(sha string) string { // coordination with the backend reader. const snapshotStateSchemaVersion = 1 +// gitStateName and gitDiffName are the git provenance sidecar basenames, uploaded +// next to the code snapshot for human/agent inspection of what was submitted. +const ( + gitStateName = "git_state.json" + gitDiffName = "git_diff.patch" +) + // defaultRemoteName is the remote consulted for merge-base and repo URL (local refs // only — the remote-fetch path is gone). const defaultRemoteName = "origin" diff --git a/experimental/air/cmd/snapshot_package.go b/experimental/air/cmd/snapshot_package.go index 672366086c9..dd043fdfa50 100644 --- a/experimental/air/cmd/snapshot_package.go +++ b/experimental/air/cmd/snapshot_package.go @@ -19,8 +19,6 @@ import ( // `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", diff --git a/experimental/air/cmd/snapshot_package_test.go b/experimental/air/cmd/snapshot_package_test.go index d895d59b98e..22561f4e2ce 100644 --- a/experimental/air/cmd/snapshot_package_test.go +++ b/experimental/air/cmd/snapshot_package_test.go @@ -49,9 +49,8 @@ func TestCreateGitArchiveSnapshot(t *testing.T) { 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. + // Every real entry is prefixed with the directory name. git archive also emits a + // `pax_global_header` pseudo-entry (no prefix) that tar ignores on extraction. assert.Contains(t, entries, dirName+"/a.txt") assert.Contains(t, entries, dirName+"/src/model.py") for _, e := range entries { @@ -75,18 +74,17 @@ func TestCreateGitArchiveSnapshot_IncludePaths(t *testing.T) { 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) + 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 +101,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 +117,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_test.go b/experimental/air/cmd/snapshot_test.go deleted file mode 100644 index d94fe005fc9..00000000000 --- 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...) -}