air run: package code_source tarball via DABs artifact upload by vinchenzo-db · Pull Request #6015 · databricks/cli

vinchenzo-db · 2026-07-21T18:59:35Z

Package the code_source working tree into a tarball and upload it through DABs' artifact-upload plumbing (libraries.ReplaceWithRemotePath + libraries.Upload over a minimal in-memory bundle), rewriting ai_runtime_task.code_source_path to the uploaded remote path. The packaging + upload orchestration is CLI-owned (experimental/air/cmd, OWNERS = us); it only reuses DABs' uploader so we don't reimplement workspace/volume upload.

snapshot_dabs.go: build the plain-tar tarball (createPlainTarball), carry it as a file-valued code_source_path on a minimal bundle, and drive the DABs upload. runsubmit.go swaps the old raw-filer snapshot upload for this. Removes the retired raw-filer upload path (snapshot.go uploader, snapshot_test.go).

Tar snapshotting only; git pinning follows in the next PR (its git helpers are removed here and reintroduced there).

Co-authored-by: Isaac

Changes

Why

Testing

Unit + acceptance

experimental/air/cmd/... and acceptance/experimental/air/run-submit — working-tree,
git-pinned, and remote-Volume submits each assert the tarball lands under .internal/
and the rewritten code_source_path rides the submitted ai_runtime_task; plus the tar
builders (.gitignore, .git exclusion, include_paths) and the no-code_source
nil-guard. All green.

Live E2E — staging `dbc-04ac0685-8857` (GPU_1xA10)

5/5 runs SUCCESS, one per packaging mode. All runs are CAN_VIEW for the workspace
users group, so every link below is openable by anyone in the workspace.

Setup — a tiny project with a gitignored file (debug.log) to prove exclusion:

mkdir -p /tmp/air-demo/proj/src/pkg && cd /tmp/air-demo/proj
printf 'import os\nprint("train ran; cwd:", os.getcwd(), "CODE_SOURCE_PATH:", os.environ.get("CODE_SOURCE_PATH"))\n' > src/train.py
echo 'def helper(): return 1' > src/pkg/util.py
echo '*.log' > src/.gitignore
echo 'noise'  > src/debug.log          # gitignored — must never be uploaded
cat > wt.yaml <<'YAML'
experiment_name: vchen_demo_wt
command: cd $CODE_SOURCE_PATH && python train.py
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
code_source: {type: snapshot, snapshot: {root_path: src}}
YAML

1. Working-tree tarball — plain tar of the working tree, honoring .gitignore.
Run: 994765091508414

dbcli experimental air run -f wt.yaml -p dbc-04ac0685-8857
# Uploading src.tar.gz... → Submitted run 994765091508414
dbcli workspace export /Workspace/Users/v.chen@databricks.com/.air/repo_snapshots/.internal/src.tar.gz --file /tmp/wt.tar.gz -p dbc-04ac0685-8857
tar tzf /tmp/wt.tar.gz
# → src/train.py, src/pkg/util.py, src/.gitignore   (NO debug.log ✓ gitignore honored)

2. Git-pinned commit — git archive of a pinned SHA. An uncommitted file is created
after the commit to prove the archive captures the commit, not the dirty working tree.
Run: 304463075281818

git init -q && git add -A && git commit -qm init
SHA=$(git rev-parse HEAD)
echo "print('uncommitted')" > src/uncommitted.py     # created AFTER the commit
cat > git.yaml <<YAML
experiment_name: vchen_demo_git
command: cd \$CODE_SOURCE_PATH && python train.py
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
code_source: {type: snapshot, snapshot: {root_path: src, git: {commit: $SHA}}}
YAML
dbcli experimental air run -f git.yaml -p dbc-04ac0685-8857
# Uploading src.tar.gz... → Submitted run 304463075281818
dbcli workspace export /Workspace/Users/v.chen@databricks.com/.air/repo_snapshots/.internal/src.tar.gz --file /tmp/git.tar.gz -p dbc-04ac0685-8857
tar tzf /tmp/git.tar.gz
# → src/train.py, src/pkg/util.py, src/.gitignore
#   (NO uncommitted.py, NO debug.log ✓ archived the commit, not the working tree)

3. UC Volume destination — remote_volume routes the upload to a UC Volume via the
Files API (/api/2.0/fs/files/...), natively, with no special-casing in the CLI.
Run: 438683652713410

dbcli volumes create main default vchen_demo MANAGED -p dbc-04ac0685-8857
cat > vol.yaml <<'YAML'
experiment_name: vchen_demo_vol
command: cd $CODE_SOURCE_PATH && python train.py
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
code_source: {type: snapshot, snapshot: {root_path: src, remote_volume: /Volumes/main/default/vchen_demo}}
YAML
dbcli experimental air run -f vol.yaml -p dbc-04ac0685-8857 --debug 2>&1 | grep -iE "submitted|code_source_path|/api/2.0/fs/files"
# → "code_source_path": "/Volumes/main/default/vchen_demo/.internal/src.tar.gz"
# → Submitted run 438683652713410
dbcli fs ls dbfs:/Volumes/main/default/vchen_demo/.internal -p dbc-04ac0685-8857
# → src.tar.gz   (uploaded to the Volume ✓)

4. include_paths subset — only the listed paths are packaged.
Run: 261250771126835

cat > inc.yaml <<'YAML'
experiment_name: vchen_demo_inc
command: cd $CODE_SOURCE_PATH && python train.py
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
code_source: {type: snapshot, snapshot: {root_path: src, include_paths: [pkg]}}
YAML
dbcli experimental air run -f inc.yaml -p dbc-04ac0685-8857
# Uploading src.tar.gz... → Submitted run 261250771126835
dbcli workspace export /Workspace/Users/v.chen@databricks.com/.air/repo_snapshots/.internal/src.tar.gz --file /tmp/inc.tar.gz -p dbc-04ac0685-8857
tar tzf /tmp/inc.tar.gz
# → src/pkg/util.py ONLY   (train.py / .gitignore excluded ✓)

5. No code_source — nothing is uploaded; code_source_path is left empty (nil-guard).
Run: 138286541552104

cat > none.yaml <<'YAML'
experiment_name: vchen_demo_none
command: echo hello
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
YAML
dbcli experimental air run -f none.yaml -p dbc-04ac0685-8857
# → Submitted run 138286541552104   (no "Uploading" line; code_source_path empty)

eng-dev-ecosystem-bot · 2026-07-29T16:47:00Z

Integration test report

Commit: dff7538

Run: 30504594220

	Env	❌FAIL	🟨KNOWN	💚RECOVERED	🙈SKIP	✅pass	🙈skip	Time
❌	aws linux	3	1	3	4	319	1067	5:27
❌	aws windows	3	1	3	4	321	1065	8:51
❌	azure linux	3	1	3	4	319	1066	5:13
❌	azure windows	3	1	3	4	321	1064	9:48
❌	gcp linux	3	1		5	318	1068	5:08
❌	gcp windows	3	1		5	320	1066	6:43

11 interesting tests: 4 SKIP, 3 FAIL, 3 RECOVERED, 1 KNOWN

	Test Name	aws linux	aws windows	azure linux	azure windows	gcp linux	gcp windows
