Package AI Runtime code_source via a tgz artifact, not a sync overlay - #6494
Merged
Merged
Conversation
A local-directory code_source_path is now turned into a `tgz` artifact and built and uploaded through the standard artifact path, instead of the aicode mutator splicing a content-addressed tarball onto the sync root. The tgz artifact is the single packing mechanism. - aicode.PackageCodeSource now synthesizes a `tgz` artifact per local-dir code_source_path (path = the dir's parent, include = its basename, so archive entries nest under the basename for the /databricks/code_source/<dir> layout) and rewrites code_source_path to the built tarball. It runs before artifacts.Prepare (initialize) instead of in the build phase. - Remove the sync-root overlay, the content-addressed packer (buildCodeSnapshot), bundle.HasAiRuntimeCodeSnapshot, bundle.AiCodeSnapshotDir, and the validateSnapshotDir guards that only existed for the overlay. - Behavior change: only .gitignore filters the packaged files now; the bundle-wide sync.include/exclude no longer apply to a code artifact. - Keep the git_source / immutable-folder / source-linked / for_each rejections. Co-authored-by: Isaac <no-reply@databricks.com>
ben-hansen-db
approved these changes
Sep 2, 2026
ben-hansen-db
left a comment
ben-hansen-db
left a comment
Contributor
There was a problem hiding this comment.
I think PR description could be a bit more clear. We are basically reverting the other PR right? That's the the majority of code changes are?
| }), nil | ||
| // artifactKey is a stable, unique artifact name for a code directory (relative to the | ||
| // bundle). Two tasks pointing at the same directory collapse to one artifact. | ||
| func artifactKey(relDir string) string { |
Contributor
There was a problem hiding this comment.
from claude:
Issue: artifactKey maps every non-alphanumeric to _, so ./a/b and ./a_b both produce air_code_source_a_b. artifacts[key] (:85) + maps.Copy (:118) collapse them to one tarball; both tasks' code_source_path point at it, so one task silently ships the other's code.
Fix: Add a hash disambiguator to the key, or error on a collision across distinct relDirs
Contributor
Author
Yeah that's the idea + using new code uploader + keeping some of the old guardrails |
artifactKey sanitizes non-alphanumerics to '_', so distinct code_source
directories ("a/b" and "a_b") could collide on one key and collapse into a
single tarball — silently shipping one task's code for another. Detect the
collision across distinct directories and error instead.
Co-authored-by: Isaac <no-reply@databricks.com>
pietern
approved these changes
Sep 3, 2026
| // path. Runs before artifacts.Prepare so the synthesized artifact is prepared | ||
| // and built like any other. Remote values and local files are left untouched. | ||
| aicode.PackageCodeSource(), | ||
|
|
Contributor
There was a problem hiding this comment.
Can you move this next to aicode.Validate()?
| // under .databricks (transient, not synced) so the built file is uploaded once via the | ||
| // artifact path and never swept into a sync or into the archive it produces. | ||
| const codeArtifactOutputDir = ".databricks/air_code_source" | ||
|
|
Contributor
There was a problem hiding this comment.
You could use b.CacheDir() instead of a const.
It resolves to the ~same and is guaranteed to not be swept up in sync.
Contributor
Author
There was a problem hiding this comment.
No literal b.CacheDir() here, the equivalent seems to be b.GetLocalStateDir, but it has the env.TempDir override, which can resolve outside the sync root.
| diags = diags.Extend(diag.FromErr(err)) | ||
| return diags | ||
| relDir := strings.TrimPrefix(filepath.ToSlash(cs.value), "./") | ||
| key := artifactKey(relDir) |
Contributor
There was a problem hiding this comment.
It might be nice to use the job name and task key in the filename instead of the src directory.
When combined, they are unique for this bundle.
Contributor
Author
There was a problem hiding this comment.
I kept the src-dir–based key on purpose: it dedups a code dir shared across tasks to a single tarball/upload, which matters for a multitask DAG where the tasks share one (potentially large like research or universe) code source. Job+task naming would re-upload identical code per task.
Per review, place the two aicode mutators adjacent. Still runs before artifacts.Prepare; ApplyArtifactsDynamicVersion only touches whl artifacts, so the intervening reorder does not affect the synthesized tgz. Co-authored-by: Isaac <no-reply@databricks.com>
vinchenzo-db
enabled auto-merge
September 3, 2026 17:37
Collaborator
Integration test reportCommit: 9a051f0
Top 6 slowest tests (at least 2 minutes):
|
Collaborator
Integration test reportCommit: 8d2500f
133 interesting tests: 131 FAIL, 1 KNOWN, 1 SKIP
Top 50 slowest tests (at least 2 minutes):
|
chenyuem-db
pushed a commit
to chenyuem-db/cli
that referenced
this pull request
Sep 9, 2026
…abricks#6532) Basically, remove the aicode mutator entirely. That way `convert-to-dabs` just calls the `artifacts` block. Clean up the code. Reverts the user-facing local-dir shorthand from databricks#6110 / databricks#6494. ## Changes <!-- Brief summary of your changes that is easy to understand --> ## Why We should simplify the UJ to limit ways to upload code. ## Tests Setup ``` # v.chen at ip-10-90-19-147 in ~ (git:) [19:17:23] $ cd ~/.worktrees/cli-rm-aicode && go build -o /tmp/dbcli-rmaicode . # v.chen at ip-10-90-19-147 in ~/.worktrees/cli-rm-aicode (git:vchen/air-rm-aicode-mutator) [19:17:24] $ mkdir -p ~/air-rmaicode-e2e/src && cd ~/air-rmaicode-e2e cat > run.yaml <<'EOF' experiment_name: rmaicode-e2e command: cd "$CODE_SOURCE_PATH" && python train.py compute: accelerator_type: GPU_1xA10 num_accelerators: 1 environment: version: 5 dependencies: - numpy code_source: type: snapshot snapshot: root_path: ./src EOF cat > src/train.py <<'EOF' import os, mlflow print("hello from delivered code_source", flush=True) print("CODE_SOURCE_PATH=" + os.environ.get("CODE_SOURCE_PATH", "<unset>"), flush=True) mlflow.log_metric("ok", 1) print("done", flush=True) EOF ``` Run convert-to-dabs: ``` $ /tmp/dbcli-rmaicode experimental air convert-to-dabs run.yaml --force cat databricks.yml Wrote a Databricks Asset Bundle to .: databricks.yml generated_artifacts/training_config.yaml generated_artifacts/command.sh To deploy and run this workload as a bundle: 1. /tmp/dbcli-rmaicode bundle validate 2. /tmp/dbcli-rmaicode bundle deploy 3. /tmp/dbcli-rmaicode bundle run rmaicode-e2e --no-wait bundle deploy uploads the code source and launch scripts automatically. To see what it deployed and where: /tmp/dbcli-rmaicode bundle summary Unlike `air run` (which submits an ephemeral run), bundle deploy creates a persistent job that is not garbage-collected. When you are done, remove the job and its uploaded files with: /tmp/dbcli-rmaicode bundle destroy bundle: name: rmaicode-e2e sync: paths: - generated_artifacts artifacts: code_source: type: tgz path: . include: - src files: - source: ./dist/code_source.tgz targets: dev: mode: development default: true resources: jobs: rmaicode-e2e: name: rmaicode-e2e tasks: - task_key: rmaicode-e2e environment_key: default max_retries: 3 ai_runtime_task: experiment: rmaicode-e2e deployments: - command_path: ./generated_artifacts/command.sh compute: accelerator_type: GPU_1xA10 accelerator_count: 1 code_source_path: ./dist/code_source.tgz environments: - environment_key: default spec: environment_version: "5" dependencies: - numpy ``` Validate and deploy: ``` $ /tmp/dbcli-rmaicode bundle validate -t dev -p e2-dogfood /tmp/dbcli-rmaicode bundle deploy -t dev -p e2-dogfood Name: rmaicode-e2e Target: dev Workspace: User: v.chen@databricks.com Path: /Workspace/Users/v.chen@databricks.com/.bundle/rmaicode-e2e/dev Validation OK! Building code_source... Uploading dist/code_source.tgz... Uploading bundle files to /Workspace/Users/v.chen@databricks.com/.bundle/rmaicode-e2e/dev/files... Files: 2 uploaded, 0 deleted Resources: 0 created, 0 changed, 0 deleted, 1 unchanged ``` Run: ``` # v.chen at ip-10-90-19-147 in ~/air-rmaicode-e2e (git:) [19:20:10] $ JOBID=$(/tmp/dbcli-rmaicode bundle summary -t dev -p e2-dogfood -o json | grep -oE '"id":[^,]*' | head -1 | grep -oE '[0-9]+') /tmp/dbcli-rmaicode jobs get $JOBID -p e2-dogfood | grep -oE '"code_source_path":[^,]*' "code_source_path": "/Workspace/Users/v.chen@databricks.com/.bundle/rmaicode-e2e/dev/artifacts/.internal/code_source.tgz" # v.chen at ip-10-90-19-147 in ~/air-rmaicode-e2e (git:) [19:20:16] $ /tmp/dbcli-rmaicode bundle run rmaicode-e2e -t dev -p e2-dogfood --no-wait Run URL: https://e2-dogfood.staging.cloud.databricks.com/jobs/1103803560415808/runs/174957625358063?o=6051921418418893 ``` <img width="1517" height="698" alt="image" src="https://github.com/user-attachments/assets/bdb5ec8d-5445-4b55-9e75-a0b4db7b45b1" /> <img width="1524" height="754" alt="image" src="https://github.com/user-attachments/assets/de66dc59-ee2d-44f2-9a97-65ed1c7e3d99" /> Co-authored-by: Isaac <no-reply@databricks.com>
Contributor
|
@vinchenzo-db are you going to delete other things added in #6110, such as libs/vfs/overlay.go ? |
pranshupand-db
pushed a commit
to pranshupand-db/cli
that referenced
this pull request
Sep 11, 2026
`vfs.Overlay` returned a `Path` that served in-memory files on top of a real tree. It was added in databricks#6110 (Aug 2026) to inject generated files into the bundle sync root for AI Runtime `code_source_path`. Its only caller was removed in databricks#6494 (Sep 2026), which packages the `code_source` as a tgz artifact instead of a sync overlay. Since then `Overlay` (and its whole `overlayPath`/`memFile` machinery) has been reachable only from its own unit tests. `deadcode -test ./...` doesn't catch this because the test suite keeps it "live"; it surfaces only when reachability is computed without test roots. This pull request and its description were written by Isaac. Co-authored-by: Isaac <no-reply@databricks.com>
janniklasrose
pushed a commit
that referenced
this pull request
Sep 15, 2026
…#6494) #6110 <- basically reverting this pr A local-directory code_source_path is now turned into a `tgz` artifact and built and uploaded through the standard artifact path, instead of the aicode mutator splicing a content-addressed tarball onto the sync root. The tgz artifact is the single packing mechanism. - aicode.PackageCodeSource now synthesizes a `tgz` artifact per local-dir code_source_path (path = the dir's parent, include = its basename, so archive entries nest under the basename for the /databricks/code_source/<dir> layout) and rewrites code_source_path to the built tarball. It runs before artifacts.Prepare (initialize) instead of in the build phase. - Remove the sync-root overlay, the content-addressed packer (buildCodeSnapshot), bundle.HasAiRuntimeCodeSnapshot, bundle.AiCodeSnapshotDir, and the validateSnapshotDir guards that only existed for the overlay. - Behavior change: only .gitignore filters the packaged files now; the bundle-wide sync.include/exclude no longer apply to a code artifact. - Keep the git_source / immutable-folder / source-linked / for_each rejections. ## Why #6428 We recently merged a DABs native uploader, which supercedes this mutator workaround ## Tests Unit tests E2E test: Setup: ``` # v.chen at ip-10-90-20-219 in /tmp/pr2-e2e (git:) [18:59:50] $ cd ~/.worktrees/cli-rmaicode && go build -o /tmp/cli-pr3 . && cd - rm -rf /tmp/pr3-e2e && mkdir -p /tmp/pr3-e2e/src && cd /tmp/pr3-e2e printf 'print("train ok")\n' > src/train.py # hand-authored command.sh cds into the extracted code dir itself printf 'cd /databricks/code_source/src\npython train.py\n' > src/command.sh cat > databricks.yml <<'YAML' bundle: name: pr3-code-source-demo resources: jobs: train: name: "[${bundle.target}] pr3 demo" tasks: - task_key: train environment_key: default ai_runtime_task: experiment: pr3_demo code_source_path: ./src deployments: - command_path: src/command.sh compute: {accelerator_type: GPU_1xA10, accelerator_count: 1} environments: - environment_key: default spec: environment_version: "5" targets: dev: {mode: development, default: true} YAML export DATABRICKS_CONFIG_PROFILE=e2-dogfood DATA /tmp/pr2-e2e # v.chen at ip-10-90-20-219 in /tmp/pr3-e2e (git:) [18:59:59] $ /tmp/cli-pr3 bundle deploy Building air_code_source_src... Uploading .databricks/air_code_source/air_code_source_src.tar.gz... Uploading bundle files to /Workspace/Users/v.chen@databricks.com/.bundle/pr3-code-source-demo/dev/files... Created jobs.train Files: 3 uploaded, 0 deleted Resources: 1 created, 0 changed, 0 deleted, 0 unchanged ``` Show deploy works: ``` # v.chen at ip-10-90-20-219 in /tmp/pr3-e2e (git:) [19:00:14] $ /tmp/cli-pr3 bundle deploy Building air_code_source_src... Uploading .databricks/air_code_source/air_code_source_src.tar.gz... Uploading bundle files to /Workspace/Users/v.chen@databricks.com/.bundle/pr3-code-source-demo/dev/files... Files: 0 uploaded, 0 deleted Resources: 0 created, 0 changed, 0 deleted, 1 unchanged # v.chen at ip-10-90-20-219 in /tmp/pr3-e2e (git:) [19:01:00] ``` Investigate what is in tar: ``` $ tar tzf .databricks/air_code_source/air_code_source_src.tar.gz src/command.sh src/train.py # v.chen at ip-10-90-20-219 in /tmp/pr3-e2e (git:) [19:01:29] $ ``` bundle summary: ``` $ /tmp/cli-pr3 bundle summary Name: pr3-code-source-demo Target: dev Workspace: User: v.chen@databricks.com Path: /Workspace/Users/v.chen@databricks.com/.bundle/pr3-code-source-demo/dev Resources: Jobs: train: Name: [dev v_chen] [dev] pr3 demo URL: https://e2-dogfood.staging.cloud.databricks.com/jobs/1015159229445184?w=6051921418418893 # v.chen at ip-10-90-20-219 in /tmp/pr3-e2e (git:) [19:01:51] $ ``` bundle run: ``` $ /tmp/cli-pr3 bundle run train --no-wait Run URL: https://e2-dogfood.staging.cloud.databricks.com/jobs/1015159229445184/runs/174551477706538?o=6051921418418893 ``` Run details: <img width="1512" height="821" alt="image" src="https://github.com/user-attachments/assets/b18e41ce-f861-42be-b247-b5af50b6ce43" /> <img width="1495" height="672" alt="image" src="https://github.com/user-attachments/assets/193f8e8e-97c0-4e0b-a12f-e5671c19f102" /> --------- Co-authored-by: Isaac <no-reply@databricks.com>
janniklasrose
pushed a commit
that referenced
this pull request
Sep 15, 2026
Basically, remove the aicode mutator entirely. That way `convert-to-dabs` just calls the `artifacts` block. Clean up the code. Reverts the user-facing local-dir shorthand from #6110 / #6494. ## Changes <!-- Brief summary of your changes that is easy to understand --> ## Why We should simplify the UJ to limit ways to upload code. ## Tests Setup ``` # v.chen at ip-10-90-19-147 in ~ (git:) [19:17:23] $ cd ~/.worktrees/cli-rm-aicode && go build -o /tmp/dbcli-rmaicode . # v.chen at ip-10-90-19-147 in ~/.worktrees/cli-rm-aicode (git:vchen/air-rm-aicode-mutator) [19:17:24] $ mkdir -p ~/air-rmaicode-e2e/src && cd ~/air-rmaicode-e2e cat > run.yaml <<'EOF' experiment_name: rmaicode-e2e command: cd "$CODE_SOURCE_PATH" && python train.py compute: accelerator_type: GPU_1xA10 num_accelerators: 1 environment: version: 5 dependencies: - numpy code_source: type: snapshot snapshot: root_path: ./src EOF cat > src/train.py <<'EOF' import os, mlflow print("hello from delivered code_source", flush=True) print("CODE_SOURCE_PATH=" + os.environ.get("CODE_SOURCE_PATH", "<unset>"), flush=True) mlflow.log_metric("ok", 1) print("done", flush=True) EOF ``` Run convert-to-dabs: ``` $ /tmp/dbcli-rmaicode experimental air convert-to-dabs run.yaml --force cat databricks.yml Wrote a Databricks Asset Bundle to .: databricks.yml generated_artifacts/training_config.yaml generated_artifacts/command.sh To deploy and run this workload as a bundle: 1. /tmp/dbcli-rmaicode bundle validate 2. /tmp/dbcli-rmaicode bundle deploy 3. /tmp/dbcli-rmaicode bundle run rmaicode-e2e --no-wait bundle deploy uploads the code source and launch scripts automatically. To see what it deployed and where: /tmp/dbcli-rmaicode bundle summary Unlike `air run` (which submits an ephemeral run), bundle deploy creates a persistent job that is not garbage-collected. When you are done, remove the job and its uploaded files with: /tmp/dbcli-rmaicode bundle destroy bundle: name: rmaicode-e2e sync: paths: - generated_artifacts artifacts: code_source: type: tgz path: . include: - src files: - source: ./dist/code_source.tgz targets: dev: mode: development default: true resources: jobs: rmaicode-e2e: name: rmaicode-e2e tasks: - task_key: rmaicode-e2e environment_key: default max_retries: 3 ai_runtime_task: experiment: rmaicode-e2e deployments: - command_path: ./generated_artifacts/command.sh compute: accelerator_type: GPU_1xA10 accelerator_count: 1 code_source_path: ./dist/code_source.tgz environments: - environment_key: default spec: environment_version: "5" dependencies: - numpy ``` Validate and deploy: ``` $ /tmp/dbcli-rmaicode bundle validate -t dev -p e2-dogfood /tmp/dbcli-rmaicode bundle deploy -t dev -p e2-dogfood Name: rmaicode-e2e Target: dev Workspace: User: v.chen@databricks.com Path: /Workspace/Users/v.chen@databricks.com/.bundle/rmaicode-e2e/dev Validation OK! Building code_source... Uploading dist/code_source.tgz... Uploading bundle files to /Workspace/Users/v.chen@databricks.com/.bundle/rmaicode-e2e/dev/files... Files: 2 uploaded, 0 deleted Resources: 0 created, 0 changed, 0 deleted, 1 unchanged ``` Run: ``` # v.chen at ip-10-90-19-147 in ~/air-rmaicode-e2e (git:) [19:20:10] $ JOBID=$(/tmp/dbcli-rmaicode bundle summary -t dev -p e2-dogfood -o json | grep -oE '"id":[^,]*' | head -1 | grep -oE '[0-9]+') /tmp/dbcli-rmaicode jobs get $JOBID -p e2-dogfood | grep -oE '"code_source_path":[^,]*' "code_source_path": "/Workspace/Users/v.chen@databricks.com/.bundle/rmaicode-e2e/dev/artifacts/.internal/code_source.tgz" # v.chen at ip-10-90-19-147 in ~/air-rmaicode-e2e (git:) [19:20:16] $ /tmp/dbcli-rmaicode bundle run rmaicode-e2e -t dev -p e2-dogfood --no-wait Run URL: https://e2-dogfood.staging.cloud.databricks.com/jobs/1103803560415808/runs/174957625358063?o=6051921418418893 ``` <img width="1517" height="698" alt="image" src="https://github.com/user-attachments/assets/bdb5ec8d-5445-4b55-9e75-a0b4db7b45b1" /> <img width="1524" height="754" alt="image" src="https://github.com/user-attachments/assets/de66dc59-ee2d-44f2-9a97-65ed1c7e3d99" /> Co-authored-by: Isaac <no-reply@databricks.com>
janniklasrose
pushed a commit
that referenced
this pull request
Sep 15, 2026
`vfs.Overlay` returned a `Path` that served in-memory files on top of a real tree. It was added in #6110 (Aug 2026) to inject generated files into the bundle sync root for AI Runtime `code_source_path`. Its only caller was removed in #6494 (Sep 2026), which packages the `code_source` as a tgz artifact instead of a sync overlay. Since then `Overlay` (and its whole `overlayPath`/`memFile` machinery) has been reachable only from its own unit tests. `deadcode -test ./...` doesn't catch this because the test suite keeps it "live"; it surfaces only when reachability is computed without test roots. This pull request and its description were written by Isaac. Co-authored-by: Isaac <no-reply@databricks.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
#6110 <- basically reverting this pr
A local-directory code_source_path is now turned into a
tgzartifact and built and uploaded through the standard artifact path, instead of the aicode mutator splicing a content-addressed tarball onto the sync root. The tgz artifact is the single packing mechanism.tgzartifact per local-dir code_source_path (path = the dir's parent, include = its basename, so archive entries nest under the basename for the /databricks/code_source/ layout) and rewrites code_source_path to the built tarball. It runs before artifacts.Prepare (initialize) instead of in the build phase.Why
#6428 We recently merged a DABs native uploader, which supercedes this mutator workaround
Tests
Unit tests
E2E test:
Setup:
Show deploy works:
Investigate what is in tar:
bundle summary:
bundle run:
Run details:
