Package AI Runtime code_source via a tgz artifact, not a sync overlay by vinchenzo-db · Pull Request #6494 · databricks/cli

vinchenzo-db · 2026-09-02T17:53:48Z

#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/ 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:

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

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?

ben-hansen-db · 2026-09-02T22:05:07Z

-	}), 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 {


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

vinchenzo-db · 2026-09-02T23:09:46Z

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?

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 · 2026-09-03T13:02:34Z

+		// 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(),
+


Can you move this next to aicode.Validate()?

pietern · 2026-09-03T13:04:11Z

+// 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"
+


You could use b.CacheDir() instead of a const.

It resolves to the ~same and is guaranteed to not be swept up in sync.

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.

pietern · 2026-09-03T13:06:19Z

-			diags = diags.Extend(diag.FromErr(err))
-			return diags
+		relDir := strings.TrimPrefix(filepath.ToSlash(cs.value), "./")
+		key := artifactKey(relDir)


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.

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>

eng-dev-ecosystem-bot · 2026-09-03T18:00:50Z

Integration test report

Commit: 9a051f0

Run: 33785863864

	Env	💚RECOVERED	🙈SKIP	✅pass	🙈skip	Time
💚	aws linux	1	1	274	1220	5:33
💚	aws windows	1	1	276	1218	4:05
💚	azure linux	1	1	273	1220	6:42
💚	azure windows	1	1	275	1218	5:48
💚	gcp linux	1	1	274	1220	6:07
💚	gcp windows	1	1	276	1218	5:00

	Test Name	aws linux	aws windows	azure linux	azure windows	gcp linux	gcp windows
💚	TestAccept	💚R	💚R	💚R	💚R	💚R	💚R
🙈	TestAccept/ssh/connection	🙈S	🙈S	🙈S	🙈S	🙈S	🙈S

Top 6 slowest tests (at least 2 minutes):

duration	env	testname
4:37	azure windows	TestAccept
4:23	gcp windows	TestAccept
4:01	aws windows	TestAccept
2:12	aws linux	TestAccept
2:05	gcp linux	TestAccept
2:03	azure linux	TestAccept

vinchenzo-db requested review from ben-hansen-db and pietern September 2, 2026 19:04

ben-hansen-db approved these changes Sep 2, 2026

View reviewed changes

pietern approved these changes Sep 3, 2026

View reviewed changes

vinchenzo-db enabled auto-merge September 3, 2026 17:37

vinchenzo-db added this pull request to the merge queue Sep 3, 2026

Merged via the queue into main with commit 8d2500f Sep 3, 2026
36 checks passed

vinchenzo-db deleted the vchen/air-code-source-via-artifact branch September 3, 2026 18:19

Navigation Menu

Uh oh!

Package AI Runtime code_source via a tgz artifact, not a sync overlay - #6494

Package AI Runtime code_source via a tgz artifact, not a sync overlay#6494
vinchenzo-db merged 3 commits into
mainfrom
vchen/air-code-source-via-artifact

vinchenzo-db commented Sep 2, 2026 •

edited

Loading

Uh oh!

ben-hansen-db left a comment

Uh oh!

ben-hansen-db Sep 2, 2026

Uh oh!

vinchenzo-db commented Sep 2, 2026

Uh oh!

pietern Sep 3, 2026

Uh oh!

pietern Sep 3, 2026

Uh oh!

vinchenzo-db Sep 3, 2026

Uh oh!

pietern Sep 3, 2026

Uh oh!

vinchenzo-db Sep 3, 2026

Uh oh!

eng-dev-ecosystem-bot commented Sep 3, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

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

Uh oh!

Why

Tests

Uh oh!

ben-hansen-db left a comment

Choose a reason for hiding this comment

Uh oh!

ben-hansen-db Sep 2, 2026

Choose a reason for hiding this comment

Uh oh!

vinchenzo-db commented Sep 2, 2026

Uh oh!

pietern Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

pietern Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

vinchenzo-db Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

pietern Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

vinchenzo-db Sep 3, 2026

Choose a reason for hiding this comment

Uh oh!

eng-dev-ecosystem-bot commented Sep 3, 2026

Integration test report

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

vinchenzo-db commented Sep 2, 2026 •

edited

Loading