🟨	TestAccept	🟨K	🟨K	🟨K	🟨K	🟨K	🟨K
❌	TestAccept/bundle/deploy/mlops-stacks	❌F	❌F	❌F	❌F	❌F	❌F
❌	TestAccept/bundle/deploy/mlops-stacks/DATABRICKS_BUNDLE_ENGINE=direct	❌F	❌F	❌F	❌F	❌F	❌F
❌	TestAccept/bundle/deploy/mlops-stacks/DATABRICKS_BUNDLE_ENGINE=terraform	❌F	❌F	❌F	❌F	❌F	❌F
🙈	TestAccept/bundle/invariant/no_drift	🙈S	🙈S	🙈S	🙈S	🙈S	🙈S
🙈	TestAccept/bundle/resources/vector_search_endpoints/drift/recreated_same_name	🙈S	🙈S	🙈S	🙈S	🙈S	🙈S
🙈	TestAccept/bundle/resources/vector_search_indexes/recreate/embedding_dimension	🙈S	🙈S	🙈S	🙈S	🙈S	🙈S
🙈	TestAccept/ssh/connection	🙈S	🙈S	🙈S	🙈S	🙈S	🙈S
💚	TestFetchRepositoryInfoAPI_FromRepo	💚R	💚R	💚R	💚R	🙈S	🙈S
💚	TestFetchRepositoryInfoAPI_FromRepo/root	💚R	💚R	💚R	💚R
💚	TestFetchRepositoryInfoAPI_FromRepo/subdir	💚R	💚R	💚R	💚R

ben-hansen-db

Few comments for me to understand

ben-hansen-db · 2026-07-29T18:40:56Z

+	}
+	cleanup := func() { _ = os.RemoveAll(tmp) }
+
+	tarball := filepath.Join(tmp, filepath.Base(repoPath)+".tar.gz")


in cli we use timestamp to make unique, this should be added here as well to prevent conflicts

plain-tar is now timestamped (<dir>_<UTC>.tar.gz), git-archive stays content-addressed.

ben-hansen-db · 2026-07-29T18:41:29Z

-// 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) {


why are these not added in the new area? this would affect cli as well, right?

Intentionally dropped: the typed AiRuntimeTask has no git_state_path/git_diff_path field and the backend never consumed them, so the DABs upload model carries only the tarball. Can re-add if the proto grows the fields.

we want this for CLI though, I care less about this for DABs

Can you document this as a TODO if we don't add it? It's superhelpful to know the git state and diff

let me try to add it here, if not ill add a todo

ben-hansen-db · 2026-07-29T18:42:51Z

-// 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 {


does the cli still have ability to use the cache?

Restored — git-archive checks the remote path and skips re-upload on a cache hit.

vinchenzo-db · 2026-07-29T20:39:00Z

ben-hansen-db

LGTM, assuming resolving new comments

Package the code_source snapshot and upload it through DABs' artifact-upload plumbing (libraries.ReplaceWithRemotePath + libraries.Upload over a minimal in-memory bundle), rewriting ai_runtime_task.code_source_path to the uploaded remote path. The packaging + upload orchestration is CLI-owned (experimental/air/cmd, OWNERS = us); it only reuses DABs' uploader so we don't reimplement workspace/volume upload. snapshot_dabs.go: - plain-tar the working tree, or `git archive` the resolved commit when code_source.snapshot.git pins a commit/branch (git resolution stays CLI-side — a train.yaml construct with no DABs config surface). - remote_volume uploads to a UC Volume natively: GetFilerForLibraries routes /Volumes destinations to filerForVolume, so no special-casing is needed. Replaces the retired raw-filer upload path. Tests cover working-tree, git-pinned, and remote_volume submits via testserver. Co-authored-by: Isaac

…ache Addresses review feedback on the DABs-upload snapshot path: - Unique tarball names: plain-tar (working-tree) snapshots are now timestamped (<dir>_<UTC>.tar.gz) instead of a fixed <dir>.tar.gz, so two concurrent submissions of the same root_path no longer clobber each other in the shared .internal artifact dir. - Restore git-archive caching: git_archive snapshots are content-addressed by (commit, include_paths); if the identical tarball is already uploaded we skip packaging and upload and reuse the remote path, checked via the same filer the uploader routes to (GetFilerForLibraries). Provenance sidecars (git_state.json / git_diff.patch) remain intentionally dropped: the typed jobs.AiRuntimeTask has no field for them and the backend never consumed them, so the DABs upload model carries only the tarball. Tests: assert plain-tar names are unique/timestamped, and that a repeated git.commit submit is a cache hit (single upload, identical remote path). Acceptance golden regenerated for the content-addressed name. Co-authored-by: Isaac

Per review, restore the git_state.json / git_diff.patch provenance sidecars for the CLI code_source path. They are files uploaded next to the run launch dir (command.sh), independent of the ai_runtime_task proto, so no proto field is needed — the earlier "no proto field" rationale for dropping them was wrong. - snapshotViaDABsUpload now writes the sidecars via the run funcDir filer after the tarball upload, git repos only, best-effort: buildGitStateSidecar + captureDirtyDiff already existed; only the uploadSidecars orchestration was gone. git_state.json records packaging_mode / tip_commit / branch / dirty / diff_status; git_diff.patch carries the working-tree diff when dirty. - snapshotResult regains GitStatePath / GitDiffPath; the gitStateName / gitDiffName constants return to snapshot_git.go. Test: a dirty git submit uploads both sidecars to the launch dir and git_state.json records dirty=true / diff_status=captured. Verified live on staging (git_state.json + git_diff.patch land next to command.sh with correct content). Co-authored-by: Isaac

vinchenzo-db force-pushed the air-cli-tar-snapshot branch 3 times, most recently from 21edbbe to cc10318 Compare July 29, 2026 16:27

databricks deleted a comment from eng-dev-ecosystem-bot Jul 29, 2026

vinchenzo-db requested review from ben-hansen-db and riddhibhagwat-db July 29, 2026 16:34

ben-hansen-db reviewed Jul 29, 2026

View reviewed changes

vinchenzo-db force-pushed the air-cli-tar-snapshot branch from a22f86f to 24e8d0c Compare July 29, 2026 20:31

ben-hansen-db self-requested a review July 29, 2026 21:33

ben-hansen-db approved these changes Jul 29, 2026

View reviewed changes

vinchenzo-db force-pushed the air-cli-tar-snapshot branch from 04bed20 to e07a6fc Compare July 29, 2026 22:16

vinchenzo-db enabled auto-merge July 29, 2026 22:31

vinchenzo-db disabled auto-merge July 29, 2026 22:44

vinchenzo-db force-pushed the air-cli-tar-snapshot branch from e07a6fc to 8df15ca Compare July 30, 2026 00:22

vinchenzo-db changed the base branch from main to air-cli July 30, 2026 00:23

vinchenzo-db force-pushed the air-cli-tar-snapshot branch from 8df15ca to 4bf7cc4 Compare July 30, 2026 00:45

vinchenzo-db added 3 commits July 30, 2026 01:02

vinchenzo-db force-pushed the air-cli-tar-snapshot branch from 4bf7cc4 to dff7538 Compare July 30, 2026 01:03

vinchenzo-db merged commit 2943f5c into air-cli Jul 30, 2026
21 of 22 checks passed

vinchenzo-db deleted the air-cli-tar-snapshot branch July 30, 2026 01:34

Uh oh!

Conversation

vinchenzo-db commented Jul 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Changes

Why

Testing

Unit + acceptance

Live E2E — staging dbc-04ac0685-8857 (GPU_1xA10)

Uh oh!

eng-dev-ecosystem-bot commented Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Integration test report

Uh oh!

ben-hansen-db left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vinchenzo-db Jul 29, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vinchenzo-db commented Jul 29, 2026

Uh oh!

ben-hansen-db left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

vinchenzo-db commented Jul 21, 2026 •

edited

Loading

Live E2E — staging `dbc-04ac0685-8857` (GPU_1xA10)

eng-dev-ecosystem-bot commented Jul 29, 2026 •

edited

Loading

vinchenzo-db Jul 29, 2026 •

edited

Loading