From 470ac256338edd516b2c4eb2b173a01512eab688 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Fri, 17 Jul 2026 21:37:38 +0000 Subject: [PATCH 01/16] Package a local directory code_source_path for AI Runtime tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Let ai_runtime_task.code_source_path point at a local directory, not just a pre-built tarball. A new aicode mutator (run in the build phase, before libraries.ExpandGlobReferences/ReplaceWithRemotePath) detects a local-directory value, packages it into a reproducible, content-addressed tarball — honoring .gitignore and the top-level sync.include/exclude globs — uploads it to the user's ~/.air/repo_snapshots directory, and rewrites code_source_path to the uploaded (de-/Workspace-prefixed) path. Because it runs first, PR #5922's artifact-style collection then sees an already-remote path and skips it, so the two compose: a directory is packaged by the CLI, a pre-built .tgz still flows through the artifact path. Also synthesizes a requirements.yaml next to the task's command_path from the job's serverless environments[] spec, so the AI Runtime harness sets up the workload environment. command_path translation itself is provided by #5922. Verified end-to-end on a live workspace: a plain bundle deploy of a directory code_source_path yields a job whose run terminates SUCCESS; sync.exclude and .gitignore entries are absent from the uploaded snapshot; unchanged code skips re-upload. Co-authored-by: Isaac --- .../bundles/ai-runtime-code-source-dir.md | 1 + .../local_code_source/databricks.yml | 29 ++ .../local_code_source/out.test.toml | 3 + .../local_code_source/output.txt | 29 ++ .../ai_runtime_task/local_code_source/script | 30 +++ .../local_code_source/src/.gitignore | 1 + .../local_code_source/src/command.sh | 2 + .../local_code_source/src/train.py | 1 + .../local_code_source/test.toml | 12 + .../config/mutator/aicode/package_upload.go | 248 ++++++++++++++++++ .../mutator/aicode/package_upload_test.go | 74 ++++++ bundle/config/mutator/aicode/requirements.go | 144 ++++++++++ .../mutator/aicode/requirements_test.go | 36 +++ .../config/mutator/aicode/snapshot_package.go | 120 +++++++++ .../mutator/aicode/snapshot_package_test.go | 121 +++++++++ bundle/config/mutator/aicode/validate.go | 101 +++++++ bundle/config/mutator/aicode/validate_test.go | 64 +++++ bundle/phases/build.go | 12 + bundle/phases/initialize.go | 7 + 19 files changed, 1035 insertions(+) create mode 100644 .nextchanges/bundles/ai-runtime-code-source-dir.md create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/output.txt create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/script create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/train.py create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/test.toml create mode 100644 bundle/config/mutator/aicode/package_upload.go create mode 100644 bundle/config/mutator/aicode/package_upload_test.go create mode 100644 bundle/config/mutator/aicode/requirements.go create mode 100644 bundle/config/mutator/aicode/requirements_test.go create mode 100644 bundle/config/mutator/aicode/snapshot_package.go create mode 100644 bundle/config/mutator/aicode/snapshot_package_test.go create mode 100644 bundle/config/mutator/aicode/validate.go create mode 100644 bundle/config/mutator/aicode/validate_test.go diff --git a/.nextchanges/bundles/ai-runtime-code-source-dir.md b/.nextchanges/bundles/ai-runtime-code-source-dir.md new file mode 100644 index 00000000000..d3c5d0aac89 --- /dev/null +++ b/.nextchanges/bundles/ai-runtime-code-source-dir.md @@ -0,0 +1 @@ +An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy` (honoring `.gitignore` and `sync.include`/`sync.exclude`, content-addressed so unchanged code skips re-upload), uploaded to the user's workspace code snapshot directory, and rewritten to the uploaded path. A `requirements.yaml` derived from the job's serverless `environments` is written next to the task's `command_path`. Complements the existing pre-built-tarball path by letting users point at a source directory directly. diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml new file mode 100644 index 00000000000..ff4a7b5dc1b --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml @@ -0,0 +1,29 @@ +bundle: + name: ai-runtime-test + +sync: + # Exclude a file from both bundle sync and the code snapshot. + exclude: + - "**/*.log" + +resources: + jobs: + train: + name: "[${bundle.target}] AI Runtime training" + tasks: + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source_path: ./src + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_8xH100 + accelerator_count: 8 + environments: + - environment_key: default + spec: + environment_version: "5" + dependencies: + - torch>=2.0.0 diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt new file mode 100644 index 00000000000..33fb18eba49 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -0,0 +1,29 @@ + +=== deploy packages the local code source + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +=== the snapshot tarball is uploaded to the user's repo_snapshots directory +/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== code_source_path points at the uploaded archive (no /Workspace prefix) +/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== command_path is translated to its absolute synced workspace path +/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh + +=== a requirements.yaml is written next to command_path +/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml + +=== re-deploying unchanged code is a cache hit: the tarball is not re-uploaded + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! +tarball uploads on redeploy: 0 diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script new file mode 100644 index 00000000000..d0a9e2ea965 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -0,0 +1,30 @@ +# A local code_source_path is packaged into a content-addressed tarball +# (_.tar.gz) and uploaded to the user's repo_snapshots directory. The +# command_path is translated to its synced workspace path, and a requirements.yaml +# derived from the job's serverless environment is written beside it. + +title "deploy packages the local code source\n" +trace $CLI bundle deploy + +title "the snapshot tarball is uploaded to the user's repo_snapshots directory\n" +# The upload appears twice: WSFS import-file 404s until the parent dir exists, then +# the filer mkdirs and retries. Collapse to the distinct path. +jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u + +title "code_source_path points at the uploaded archive (no /Workspace prefix)\n" +jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt + +title "command_path is translated to its absolute synced workspace path\n" +jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.deployments[0].command_path' out.requests.txt + +title "a requirements.yaml is written next to command_path\n" +jq -r 'select(.path | test("/import-file/.*/files/src/requirements.yaml$")) | .path' out.requests.txt | sort -u + +rm out.requests.txt + +title "re-deploying unchanged code is a cache hit: the tarball is not re-uploaded\n" +trace $CLI bundle deploy +echo -n "tarball uploads on redeploy: " +jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u | wc -l | tr -d ' ' + +rm out.requests.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore new file mode 100644 index 00000000000..7a79c3b103e --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore @@ -0,0 +1 @@ +ignored_by_git.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh b/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh new file mode 100644 index 00000000000..7ce3949767c --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh @@ -0,0 +1,2 @@ +cd $CODE_SOURCE_PATH +python train.py diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py b/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py new file mode 100644 index 00000000000..d8b062e5dfe --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py @@ -0,0 +1 @@ +print("training") diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml new file mode 100644 index 00000000000..6e4639a532a --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml @@ -0,0 +1,12 @@ +RecordRequests = true + +Ignore = [ + '.databricks', +] + +# The archive is content-addressed: _.tar.gz. The hash is stable +# given the committed inputs, but collapse it to a token so the test does not pin a +# specific digest. +[[Repls]] +Old = 'src_[0-9a-f]{16}\.tar\.gz' +New = 'src_[SNAPSHOT].tar.gz' diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go new file mode 100644 index 00000000000..a4f3c4a7bbe --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload.go @@ -0,0 +1,248 @@ +// Package aicode packages a local code directory referenced by an AI Runtime +// task's code_source_path and uploads it to the workspace during deploy. +// +// The SDK jobs.AiRuntimeTask.code_source_path field expects a workspace or UC +// volume path to an uploaded code archive; its doc comment states that the CLI +// is responsible for packaging the user's local code directory into that +// archive. This mutator implements that contract for DABs: when a user points +// code_source_path at a local directory, it packages the directory into a +// reproducible tarball (.git and gitignored files excluded), uploads the archive +// to the user's workspace code snapshot directory, and rewrites the field to the +// resulting remote path so the deployed job runs against the uploaded code. Values +// that are already remote are left untouched. +// +// The archive is content-addressed: its name embeds the SHA-256 of the +// (reproducible) tarball, so an unchanged code directory resolves to the same +// remote path across deploys and re-uploads are skipped (see snapshot_package.go). +package aicode + +import ( + "bytes" + "context" + "errors" + "fmt" + "io/fs" + "path" + "path/filepath" + "strings" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/deploy/files" + "github.com/databricks/cli/bundle/libraries" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/fileset" + "github.com/databricks/cli/libs/log" + libsync "github.com/databricks/cli/libs/sync" +) + +// codeSourcePatterns are the config locations of an AI Runtime task's +// code_source_path, both as a direct task and nested under a for_each_task. +var codeSourcePatterns = []dyn.Pattern{ + dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"), + ), + dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + dyn.Key("for_each_task"), dyn.Key("task"), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"), + ), +} + +// codeSource is a single local code_source_path occurrence to package. +type codeSource struct { + configPath dyn.Path + location dyn.Location + // value is the raw code_source_path string as written in config. + value string +} + +func PackageAndUpload() bundle.Mutator { + return &packageAndUpload{} +} + +type packageAndUpload struct { + // client is the filer used for uploads. When nil (the normal case) a filer + // rooted at the code snapshot cache is built per code source. It is only set + // in tests, to inject a recording filer. + client filer.Filer +} + +func (m *packageAndUpload) Name() string { + return "aicode.PackageAndUpload" +} + +func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + sources, diags := collectLocalCodeSources(b) + if diags.HasError() { + return diags + } + if len(sources) == 0 { + return diags + } + + userDir, err := userWorkspaceHome(b) + if err != nil { + return diags.Extend(diag.FromErr(err)) + } + + // remotePaths maps each config location to the remote archive path it should + // point to after upload. Built outside the Mutate closure so upload failures + // are reported before any config is rewritten. + remotePaths := make(map[string]string, len(sources)) + for _, cs := range sources { + remote, err := m.packageOne(ctx, b, cs, userDir) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + return diags + } + remotePaths[cs.configPath.String()] = remote + } + + err = b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { + for _, cs := range sources { + remote := remotePaths[cs.configPath.String()] + root, err = dyn.SetByPath(root, cs.configPath, dyn.NewValue(remote, []dyn.Location{cs.location})) + if err != nil { + return root, fmt.Errorf("failed to update code_source_path %q to %q: %w", cs.value, remote, err) + } + } + return root, nil + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + + return diags +} + +// repoSnapshotsSubdir is the per-user workspace location for code snapshots, +// under the user's home. It matches the Python air CLI (and PR #5897) and is +// deliberately NOT /.internal, which artifacts.CleanUp() deletes at +// the start of every deploy. +const repoSnapshotsSubdir = ".air/repo_snapshots" + +// packageOne packages the local directory for a single code source into a +// reproducible tarball, uploads it to the user's repo_snapshots dir (skipping the +// upload when a content-identical archive already exists there), and returns the +// remote path the config should point to. +func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, userDir string) (string, error) { + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) + dirName := filepath.Base(localDir) + + // relBase is the code directory relative to the sync root, used both to scope the + // sync file list to this directory and to re-base archive entry names under it. + relBase, err := filepath.Rel(b.SyncRootPath, localDir) + if err != nil { + return "", fmt.Errorf("code_source_path %q: %w", cs.value, err) + } + relBase = filepath.ToSlash(relBase) + + files, err := codeSourceFiles(ctx, b, relBase) + if err != nil { + return "", fmt.Errorf("failed to list files for code_source_path %q: %w", cs.value, err) + } + + uploadPath := path.Join(userDir, repoSnapshotsSubdir, dirName) + client := m.client + if client == nil { + client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath) + if err != nil { + return "", err + } + } + + // Build the archive in memory so its content hash can name the upload; the hash + // is computed while gzipping, so this adds no extra pass over the files. + var buf bytes.Buffer + sha, err := buildCodeSnapshot(b.SyncRoot, relBase, files, dirName, &buf) + if err != nil { + return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) + } + + archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16]) + // The AI Runtime snapshot fetcher expects code_source_path in the legacy + // "/Users/..." form (no "/Workspace" prefix), matching the Python air CLI. The + // filer needs the "/Workspace/Users/..." form to upload, so upload to uploadPath + // but record the de-prefixed path on the task. + remotePath := strings.TrimPrefix(path.Join(uploadPath, archiveName), "/Workspace") + + // The archive is reproducible, so a matching name means identical content is + // already uploaded: skip the upload and just point the config at it. + if _, err := client.Stat(ctx, archiveName); err == nil { + log.Debugf(ctx, "code snapshot already present at %s, skipping upload", remotePath) + return remotePath, nil + } else if !errors.Is(err, fs.ErrNotExist) { + return "", fmt.Errorf("failed to check for existing code snapshot %q: %w", remotePath, err) + } + + if err := client.Write(ctx, archiveName, &buf, filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return "", fmt.Errorf("failed to upload code snapshot %q: %w", remotePath, err) + } + return remotePath, nil +} + +// codeSourceFiles returns the files under the code directory (relBase, relative to +// the sync root) that should go into the snapshot. It reuses the bundle's sync +// options so the file list is filtered exactly like bundle file sync: .gitignore +// aware, plus the top-level sync.include/exclude globs. Scoping Paths to relBase +// restricts the walk (and the returned relative paths) to the code directory. +func codeSourceFiles(ctx context.Context, b *bundle.Bundle, relBase string) ([]fileset.File, error) { + opts, err := files.GetSyncOptions(ctx, b) + if err != nil { + return nil, err + } + // Scope the file list to the code directory (relBase) while keeping the + // bundle's include/exclude globs, so filtering matches bundle file sync. + fl, err := libsync.NewFileList(ctx, opts.WorktreeRoot, opts.LocalRoot, []string{relBase}, opts.Include, opts.Exclude) + if err != nil { + return nil, err + } + return fl.Files(ctx) +} + +// userWorkspaceHome returns the current user's workspace home directory +// (/Workspace/Users/), the root under which code snapshots are stored. +func userWorkspaceHome(b *bundle.Bundle) (string, error) { + u := b.Config.Workspace.CurrentUser + if u == nil || u.User == nil || u.UserName == "" { + return "", errors.New("unable to resolve code snapshot location: current user not set") + } + return "/Workspace/Users/" + u.UserName, nil +} + +// collectLocalCodeSources returns every AI Runtime task code_source_path that +// points at a local directory. Already-remote values are skipped. +func collectLocalCodeSources(b *bundle.Bundle) ([]codeSource, diag.Diagnostics) { + var sources []codeSource + var diags diag.Diagnostics + + for _, pattern := range codeSourcePatterns { + err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { + return dyn.MapByPattern(root, pattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) { + value, ok := v.AsString() + if !ok { + return v, fmt.Errorf("expected string, got %s", v.Kind()) + } + if !libraries.IsLocalPath(value) { + return v, nil + } + sources = append(sources, codeSource{ + configPath: p, + location: v.Location(), + value: value, + }) + return v, nil + }) + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + } + + return sources, diags +} diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go new file mode 100644 index 00000000000..25192332f81 --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -0,0 +1,74 @@ +package aicode + +import ( + "path/filepath" + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/bundle/internal/bundletest" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// bundleWithCodeSource builds a bundle rooted at dir whose single AI Runtime task +// points at codeSourcePath. +// +// The end-to-end package/upload/rewrite behavior (local dir -> tarball -> upload -> +// rewritten code_source_path) runs the full mutator pipeline (sync file list, +// workspace filer) and is covered by acceptance tests under +// acceptance/bundle/ai_runtime_task. This unit test covers only the pure +// config-collection seam that does not touch the pipeline. +func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bundle { + t.Helper() + b := &bundle.Bundle{ + BundleRootPath: dir, + SyncRootPath: dir, + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Workspace: config.Workspace{ + CurrentUser: &config.User{User: &iam.User{UserName: "me@databricks.com"}}, + }, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "train": { + JobSettings: jobs.JobSettings{ + Tasks: []jobs.Task{ + { + TaskKey: "train", + AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp", CodeSourcePath: codeSourcePath}, + }, + }, + }, + }, + }, + }, + }, + } + bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + return b +} + +func TestCollectLocalCodeSourcesFindsLocalPath(t *testing.T) { + b := bundleWithCodeSource(t, t.TempDir(), "./src") + sources, diags := collectLocalCodeSources(b) + require.Empty(t, diags) + require.Len(t, sources, 1) + assert.Equal(t, "./src", sources[0].value) +} + +func TestCollectLocalCodeSourcesSkipsRemotePaths(t *testing.T) { + for _, remote := range []string{ + "/Workspace/Users/me/code.tar.gz", + "/Volumes/main/default/code/existing.tar.gz", + } { + b := bundleWithCodeSource(t, t.TempDir(), remote) + sources, diags := collectLocalCodeSources(b) + require.Empty(t, diags) + assert.Empty(t, sources, "remote code_source_path %q must not be collected", remote) + } +} diff --git a/bundle/config/mutator/aicode/requirements.go b/bundle/config/mutator/aicode/requirements.go new file mode 100644 index 00000000000..df2eee3cbfa --- /dev/null +++ b/bundle/config/mutator/aicode/requirements.go @@ -0,0 +1,144 @@ +package aicode + +import ( + "bytes" + "context" + "fmt" + "path" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" + "go.yaml.in/yaml/v3" +) + +// requirementsFileName is the file the AI Runtime entry script reads from the +// command_path directory to install the workload's pip dependencies. The server +// derives its path from command_path, so it must sit next to command.sh. +const requirementsFileName = "requirements.yaml" + +// requirementsSpec is the requirements.yaml shape the AI Runtime consumes: a +// client image version and a list of pip requirement lines. It mirrors what the +// Python air CLI synthesizes. +type requirementsSpec struct { + Version string `yaml:"version,omitempty"` + Dependencies []string `yaml:"dependencies,omitempty"` +} + +// SynthesizeRequirements uploads a requirements.yaml next to each AI Runtime task's +// command_path, derived from the job-level serverless environment referenced by the +// task's environment_key. The AI Runtime entry script reads this file (resolved +// from command_path's directory) to set up the workload environment; without it the +// run fails during setup. It runs at deploy time, after command_path has been +// translated to its absolute workspace path. +func SynthesizeRequirements() bundle.Mutator { + return &synthesizeRequirements{} +} + +type synthesizeRequirements struct { + // client is the filer used for uploads, keyed by the command_path directory. + // When nil (normal case) a workspace filer is built per directory; only set in + // tests to inject a recording filer. + client filer.Filer +} + +func (m *synthesizeRequirements) Name() string { + return "aicode.SynthesizeRequirements" +} + +func (m *synthesizeRequirements) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + + for name, job := range b.Config.Resources.Jobs { + envs := environmentsByKey(job.Environments) + for i := range job.Tasks { + task := &job.Tasks[i] + if task.AiRuntimeTask == nil { + continue + } + if err := m.synthesizeForTask(ctx, b, name, task, envs); err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + } + } + + return diags +} + +// synthesizeForTask uploads requirements.yaml next to task's command_path. It is a +// no-op when the task has no deployment command_path or no matching environment +// spec (the run can still supply deps another way, so this is not an error). +func (m *synthesizeRequirements) synthesizeForTask(ctx context.Context, b *bundle.Bundle, jobName string, task *jobs.Task, envs map[string]*compute.Environment) error { + if len(task.AiRuntimeTask.Deployments) == 0 { + return nil + } + commandPath := task.AiRuntimeTask.Deployments[0].CommandPath + if commandPath == "" { + return nil + } + + env := envs[task.EnvironmentKey] + if env == nil { + return nil + } + + content, err := renderRequirements(env) + if err != nil { + return err + } + + // command_path is an absolute workspace path (translated during initialize); the + // entry script reads requirements.yaml from its directory. + dir := path.Dir(commandPath) + client := m.client + if client == nil { + client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), dir) + if err != nil { + return err + } + } + + log.Debugf(ctx, "writing %s for job %s task %s to %s", requirementsFileName, jobName, task.TaskKey, dir) + if err := client.Write(ctx, requirementsFileName, bytes.NewReader(content), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return fmt.Errorf("failed to upload %s next to command_path %q: %w", requirementsFileName, commandPath, err) + } + return nil +} + +// renderRequirements builds the requirements.yaml content from a serverless +// environment spec: version from environment_version (or the legacy client field), +// dependencies from its pip dependency list. +func renderRequirements(env *compute.Environment) ([]byte, error) { + version := env.EnvironmentVersion + if version == "" { + version = env.Client + } + spec := requirementsSpec{ + Version: version, + Dependencies: env.Dependencies, + } + var buf bytes.Buffer + enc := yaml.NewEncoder(&buf) + enc.SetIndent(2) + if err := enc.Encode(spec); err != nil { + return nil, err + } + if err := enc.Close(); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +// environmentsByKey indexes a job's environments by their key for task lookup. +func environmentsByKey(envs []jobs.JobEnvironment) map[string]*compute.Environment { + out := make(map[string]*compute.Environment, len(envs)) + for i := range envs { + if envs[i].Spec != nil { + out[envs[i].EnvironmentKey] = envs[i].Spec + } + } + return out +} diff --git a/bundle/config/mutator/aicode/requirements_test.go b/bundle/config/mutator/aicode/requirements_test.go new file mode 100644 index 00000000000..a8044e3e35c --- /dev/null +++ b/bundle/config/mutator/aicode/requirements_test.go @@ -0,0 +1,36 @@ +package aicode + +import ( + "testing" + + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestRenderRequirementsFromEnvironmentVersion(t *testing.T) { + out, err := renderRequirements(&compute.Environment{ + EnvironmentVersion: "5", + Dependencies: []string{"numpy", "torch>=2.0.0"}, + }) + require.NoError(t, err) + assert.Equal(t, "version: \"5\"\ndependencies:\n - numpy\n - torch>=2.0.0\n", string(out)) +} + +func TestRenderRequirementsFallsBackToClient(t *testing.T) { + out, err := renderRequirements(&compute.Environment{Client: "5"}) + require.NoError(t, err) + assert.Equal(t, "version: \"5\"\n", string(out)) +} + +func TestEnvironmentsByKey(t *testing.T) { + envs := []jobs.JobEnvironment{ + {EnvironmentKey: "default", Spec: &compute.Environment{EnvironmentVersion: "5"}}, + {EnvironmentKey: "nospec"}, + } + got := environmentsByKey(envs) + require.Contains(t, got, "default") + assert.Equal(t, "5", got["default"].EnvironmentVersion) + assert.NotContains(t, got, "nospec", "environments without a spec are skipped") +} diff --git a/bundle/config/mutator/aicode/snapshot_package.go b/bundle/config/mutator/aicode/snapshot_package.go new file mode 100644 index 00000000000..ee545df1d6c --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -0,0 +1,120 @@ +package aicode + +import ( + "archive/tar" + "compress/gzip" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "path" + "slices" + "strings" + "time" + + "github.com/databricks/cli/libs/fileset" + "github.com/databricks/cli/libs/vfs" +) + +// tarEpoch is a fixed modification time stamped on every tar entry so the archive +// is content-addressed: identical file contents always produce identical bytes +// (and therefore an identical SHA-256), regardless of file mtimes or when the +// archive was built. This is what lets an unchanged code directory resolve to the +// same uploaded filename across deploys and skip re-upload. The technique mirrors +// bundle/deploy/snapshot/path.go (which does the same for the immutable-folder zip). +var tarEpoch = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) + +// appleDoublePrefix is the basename prefix of macOS AppleDouble metadata files. +// The AIR CLI excludes these; we match it so archives are identical on macOS and Linux. +const appleDoublePrefix = "._" + +// buildCodeSnapshot writes a reproducible gzipped tarball of the given files to out +// and returns its SHA-256 hex digest. syncRoot is the root the files' Relative paths +// are against (the bundle sync root); relBase is the code directory relative to that +// root; prefix is the archive's top-level directory name. Each file at +// "/" is written to the archive as "/", so the archive +// expands to /... — matching the runtime's /databricks/code_source/ +// extraction contract. +// +// The file list is produced by the bundle's sync walker, so it honors .gitignore +// (including nested files) and the top-level sync.include/exclude globs — the same +// filtering as bundle file sync. +func buildCodeSnapshot(syncRoot vfs.Path, relBase string, files []fileset.File, prefix string, out io.Writer) (string, error) { + // Sort by relative path so the archive byte stream (and thus its hash) does not + // depend on iteration order. + slices.SortFunc(files, func(a, b fileset.File) int { + return strings.Compare(a.Relative, b.Relative) + }) + + hash := sha256.New() + gzw := gzip.NewWriter(io.MultiWriter(out, hash)) + tw := tar.NewWriter(gzw) + + for _, f := range files { + if err := addFileToArchive(tw, syncRoot, relBase, f, prefix); err != nil { + return "", err + } + } + + if err := tw.Close(); err != nil { + return "", err + } + if err := gzw.Close(); err != nil { + return "", err + } + return hex.EncodeToString(hash.Sum(nil)), nil +} + +func addFileToArchive(tw *tar.Writer, syncRoot vfs.Path, relBase string, f fileset.File, prefix string) error { + // f.Relative is relative to syncRoot and slash-separated. Re-base it under the + // code directory so the entry nests under the archive prefix. + rel := f.Relative + if relBase != "." { + trimmed, ok := strings.CutPrefix(rel, relBase+"/") + if !ok { + // Not under the code dir; the sync file list is scoped to it, so this + // should not happen, but skip defensively rather than mis-place a file. + return nil + } + rel = trimmed + } + + if strings.HasPrefix(path.Base(rel), appleDoublePrefix) { + return nil + } + + rc, err := syncRoot.Open(f.Relative) + if err != nil { + return fmt.Errorf("open %s: %w", f.Relative, err) + } + defer rc.Close() + + info, err := rc.Stat() + if err != nil { + return fmt.Errorf("stat %s: %w", f.Relative, err) + } + + // Only regular files are archived. The walker never yields directories, and + // symlinks inside a code snapshot are out of scope. + if !info.Mode().IsRegular() { + return nil + } + + hdr := &tar.Header{ + Typeflag: tar.TypeReg, + Name: path.Join(prefix, rel), + Size: info.Size(), + // Normalize permissions and zero the mtime so the archive is reproducible + // across machines. The runtime invokes code via an interpreter, not by + // executing files from the snapshot, so execute bits are not preserved. + Mode: 0o644, + ModTime: tarEpoch, + } + if err := tw.WriteHeader(hdr); err != nil { + return fmt.Errorf("tar header for %s: %w", rel, err) + } + if _, err := io.Copy(tw, rc); err != nil { + return fmt.Errorf("write %s: %w", rel, err) + } + return nil +} diff --git a/bundle/config/mutator/aicode/snapshot_package_test.go b/bundle/config/mutator/aicode/snapshot_package_test.go new file mode 100644 index 00000000000..feb2e10c1e7 --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package_test.go @@ -0,0 +1,121 @@ +package aicode + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "io" + "os" + "path/filepath" + "testing" + + "github.com/databricks/cli/libs/fileset" + "github.com/databricks/cli/libs/vfs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// writeTree materializes files (relative slash path -> content) under a fresh +// temp dir and returns a vfs.Path rooted at it plus the fileset for its contents. +func writeTree(t *testing.T, files map[string]string) (vfs.Path, []fileset.File) { + t.Helper() + dir := t.TempDir() + for name, content := range files { + p := filepath.Join(dir, filepath.FromSlash(name)) + require.NoError(t, os.MkdirAll(filepath.Dir(p), 0o755)) + require.NoError(t, os.WriteFile(p, []byte(content), 0o644)) + } + root := vfs.MustNew(dir) + fs, err := fileset.New(root).Files() + require.NoError(t, err) + return root, fs +} + +// tarEntries reads a gzipped tarball and returns entry name -> content. +func tarEntries(t *testing.T, b []byte) map[string]string { + t.Helper() + gzr, err := gzip.NewReader(bytes.NewReader(b)) + require.NoError(t, err) + tr := tar.NewReader(gzr) + out := map[string]string{} + for { + hdr, err := tr.Next() + if err == io.EOF { + break + } + require.NoError(t, err) + content, err := io.ReadAll(tr) + require.NoError(t, err) + out[hdr.Name] = string(content) + } + return out +} + +func TestBuildCodeSnapshotPrefixesEntries(t *testing.T) { + root, files := writeTree(t, map[string]string{ + "train.py": "print('train')", + "pkg/util.py": "x = 1", + "._resource_fork": "apple double", + }) + + var buf bytes.Buffer + sha, err := buildCodeSnapshot(root, ".", files, "mycode", &buf) + require.NoError(t, err) + require.NotEmpty(t, sha) + + entries := tarEntries(t, buf.Bytes()) + // Entries are prefixed with the code dir basename (runtime extracts to + // /databricks/code_source/). + assert.Equal(t, "print('train')", entries["mycode/train.py"]) + assert.Equal(t, "x = 1", entries["mycode/pkg/util.py"]) + assert.NotContains(t, entries, "mycode/._resource_fork", "AppleDouble metadata must be excluded") +} + +func TestBuildCodeSnapshotRebasesUnderRelBase(t *testing.T) { + // Files listed relative to a sync root; only the "src" subtree is packaged and + // re-based so entries nest under the prefix (not the intermediate "src/"). + root, files := writeTree(t, map[string]string{ + "src/train.py": "t", + "src/pkg/util.py": "u", + "other/ignore.py": "o", + }) + + var buf bytes.Buffer + _, err := buildCodeSnapshot(root, "src", files, "src", &buf) + require.NoError(t, err) + + entries := tarEntries(t, buf.Bytes()) + assert.Contains(t, entries, "src/train.py") + assert.Contains(t, entries, "src/pkg/util.py") + // A file outside relBase is not under "src/", so it is skipped. + assert.NotContains(t, entries, "src/other/ignore.py") + assert.NotContains(t, entries, "other/ignore.py") +} + +func TestBuildCodeSnapshotIsReproducible(t *testing.T) { + files := map[string]string{"a.py": "aaa", "sub/b.py": "bbb"} + root1, fs1 := writeTree(t, files) + root2, fs2 := writeTree(t, files) + + var buf1, buf2 bytes.Buffer + sha1, err := buildCodeSnapshot(root1, ".", fs1, "code", &buf1) + require.NoError(t, err) + sha2, err := buildCodeSnapshot(root2, ".", fs2, "code", &buf2) + require.NoError(t, err) + + assert.Equal(t, sha1, sha2, "identical content must produce an identical hash") + assert.Equal(t, buf1.Bytes(), buf2.Bytes(), "identical content must produce identical bytes") +} + +func TestBuildCodeSnapshotHashChangesWithContent(t *testing.T) { + root1, fs1 := writeTree(t, map[string]string{"main.py": "v1"}) + root2, fs2 := writeTree(t, map[string]string{"main.py": "v2"}) + + var buf1, buf2 bytes.Buffer + sha1, err := buildCodeSnapshot(root1, ".", fs1, "code", &buf1) + require.NoError(t, err) + sha2, err := buildCodeSnapshot(root2, ".", fs2, "code", &buf2) + require.NoError(t, err) + + assert.NotEqual(t, sha1, sha2) +} diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go new file mode 100644 index 00000000000..e5d7a0cdb8f --- /dev/null +++ b/bundle/config/mutator/aicode/validate.go @@ -0,0 +1,101 @@ +package aicode + +import ( + "context" + "fmt" + "os" + "path/filepath" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/libraries" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/dyn" +) + +// Validate checks AI Runtime tasks that reference a local code_source_path so +// that misconfigurations surface at `bundle validate` time with an actionable +// message, rather than as an obscure failure mid-deploy. It performs no uploads. +func Validate() bundle.ReadOnlyMutator { + return &validate{} +} + +type validate struct{ bundle.RO } + +func (v *validate) Name() string { + return "aicode.Validate" +} + +func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + + jobsPath := dyn.NewPath(dyn.Key("resources"), dyn.Key("jobs")) + + for name, job := range b.Config.Resources.Jobs { + jobPath := jobsPath.Append(dyn.Key(name)) + + for i, task := range job.Tasks { + if task.AiRuntimeTask == nil { + continue + } + codePath := jobPath.Append(dyn.Key("tasks"), dyn.Index(i), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) + diags = diags.Extend(v.validateTask(b, job.GitSource != nil, task.AiRuntimeTask.CodeSourcePath, codePath)) + } + } + + return diags +} + +func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { + // Only local code_source_path values are packaged at deploy; remote values + // are used as-is and need no validation here. + if codeSourcePath == "" || !libraries.IsLocalPath(codeSourcePath) { + return nil + } + + locations := b.Config.GetLocations(codePath.String()) + + // The deploy engine retrieves task files from git when git_source is set, so + // packaging a local directory would be silently ignored. Reject the combination. + if jobHasGitSource { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path cannot be combined with git_source", + Detail: "Remove git_source or set code_source_path to a workspace or volume path", + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + // Immutable-folder deployments upload a single content-addressed snapshot and + // do not support the per-task code packaging this mutator performs. + if b.IsImmutableFolder() { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path is not supported with experimental.immutable_folder", + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath)) + info, err := os.Stat(localDir) + if err != nil { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: fmt.Sprintf("code_source_path %q not found", codeSourcePath), + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + if !info.IsDir() { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: fmt.Sprintf("code_source_path %q must be a directory", codeSourcePath), + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + return nil +} diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go new file mode 100644 index 00000000000..cd3548372e1 --- /dev/null +++ b/bundle/config/mutator/aicode/validate_test.go @@ -0,0 +1,64 @@ +package aicode + +import ( + "path/filepath" + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/bundle/internal/bundletest" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func bundleForValidate(t *testing.T, codeSourcePath string, gitSource *jobs.GitSource) *bundle.Bundle { + t.Helper() + dir := t.TempDir() + b := &bundle.Bundle{ + BundleRootPath: dir, + SyncRootPath: dir, + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "train": { + JobSettings: jobs.JobSettings{ + GitSource: gitSource, + Tasks: []jobs.Task{ + { + TaskKey: "train", + AiRuntimeTask: &jobs.AiRuntimeTask{CodeSourcePath: codeSourcePath}, + }, + }, + }, + }, + }, + }, + }, + } + bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + return b +} + +func TestValidateMissingCodeSourceDir(t *testing.T) { + b := bundleForValidate(t, "does-not-exist", nil) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Equal(t, `code_source_path "does-not-exist" not found`, diags[0].Summary) +} + +func TestValidateGitSourceConflict(t *testing.T) { + b := bundleForValidate(t, "src", &jobs.GitSource{GitUrl: "https://example.invalid/repo"}) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "cannot be combined with git_source") +} + +func TestValidateRemoteCodeSourceIsSkipped(t *testing.T) { + b := bundleForValidate(t, "/Volumes/main/default/code/x.tar.gz", nil) + diags := Validate().Apply(t.Context(), b) + assert.Empty(t, diags) +} diff --git a/bundle/phases/build.go b/bundle/phases/build.go index db376e07e28..a386093bc24 100644 --- a/bundle/phases/build.go +++ b/bundle/phases/build.go @@ -7,6 +7,7 @@ import ( "github.com/databricks/cli/bundle/artifacts" "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/config/mutator" + "github.com/databricks/cli/bundle/config/mutator/aicode" "github.com/databricks/cli/bundle/libraries" "github.com/databricks/cli/bundle/scripts" "github.com/databricks/cli/bundle/trampoline" @@ -27,6 +28,17 @@ func Build(ctx context.Context, b *bundle.Bundle) LibLocationMap { artifacts.Build(), scripts.Execute(config.ScriptPostBuild), + // Package any AI Runtime task code_source_path that points at a local + // directory into a tarball, upload it, and rewrite the field to the remote + // path. Runs before ExpandGlobReferences/ReplaceWithRemotePath below so those + // see an already-remote code_source_path (a directory would otherwise be + // collected as a local "library" and fail to upload as a file). + aicode.PackageAndUpload(), + // Write requirements.yaml next to the (already translated) command_path, + // derived from the job's serverless environment, so the AI Runtime harness + // can set up the workload environment. + aicode.SynthesizeRequirements(), + mutator.ResolveVariableReferencesWithoutResources( "artifacts", ), diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index 31c37064029..6a8fd3da224 100644 --- a/bundle/phases/initialize.go +++ b/bundle/phases/initialize.go @@ -10,6 +10,7 @@ import ( "github.com/databricks/cli/bundle/artifacts" "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/config/mutator" + "github.com/databricks/cli/bundle/config/mutator/aicode" pythonmutator "github.com/databricks/cli/bundle/config/mutator/python" "github.com/databricks/cli/bundle/config/validate" "github.com/databricks/cli/bundle/deploy/metadata" @@ -195,6 +196,12 @@ func Initialize(ctx context.Context, b *bundle.Bundle) { mutator.TranslatePaths(), + // Reads (typed): resources.jobs.*.tasks[*].ai_runtime_task.code_source_path, job git_source + // Validates that AI Runtime tasks referencing a local code_source_path point at an existing + // directory and are not combined with git_source or immutable-folder deployments, so these + // misconfigurations are caught at validate time rather than mid-deploy. + aicode.Validate(), + // Reads (typed): b.Config.Experimental.PythonWheelWrapper, b.Config.Presets.SourceLinkedDeployment (checks Python wheel wrapper and deployment mode settings) // Reads (dynamic): resources.jobs.*.tasks (checks for tasks with local libraries and incompatible DBR versions) // Provides warnings when Python wheel tasks are used with DBR < 13.3 or when wheel wrapper is incompatible with source-linked deployment From 2acec817b56b3745c4d289b60fae0b4c7dfb5b53 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 06:16:16 +0000 Subject: [PATCH 02/16] aicode: only claim a local *directory* code_source_path, skip files MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The aicode mutator packages a local directory code_source_path at deploy. A pre-built tarball delivered via an `artifacts` block (code_source_path pointing at a local *file*) must instead flow through the standard artifact-upload path, exactly as the two-path design intends. Validate rejected such a path outright ("code_source_path not found" — the artifact tarball does not exist yet at initialize time), breaking the existing bundle/artifacts/ai_runtime_code_source acceptance test. Both Validate and collectLocalCodeSources now skip any local code_source_path that does not resolve to an existing directory, leaving it for the artifact uploader. The packaging-specific constraints (git_source, immutable_folder) only apply once the path is confirmed to be a directory this mutator will package. Co-authored-by: Isaac --- .../config/mutator/aicode/package_upload.go | 14 ++++++++ .../mutator/aicode/package_upload_test.go | 19 +++++++++-- bundle/config/mutator/aicode/validate.go | 32 +++++++------------ bundle/config/mutator/aicode/validate_test.go | 26 ++++++++++++--- 4 files changed, 65 insertions(+), 26 deletions(-) diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index a4f3c4a7bbe..a8cbef4f88a 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -22,6 +22,7 @@ import ( "errors" "fmt" "io/fs" + "os" "path" "path/filepath" "strings" @@ -231,6 +232,19 @@ func collectLocalCodeSources(b *bundle.Bundle) ([]codeSource, diag.Diagnostics) if !libraries.IsLocalPath(value) { return v, nil } + // Only package a local *directory*. A local file (e.g. a pre-built + // tarball delivered via an `artifacts` block) is left alone so it flows + // through the standard artifact-upload path as a file. aicode.Validate + // applies the same directory check, so the two stay in agreement. + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(value)) + info, statErr := os.Stat(localDir) + isDir := statErr == nil && info.IsDir() + if !isDir { + // Not an existing directory: leave the value untouched for the + // artifact-upload path (a stat error here is not fatal — a missing + // path is simply not something this mutator packages). + return v, nil + } sources = append(sources, codeSource{ configPath: p, location: v.Location(), diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go index 25192332f81..fca33a0a7bf 100644 --- a/bundle/config/mutator/aicode/package_upload_test.go +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -1,6 +1,7 @@ package aicode import ( + "os" "path/filepath" "testing" @@ -53,8 +54,10 @@ func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bund return b } -func TestCollectLocalCodeSourcesFindsLocalPath(t *testing.T) { - b := bundleWithCodeSource(t, t.TempDir(), "./src") +func TestCollectLocalCodeSourcesFindsLocalDir(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o700)) + b := bundleWithCodeSource(t, dir, "./src") sources, diags := collectLocalCodeSources(b) require.Empty(t, diags) require.Len(t, sources, 1) @@ -72,3 +75,15 @@ func TestCollectLocalCodeSourcesSkipsRemotePaths(t *testing.T) { assert.Empty(t, sources, "remote code_source_path %q must not be collected", remote) } } + +// A local path that resolves to a file (not a directory) — e.g. a pre-built +// tarball delivered via an `artifacts` block — is NOT collected: it flows through +// the standard artifact-upload path as a file rather than being packaged here. +func TestCollectLocalCodeSourcesSkipsLocalFile(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "code.tgz"), []byte("x"), 0o600)) + b := bundleWithCodeSource(t, dir, "code.tgz") + sources, diags := collectLocalCodeSources(b) + require.Empty(t, diags) + assert.Empty(t, sources, "a local tarball file must flow through artifact upload, not aicode packaging") +} diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go index e5d7a0cdb8f..8e742f0e981 100644 --- a/bundle/config/mutator/aicode/validate.go +++ b/bundle/config/mutator/aicode/validate.go @@ -2,7 +2,6 @@ package aicode import ( "context" - "fmt" "os" "path/filepath" @@ -53,6 +52,18 @@ func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSour return nil } + // This mutator packages a local *directory*. A local path that is not an existing + // directory is left alone so it flows through the standard artifact-upload path: + // a pre-built tarball delivered via an `artifacts` block is produced during the + // build phase (so it does not exist yet at validate time) and is uploaded as a + // file, not packaged here. Only when the path is an existing directory do the + // packaging-specific constraints below apply. + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath)) + info, statErr := os.Stat(localDir) + if statErr != nil || !info.IsDir() { + return nil + } + locations := b.Config.GetLocations(codePath.String()) // The deploy engine retrieves task files from git when git_source is set, so @@ -78,24 +89,5 @@ func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSour }} } - localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath)) - info, err := os.Stat(localDir) - if err != nil { - return diag.Diagnostics{{ - Severity: diag.Error, - Summary: fmt.Sprintf("code_source_path %q not found", codeSourcePath), - Locations: locations, - Paths: []dyn.Path{codePath}, - }} - } - if !info.IsDir() { - return diag.Diagnostics{{ - Severity: diag.Error, - Summary: fmt.Sprintf("code_source_path %q must be a directory", codeSourcePath), - Locations: locations, - Paths: []dyn.Path{codePath}, - }} - } - return nil } diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go index cd3548372e1..bddcc8f309a 100644 --- a/bundle/config/mutator/aicode/validate_test.go +++ b/bundle/config/mutator/aicode/validate_test.go @@ -1,6 +1,7 @@ package aicode import ( + "os" "path/filepath" "testing" @@ -43,15 +44,32 @@ func bundleForValidate(t *testing.T, codeSourcePath string, gitSource *jobs.GitS return b } -func TestValidateMissingCodeSourceDir(t *testing.T) { +// mkCodeDir creates a code_source directory (with one file) under the bundle's +// sync root, so the path resolves to an existing directory this mutator packages. +func mkCodeDir(t *testing.T, b *bundle.Bundle, rel string) { + t.Helper() + dir := filepath.Join(b.SyncRootPath, rel) + require.NoError(t, os.MkdirAll(dir, 0o700)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "train.py"), []byte("print()\n"), 0o600)) +} + +// A local path that is not an existing directory is left alone: it flows through +// the standard artifact-upload path (e.g. a pre-built tarball built by an +// `artifacts` block, which does not exist yet at validate time). +func TestValidateNonDirectoryCodeSourceIsSkipped(t *testing.T) { + // Missing path (nothing on disk yet). b := bundleForValidate(t, "does-not-exist", nil) - diags := Validate().Apply(t.Context(), b) - require.Len(t, diags, 1) - assert.Equal(t, `code_source_path "does-not-exist" not found`, diags[0].Summary) + assert.Empty(t, Validate().Apply(t.Context(), b)) + + // Existing local file (a pre-built tarball), not a directory. + b = bundleForValidate(t, "code.tgz", nil) + require.NoError(t, os.WriteFile(filepath.Join(b.SyncRootPath, "code.tgz"), []byte("x"), 0o600)) + assert.Empty(t, Validate().Apply(t.Context(), b)) } func TestValidateGitSourceConflict(t *testing.T) { b := bundleForValidate(t, "src", &jobs.GitSource{GitUrl: "https://example.invalid/repo"}) + mkCodeDir(t, b, "src") diags := Validate().Apply(t.Context(), b) require.Len(t, diags, 1) assert.Contains(t, diags[0].Summary, "cannot be combined with git_source") From ef37d80ac2e81d9c2617422372f866d8b91bdfc3 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 12:18:56 +0000 Subject: [PATCH 03/16] =?UTF-8?q?aicode:=20address=20review=20=E2=80=94=20?= =?UTF-8?q?drop=20test-only=20filer=20field,=20tighten=20helpers?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Remove the never-set `client` filer field from packageAndUpload and synthesizeRequirements (only ever nil); build the workspace filer directly. - Read the current user directly (resolved in the initialize phase) instead of the defensive nil-guarded helper. - Surface a stat error other than not-exist (e.g. unreadable directory) instead of silently treating it as "not a directory to package"; shared isExistingDir helper used by collect and Validate. - Pass GitSource into validateTask (nil-check inside) to keep validation logic contained. - Collect a single direct-task code_source_path pattern, matching Validate and SynthesizeRequirements (for_each_task is unsupported until all three gain it). - SynthesizeRequirements writes requirements.yaml next to every deployment's command_path (deduped by directory), not just the first. Co-authored-by: Isaac --- .../local_code_source/src/data/model.bin | 1 + .../local_code_source/src/data/notes.txt | 1 + .../local_code_source/src/debug.log | 1 + .../local_code_source/src/ignored_by_git.txt | 1 + .../config/mutator/aicode/package_upload.go | 139 ++++++++---------- bundle/config/mutator/aicode/requirements.go | 57 ++++--- bundle/config/mutator/aicode/validate.go | 24 ++- 7 files changed, 112 insertions(+), 112 deletions(-) create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/data/model.bin create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/data/notes.txt create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/debug.log create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/ignored_by_git.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/data/model.bin b/acceptance/bundle/ai_runtime_task/local_code_source/src/data/model.bin new file mode 100644 index 00000000000..05424f2a4c8 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/data/model.bin @@ -0,0 +1 @@ +weights diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/data/notes.txt b/acceptance/bundle/ai_runtime_task/local_code_source/src/data/notes.txt new file mode 100644 index 00000000000..fb188b9ecf0 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/data/notes.txt @@ -0,0 +1 @@ +scratch diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/debug.log b/acceptance/bundle/ai_runtime_task/local_code_source/src/debug.log new file mode 100644 index 00000000000..6bfe6b19e37 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/debug.log @@ -0,0 +1 @@ +log diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/ignored_by_git.txt b/acceptance/bundle/ai_runtime_task/local_code_source/src/ignored_by_git.txt new file mode 100644 index 00000000000..bd930095363 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/ignored_by_git.txt @@ -0,0 +1 @@ +kept diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index a8cbef4f88a..97959996b9c 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -38,21 +38,16 @@ import ( libsync "github.com/databricks/cli/libs/sync" ) -// codeSourcePatterns are the config locations of an AI Runtime task's -// code_source_path, both as a direct task and nested under a for_each_task. -var codeSourcePatterns = []dyn.Pattern{ - dyn.NewPattern( - dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), - dyn.Key("tasks"), dyn.AnyIndex(), - dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"), - ), - dyn.NewPattern( - dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), - dyn.Key("tasks"), dyn.AnyIndex(), - dyn.Key("for_each_task"), dyn.Key("task"), - dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"), - ), -} +// codeSourcePattern is the config location of an AI Runtime task's +// code_source_path. It matches a direct task only — the same scope aicode.Validate +// and aicode.SynthesizeRequirements operate on. ai_runtime_task nested under a +// for_each_task is not a supported combination yet; when it is, all three mutators +// should gain it together. +var codeSourcePattern = dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"), +) // codeSource is a single local code_source_path occurrence to package. type codeSource struct { @@ -66,12 +61,7 @@ func PackageAndUpload() bundle.Mutator { return &packageAndUpload{} } -type packageAndUpload struct { - // client is the filer used for uploads. When nil (the normal case) a filer - // rooted at the code snapshot cache is built per code source. It is only set - // in tests, to inject a recording filer. - client filer.Filer -} +type packageAndUpload struct{} func (m *packageAndUpload) Name() string { return "aicode.PackageAndUpload" @@ -86,17 +76,16 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia return diags } - userDir, err := userWorkspaceHome(b) - if err != nil { - return diags.Extend(diag.FromErr(err)) - } + // The current user is resolved during the initialize phase, which runs before + // this build-phase mutator, so it is always set here. + userDir := "/Workspace/Users/" + b.Config.Workspace.CurrentUser.UserName // remotePaths maps each config location to the remote archive path it should // point to after upload. Built outside the Mutate closure so upload failures // are reported before any config is rewritten. remotePaths := make(map[string]string, len(sources)) for _, cs := range sources { - remote, err := m.packageOne(ctx, b, cs, userDir) + remote, err := packageOne(ctx, b, cs, userDir) if err != nil { diags = diags.Extend(diag.FromErr(err)) return diags @@ -104,9 +93,10 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia remotePaths[cs.configPath.String()] = remote } - err = b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { + err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { for _, cs := range sources { remote := remotePaths[cs.configPath.String()] + var err error root, err = dyn.SetByPath(root, cs.configPath, dyn.NewValue(remote, []dyn.Location{cs.location})) if err != nil { return root, fmt.Errorf("failed to update code_source_path %q to %q: %w", cs.value, remote, err) @@ -131,7 +121,7 @@ const repoSnapshotsSubdir = ".air/repo_snapshots" // reproducible tarball, uploads it to the user's repo_snapshots dir (skipping the // upload when a content-identical archive already exists there), and returns the // remote path the config should point to. -func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, userDir string) (string, error) { +func packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, userDir string) (string, error) { localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) dirName := filepath.Base(localDir) @@ -149,12 +139,9 @@ func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs } uploadPath := path.Join(userDir, repoSnapshotsSubdir, dirName) - client := m.client - if client == nil { - client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath) - if err != nil { - return "", err - } + client, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath) + if err != nil { + return "", err } // Build the archive in memory so its content hash can name the upload; the hash @@ -206,57 +193,59 @@ func codeSourceFiles(ctx context.Context, b *bundle.Bundle, relBase string) ([]f return fl.Files(ctx) } -// userWorkspaceHome returns the current user's workspace home directory -// (/Workspace/Users/), the root under which code snapshots are stored. -func userWorkspaceHome(b *bundle.Bundle) (string, error) { - u := b.Config.Workspace.CurrentUser - if u == nil || u.User == nil || u.UserName == "" { - return "", errors.New("unable to resolve code snapshot location: current user not set") - } - return "/Workspace/Users/" + u.UserName, nil -} - // collectLocalCodeSources returns every AI Runtime task code_source_path that // points at a local directory. Already-remote values are skipped. func collectLocalCodeSources(b *bundle.Bundle) ([]codeSource, diag.Diagnostics) { var sources []codeSource var diags diag.Diagnostics - for _, pattern := range codeSourcePatterns { - err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { - return dyn.MapByPattern(root, pattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) { - value, ok := v.AsString() - if !ok { - return v, fmt.Errorf("expected string, got %s", v.Kind()) - } - if !libraries.IsLocalPath(value) { - return v, nil - } - // Only package a local *directory*. A local file (e.g. a pre-built - // tarball delivered via an `artifacts` block) is left alone so it flows - // through the standard artifact-upload path as a file. aicode.Validate - // applies the same directory check, so the two stay in agreement. - localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(value)) - info, statErr := os.Stat(localDir) - isDir := statErr == nil && info.IsDir() - if !isDir { - // Not an existing directory: leave the value untouched for the - // artifact-upload path (a stat error here is not fatal — a missing - // path is simply not something this mutator packages). - return v, nil - } - sources = append(sources, codeSource{ - configPath: p, - location: v.Location(), - value: value, - }) + err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { + return dyn.MapByPattern(root, codeSourcePattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) { + value, ok := v.AsString() + if !ok { + return v, fmt.Errorf("expected string, got %s", v.Kind()) + } + if !libraries.IsLocalPath(value) { + return v, nil + } + // Only package a local *directory*. A local file (e.g. a pre-built + // tarball delivered via an `artifacts` block) is left alone so it flows + // through the standard artifact-upload path as a file. aicode.Validate + // applies the same directory check, so the two stay in agreement. + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(value)) + isDir, err := isExistingDir(localDir) + if err != nil { + return v, fmt.Errorf("code_source_path %q: %w", value, err) + } + if !isDir { return v, nil + } + sources = append(sources, codeSource{ + configPath: p, + location: v.Location(), + value: value, }) + return v, nil }) - if err != nil { - diags = diags.Extend(diag.FromErr(err)) - } + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) } return sources, diags } + +// isExistingDir reports whether path is an existing directory. A not-exist error +// is not an error here (the path is simply not a directory this mutator packages), +// but any other stat failure — notably a permission error on the parent — is +// surfaced so it is not silently swallowed into "skip". +func isExistingDir(path string) (bool, error) { + info, err := os.Stat(path) + if err != nil { + if errors.Is(err, fs.ErrNotExist) { + return false, nil + } + return false, err + } + return info.IsDir(), nil +} diff --git a/bundle/config/mutator/aicode/requirements.go b/bundle/config/mutator/aicode/requirements.go index df2eee3cbfa..8927500d78d 100644 --- a/bundle/config/mutator/aicode/requirements.go +++ b/bundle/config/mutator/aicode/requirements.go @@ -38,12 +38,7 @@ func SynthesizeRequirements() bundle.Mutator { return &synthesizeRequirements{} } -type synthesizeRequirements struct { - // client is the filer used for uploads, keyed by the command_path directory. - // When nil (normal case) a workspace filer is built per directory; only set in - // tests to inject a recording filer. - client filer.Filer -} +type synthesizeRequirements struct{} func (m *synthesizeRequirements) Name() string { return "aicode.SynthesizeRequirements" @@ -59,7 +54,7 @@ func (m *synthesizeRequirements) Apply(ctx context.Context, b *bundle.Bundle) di if task.AiRuntimeTask == nil { continue } - if err := m.synthesizeForTask(ctx, b, name, task, envs); err != nil { + if err := synthesizeForTask(ctx, b, name, task, envs); err != nil { diags = diags.Extend(diag.FromErr(err)) } } @@ -68,18 +63,13 @@ func (m *synthesizeRequirements) Apply(ctx context.Context, b *bundle.Bundle) di return diags } -// synthesizeForTask uploads requirements.yaml next to task's command_path. It is a -// no-op when the task has no deployment command_path or no matching environment -// spec (the run can still supply deps another way, so this is not an error). -func (m *synthesizeRequirements) synthesizeForTask(ctx context.Context, b *bundle.Bundle, jobName string, task *jobs.Task, envs map[string]*compute.Environment) error { - if len(task.AiRuntimeTask.Deployments) == 0 { - return nil - } - commandPath := task.AiRuntimeTask.Deployments[0].CommandPath - if commandPath == "" { - return nil - } - +// synthesizeForTask uploads requirements.yaml next to each deployment's +// command_path. It is a no-op when the task has no matching environment spec (the +// run can still supply deps another way, so this is not an error). Each deployment +// carries its own command_path and the entry script reads requirements.yaml from +// that directory, so every deployment's directory needs the file; directories are +// deduplicated so a shared command_path directory is written only once. +func synthesizeForTask(ctx context.Context, b *bundle.Bundle, jobName string, task *jobs.Task, envs map[string]*compute.Environment) error { env := envs[task.EnvironmentKey] if env == nil { return nil @@ -90,20 +80,27 @@ func (m *synthesizeRequirements) synthesizeForTask(ctx context.Context, b *bundl return err } - // command_path is an absolute workspace path (translated during initialize); the - // entry script reads requirements.yaml from its directory. - dir := path.Dir(commandPath) - client := m.client - if client == nil { - client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), dir) + written := make(map[string]bool) + for _, dep := range task.AiRuntimeTask.Deployments { + if dep.CommandPath == "" { + continue + } + // command_path is an absolute workspace path (translated during initialize); + // the entry script reads requirements.yaml from its directory. + dir := path.Dir(dep.CommandPath) + if written[dir] { + continue + } + written[dir] = true + + client, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), dir) if err != nil { return err } - } - - log.Debugf(ctx, "writing %s for job %s task %s to %s", requirementsFileName, jobName, task.TaskKey, dir) - if err := client.Write(ctx, requirementsFileName, bytes.NewReader(content), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { - return fmt.Errorf("failed to upload %s next to command_path %q: %w", requirementsFileName, commandPath, err) + log.Debugf(ctx, "writing %s for job %s task %s to %s", requirementsFileName, jobName, task.TaskKey, dir) + if err := client.Write(ctx, requirementsFileName, bytes.NewReader(content), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return fmt.Errorf("failed to upload %s next to command_path %q: %w", requirementsFileName, dep.CommandPath, err) + } } return nil } diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go index 8e742f0e981..1b2777dfee6 100644 --- a/bundle/config/mutator/aicode/validate.go +++ b/bundle/config/mutator/aicode/validate.go @@ -2,13 +2,14 @@ package aicode import ( "context" - "os" + "fmt" "path/filepath" "github.com/databricks/cli/bundle" "github.com/databricks/cli/bundle/libraries" "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/dyn" + "github.com/databricks/databricks-sdk-go/service/jobs" ) // Validate checks AI Runtime tasks that reference a local code_source_path so @@ -38,14 +39,14 @@ func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics } codePath := jobPath.Append(dyn.Key("tasks"), dyn.Index(i), dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) - diags = diags.Extend(v.validateTask(b, job.GitSource != nil, task.AiRuntimeTask.CodeSourcePath, codePath)) + diags = diags.Extend(v.validateTask(b, job.GitSource, task.AiRuntimeTask.CodeSourcePath, codePath)) } } return diags } -func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { +func (v *validate) validateTask(b *bundle.Bundle, gitSource *jobs.GitSource, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { // Only local code_source_path values are packaged at deploy; remote values // are used as-is and need no validation here. if codeSourcePath == "" || !libraries.IsLocalPath(codeSourcePath) { @@ -57,10 +58,19 @@ func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSour // a pre-built tarball delivered via an `artifacts` block is produced during the // build phase (so it does not exist yet at validate time) and is uploaded as a // file, not packaged here. Only when the path is an existing directory do the - // packaging-specific constraints below apply. + // packaging-specific constraints below apply. A stat failure other than not-exist + // (e.g. unreadable parent) is surfaced rather than silently skipped. localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath)) - info, statErr := os.Stat(localDir) - if statErr != nil || !info.IsDir() { + isDir, err := isExistingDir(localDir) + if err != nil { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: fmt.Sprintf("failed to inspect code_source_path %q: %v", codeSourcePath, err), + Locations: b.Config.GetLocations(codePath.String()), + Paths: []dyn.Path{codePath}, + }} + } + if !isDir { return nil } @@ -68,7 +78,7 @@ func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSour // The deploy engine retrieves task files from git when git_source is set, so // packaging a local directory would be silently ignored. Reject the combination. - if jobHasGitSource { + if gitSource != nil { return diag.Diagnostics{{ Severity: diag.Error, Summary: "ai_runtime_task with a local code_source_path cannot be combined with git_source", From b7a1a160109e08f25fc9bd5cc412282a865c1242 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 12:19:15 +0000 Subject: [PATCH 04/16] acceptance/ai_runtime: assert snapshot contents, filtering, and cache Reworks the local_code_source acceptance test per review: - Assert the uploaded tarball's contents via list_code_snapshot.py: gitignored files (ignored_by_git.txt, data/) are excluded, *.log is excluded via sync.exclude, and data/model.bin is force-included via sync.include. - Use print_requests.py instead of hand-rolled jq (testing.md rule); extend its --del-body to also strip the top-level raw_body so the binary tarball upload doesn't dump into the golden. - Add a cache-miss case: editing a file changes the content hash and re-uploads. Co-authored-by: Isaac --- acceptance/bin/list_code_snapshot.py | 56 +++++++++ acceptance/bin/print_requests.py | 7 +- .../local_code_source/databricks.yml | 4 +- .../local_code_source/output.txt | 115 ++++++++++++++++-- .../ai_runtime_task/local_code_source/script | 40 +++--- .../local_code_source/src/.gitignore | 1 + 6 files changed, 186 insertions(+), 37 deletions(-) create mode 100755 acceptance/bin/list_code_snapshot.py diff --git a/acceptance/bin/list_code_snapshot.py b/acceptance/bin/list_code_snapshot.py new file mode 100755 index 00000000000..67dd6dcdd7d --- /dev/null +++ b/acceptance/bin/list_code_snapshot.py @@ -0,0 +1,56 @@ +#!/usr/bin/env python3 +""" +List the entries of the AI Runtime code snapshot tarball uploaded during deploy. + +Reads out.requests.txt, takes code_source_path from the jobs/create request, +exports that workspace file via the CLI, and prints its sorted tar entries. Used +to assert which local files the snapshot includes (gitignore / sync rules). +""" + +import gzip +import io +import json +import os +import subprocess +import sys +import tarfile + +with open("out.requests.txt") as f: + text = f.read() + +# out.requests.txt is a stream of concatenated (pretty-printed) JSON objects. +decoder = json.JSONDecoder() +requests = [] +pos = 0 +while pos < len(text): + while pos < len(text) and text[pos].isspace(): + pos += 1 + if pos >= len(text): + break + obj, pos = decoder.raw_decode(text, pos) + requests.append(obj) + +code_source_path = None +for req in requests: + body = req.get("body") + if isinstance(body, dict) and req.get("path", "").endswith("/jobs/create"): + code_source_path = body["tasks"][0]["ai_runtime_task"]["code_source_path"] + +if not code_source_path: + sys.exit("no jobs/create request with code_source_path in out.requests.txt") + +# code_source_path is recorded de-prefixed (/Users/...); export needs /Workspace. +remote = "/Workspace" + code_source_path +cli = os.environ["CLI"] +local = "code_snapshot.tar.gz" +subprocess.run( + [cli, "workspace", "export", remote, "--format", "AUTO", "--file", local], + check=True, +) +with open(local, "rb") as f: + data = gzip.decompress(f.read()) +os.remove(local) + +with tarfile.open(fileobj=io.BytesIO(data)) as tar: + for name in sorted(tar.getnames()): + print(name) diff --git a/acceptance/bin/print_requests.py b/acceptance/bin/print_requests.py index 82720ac5b23..96c3fcb56ad 100755 --- a/acceptance/bin/print_requests.py +++ b/acceptance/bin/print_requests.py @@ -226,9 +226,12 @@ def main(): for req in filtered_requests: body = req.get("body") - if isinstance(body, dict): - for field in del_body_fields: + for field in del_body_fields: + # Drop the field from the parsed body, and (for non-JSON bodies) the + # top-level raw_body — used to strip binary/volatile upload payloads. + if isinstance(body, dict): body.pop(field, None) + req.pop(field, None) if args.verbose: print( f"Read {len(data)} chars, {len(requests)} requests, {len(filtered_requests)} after filtering", diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml index ff4a7b5dc1b..97401944141 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml +++ b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml @@ -2,9 +2,11 @@ bundle: name: ai-runtime-test sync: - # Exclude a file from both bundle sync and the code snapshot. + # *.log excluded; data/*.bin force-included despite .gitignore. exclude: - "**/*.log" + include: + - src/data/*.bin resources: jobs: diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt index 33fb18eba49..22d036f8436 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -1,5 +1,5 @@ -=== deploy packages the local code source +=== deploy packages and uploads the local code source >>> [CLI] bundle deploy Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... @@ -7,23 +7,116 @@ Deploying resources... Updating deployment state... Deployment complete! -=== the snapshot tarball is uploaded to the user's repo_snapshots directory -/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz +=== the tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded -=== code_source_path points at the uploaded archive (no /Workspace prefix) -/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz +>>> list_code_snapshot.py +src/.gitignore +src/command.sh +src/data/model.bin +src/train.py -=== command_path is translated to its absolute synced workspace path -/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh +=== code_source_path and command_path are rewritten; requirements.yaml is written next to command_path -=== a requirements.yaml is written next to command_path -/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml +>>> print_requests.py --sort --del-body raw_body /repo_snapshots/ /jobs/create /files/src/requirements.yaml +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", + "q": { + "overwrite": "true" + } +} +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", + "q": { + "overwrite": "true" + } +} +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml", + "q": { + "overwrite": "true" + } +} +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml", + "q": { + "overwrite": "true" + } +} +{ + "method": "POST", + "path": "/api/2.2/jobs/create", + "body": { + "deployment": { + "kind": "BUNDLE", + "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/state/metadata.json" + }, + "edit_mode": "UI_LOCKED", + "environments": [ + { + "environment_key": "default", + "spec": { + "dependencies": [ + "torch>=2.0.0" + ], + "environment_version": "5" + } + } + ], + "format": "MULTI_TASK", + "max_concurrent_runs": 1, + "name": "[default] AI Runtime training", + "queue": { + "enabled": true + }, + "tasks": [ + { + "ai_runtime_task": { + "code_source_path": "/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh", + "compute": { + "accelerator_count": 8, + "accelerator_type": "GPU_8xH100" + } + } + ], + "experiment": "my-training" + }, + "environment_key": "default", + "task_key": "train" + } + ] + } +} -=== re-deploying unchanged code is a cache hit: the tarball is not re-uploaded +=== re-deploy of unchanged code is a cache hit (no tarball upload) >>> [CLI] bundle deploy Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... Deploying resources... Updating deployment state... Deployment complete! -tarball uploads on redeploy: 0 + +>>> print_requests.py --sort --del-body raw_body /repo_snapshots/ + +=== editing a file changes the hash and re-uploads (cache miss) + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> print_requests.py --sort --del-body raw_body /repo_snapshots/ +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", + "q": { + "overwrite": "true" + } +} diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script index d0a9e2ea965..b2e55d41623 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/script +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -1,30 +1,24 @@ -# A local code_source_path is packaged into a content-addressed tarball -# (_.tar.gz) and uploaded to the user's repo_snapshots directory. The -# command_path is translated to its synced workspace path, and a requirements.yaml -# derived from the job's serverless environment is written beside it. +# A local code_source_path is packaged into a content-addressed tarball, uploaded +# to repo_snapshots, and code_source_path/command_path rewritten; a requirements.yaml +# is written next to command_path. -title "deploy packages the local code source\n" +title "deploy packages and uploads the local code source\n" trace $CLI bundle deploy -title "the snapshot tarball is uploaded to the user's repo_snapshots directory\n" -# The upload appears twice: WSFS import-file 404s until the parent dir exists, then -# the filer mkdirs and retries. Collapse to the distinct path. -jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u +title "the tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded\n" +trace list_code_snapshot.py -title "code_source_path points at the uploaded archive (no /Workspace prefix)\n" -jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt +title "code_source_path and command_path are rewritten; requirements.yaml is written next to command_path\n" +# --del-body raw_body drops the binary tarball upload body (kept readable); the +# tarball path itself is asserted via list_code_snapshot.py above. --keep is not +# passed, so print_requests.py consumes out.requests.txt. +trace print_requests.py --sort --del-body raw_body '/repo_snapshots/' '/jobs/create' '/files/src/requirements.yaml' -title "command_path is translated to its absolute synced workspace path\n" -jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.deployments[0].command_path' out.requests.txt - -title "a requirements.yaml is written next to command_path\n" -jq -r 'select(.path | test("/import-file/.*/files/src/requirements.yaml$")) | .path' out.requests.txt | sort -u - -rm out.requests.txt - -title "re-deploying unchanged code is a cache hit: the tarball is not re-uploaded\n" +title "re-deploy of unchanged code is a cache hit (no tarball upload)\n" trace $CLI bundle deploy -echo -n "tarball uploads on redeploy: " -jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u | wc -l | tr -d ' ' +trace print_requests.py --sort --del-body raw_body '/repo_snapshots/' -rm out.requests.txt +title "editing a file changes the hash and re-uploads (cache miss)\n" +update_file.py src/train.py 'print("training")' 'print("training v2")' +trace $CLI bundle deploy +trace print_requests.py --sort --del-body raw_body '/repo_snapshots/' diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore index 7a79c3b103e..54553edaac1 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore @@ -1 +1,2 @@ ignored_by_git.txt +data/ From a164d705140b7b35952423d589945c3251528cf0 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 12:53:53 +0000 Subject: [PATCH 05/16] acceptance/ai_runtime: fix Git Bash path conversion on Windows print_requests.py path filters and the workspace export in list_code_snapshot.py were rewritten by Git Bash's MSYS path conversion on Windows (e.g. /repo_snapshots/ -> C:/Program Files/Git/repo_snapshots/), failing the test. Use // leading filters (matching the existing convention) and set MSYS_NO_PATHCONV for the export. Co-authored-by: Isaac --- acceptance/bin/list_code_snapshot.py | 4 +++- .../bundle/ai_runtime_task/local_code_source/output.txt | 6 +++--- .../bundle/ai_runtime_task/local_code_source/script | 9 +++++---- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/acceptance/bin/list_code_snapshot.py b/acceptance/bin/list_code_snapshot.py index 67dd6dcdd7d..4f798c63a42 100755 --- a/acceptance/bin/list_code_snapshot.py +++ b/acceptance/bin/list_code_snapshot.py @@ -43,9 +43,11 @@ remote = "/Workspace" + code_source_path cli = os.environ["CLI"] local = "code_snapshot.tar.gz" +# MSYS_NO_PATHCONV stops Git Bash on Windows from rewriting the /Workspace path. +env = {**os.environ, "MSYS_NO_PATHCONV": "1"} subprocess.run( [cli, "workspace", "export", remote, "--format", "AUTO", "--file", local], - check=True, + check=True, env=env, ) with open(local, "rb") as f: data = gzip.decompress(f.read()) diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt index 22d036f8436..e7f4e7bbab9 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -17,7 +17,7 @@ src/train.py === code_source_path and command_path are rewritten; requirements.yaml is written next to command_path ->>> print_requests.py --sort --del-body raw_body /repo_snapshots/ /jobs/create /files/src/requirements.yaml +>>> print_requests.py --sort --del-body raw_body //repo_snapshots/ //jobs/create //files/src/requirements.yaml { "method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", @@ -102,7 +102,7 @@ Deploying resources... Updating deployment state... Deployment complete! ->>> print_requests.py --sort --del-body raw_body /repo_snapshots/ +>>> print_requests.py --sort --del-body raw_body //repo_snapshots/ === editing a file changes the hash and re-uploads (cache miss) @@ -112,7 +112,7 @@ Deploying resources... Updating deployment state... Deployment complete! ->>> print_requests.py --sort --del-body raw_body /repo_snapshots/ +>>> print_requests.py --sort --del-body raw_body //repo_snapshots/ { "method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script index b2e55d41623..edee046aac1 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/script +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -10,15 +10,16 @@ trace list_code_snapshot.py title "code_source_path and command_path are rewritten; requirements.yaml is written next to command_path\n" # --del-body raw_body drops the binary tarball upload body (kept readable); the -# tarball path itself is asserted via list_code_snapshot.py above. --keep is not +# tarball path itself is asserted via list_code_snapshot.py above. Filters use a +# leading // so Git Bash on Windows does not path-convert them. --keep is not # passed, so print_requests.py consumes out.requests.txt. -trace print_requests.py --sort --del-body raw_body '/repo_snapshots/' '/jobs/create' '/files/src/requirements.yaml' +trace print_requests.py --sort --del-body raw_body '//repo_snapshots/' '//jobs/create' '//files/src/requirements.yaml' title "re-deploy of unchanged code is a cache hit (no tarball upload)\n" trace $CLI bundle deploy -trace print_requests.py --sort --del-body raw_body '/repo_snapshots/' +trace print_requests.py --sort --del-body raw_body '//repo_snapshots/' title "editing a file changes the hash and re-uploads (cache miss)\n" update_file.py src/train.py 'print("training")' 'print("training v2")' trace $CLI bundle deploy -trace print_requests.py --sort --del-body raw_body '/repo_snapshots/' +trace print_requests.py --sort --del-body raw_body '//repo_snapshots/' From b0ee83790d8760d52bf813962562018c358845bc Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 13:04:11 +0000 Subject: [PATCH 06/16] acceptance/ai_runtime: ruff-format list_code_snapshot.py Split subprocess.run kwargs onto separate lines to satisfy ruff format. Co-authored-by: Isaac --- acceptance/bin/list_code_snapshot.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/acceptance/bin/list_code_snapshot.py b/acceptance/bin/list_code_snapshot.py index 4f798c63a42..8b0ff1bbb93 100755 --- a/acceptance/bin/list_code_snapshot.py +++ b/acceptance/bin/list_code_snapshot.py @@ -47,7 +47,8 @@ env = {**os.environ, "MSYS_NO_PATHCONV": "1"} subprocess.run( [cli, "workspace", "export", remote, "--format", "AUTO", "--file", local], - check=True, env=env, + check=True, + env=env, ) with open(local, "rb") as f: data = gzip.decompress(f.read()) From 3b86ff51482a50d20094077b8ca0a9cdfa7ba90f Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 13:39:56 +0000 Subject: [PATCH 07/16] aicode: package snapshot into the bundle, not a home-dir cache Reworks PackageAndUpload to place the code snapshot inside the bundle instead of uploading it to a shared ~/.air/repo_snapshots workspace cache. The mutator now writes the content-addressed tarball into the bundle's sync tree (.air_snapshots/) and rewrites code_source_path to the workspace path it will occupy once synced; it performs no workspace write itself. Normal bundle file sync uploads it during the deploy phase. Why: - Build phase must not mutate the workspace (it runs before `bundle plan`). The old filer upload violated that; now the only build-phase work is local prepare-and-archive, and the upload happens in the deploy phase via file sync. - No deploy-time side-effect outside the bundle: `bundle destroy` removes the snapshot like any other bundle file, so the home-dir cache can no longer compound across bundles/deploys. - Dedup is preserved: the name is content-addressed and bundle file sync is incremental, so unchanged code keeps the same synced path and is not re-uploaded. Tradeoff: the cache is now per-bundle rather than shared across bundles. See the package doc for the (rejected) TranslatePaths alternative and why archiving stays in the build phase. Co-authored-by: Isaac --- acceptance/bin/list_code_snapshot.py | 4 +- .../local_code_source/output.txt | 31 ++--- .../ai_runtime_task/local_code_source/script | 24 ++-- .../local_code_source/test.toml | 3 + .../config/mutator/aicode/package_upload.go | 113 +++++++++--------- 5 files changed, 77 insertions(+), 98 deletions(-) diff --git a/acceptance/bin/list_code_snapshot.py b/acceptance/bin/list_code_snapshot.py index 8b0ff1bbb93..088829a4456 100755 --- a/acceptance/bin/list_code_snapshot.py +++ b/acceptance/bin/list_code_snapshot.py @@ -39,8 +39,8 @@ if not code_source_path: sys.exit("no jobs/create request with code_source_path in out.requests.txt") -# code_source_path is recorded de-prefixed (/Users/...); export needs /Workspace. -remote = "/Workspace" + code_source_path +# code_source_path is an absolute workspace path (/Workspace/Users/.../files/...). +remote = code_source_path cli = os.environ["CLI"] local = "code_snapshot.tar.gz" # MSYS_NO_PATHCONV stops Git Bash on Windows from rewriting the /Workspace path. diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt index e7f4e7bbab9..c9d0380d7f3 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -15,19 +15,12 @@ src/command.sh src/data/model.bin src/train.py -=== code_source_path and command_path are rewritten; requirements.yaml is written next to command_path +=== code_source_path points into the bundle (.air_snapshots), command_path is rewritten, requirements.yaml written ->>> print_requests.py --sort --del-body raw_body //repo_snapshots/ //jobs/create //files/src/requirements.yaml +>>> print_requests.py --sort --del-body raw_body //.air_snapshots/ //jobs/create //files/src/requirements.yaml { "method": "POST", - "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", - "q": { - "overwrite": "true" - } -} -{ - "method": "POST", - "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", "q": { "overwrite": "true" } @@ -75,7 +68,7 @@ src/train.py "tasks": [ { "ai_runtime_task": { - "code_source_path": "/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", + "code_source_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", "deployments": [ { "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh", @@ -94,17 +87,7 @@ src/train.py } } -=== re-deploy of unchanged code is a cache hit (no tarball upload) - ->>> [CLI] bundle deploy -Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... -Deploying resources... -Updating deployment state... -Deployment complete! - ->>> print_requests.py --sort --del-body raw_body //repo_snapshots/ - -=== editing a file changes the hash and re-uploads (cache miss) +=== editing a file changes the snapshot hash (content-addressed name changes) >>> [CLI] bundle deploy Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... @@ -112,10 +95,10 @@ Deploying resources... Updating deployment state... Deployment complete! ->>> print_requests.py --sort --del-body raw_body //repo_snapshots/ +>>> print_requests.py --sort --del-body raw_body //.air_snapshots/ { "method": "POST", - "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", "q": { "overwrite": "true" } diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script index edee046aac1..c417b896cc8 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/script +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -1,6 +1,7 @@ -# A local code_source_path is packaged into a content-addressed tarball, uploaded -# to repo_snapshots, and code_source_path/command_path rewritten; a requirements.yaml -# is written next to command_path. +# A local directory code_source_path is packaged into a content-addressed tarball +# written into the bundle (.air_snapshots/), uploaded by normal bundle file sync, +# and code_source_path/command_path rewritten; a requirements.yaml is written next +# to command_path. title "deploy packages and uploads the local code source\n" trace $CLI bundle deploy @@ -8,18 +9,13 @@ trace $CLI bundle deploy title "the tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded\n" trace list_code_snapshot.py -title "code_source_path and command_path are rewritten; requirements.yaml is written next to command_path\n" -# --del-body raw_body drops the binary tarball upload body (kept readable); the -# tarball path itself is asserted via list_code_snapshot.py above. Filters use a -# leading // so Git Bash on Windows does not path-convert them. --keep is not +title "code_source_path points into the bundle (.air_snapshots), command_path is rewritten, requirements.yaml written\n" +# --del-body raw_body drops the binary tarball upload body (kept readable). Filters +# use a leading // so Git Bash on Windows does not path-convert them. --keep is not # passed, so print_requests.py consumes out.requests.txt. -trace print_requests.py --sort --del-body raw_body '//repo_snapshots/' '//jobs/create' '//files/src/requirements.yaml' +trace print_requests.py --sort --del-body raw_body '//.air_snapshots/' '//jobs/create' '//files/src/requirements.yaml' -title "re-deploy of unchanged code is a cache hit (no tarball upload)\n" -trace $CLI bundle deploy -trace print_requests.py --sort --del-body raw_body '//repo_snapshots/' - -title "editing a file changes the hash and re-uploads (cache miss)\n" +title "editing a file changes the snapshot hash (content-addressed name changes)\n" update_file.py src/train.py 'print("training")' 'print("training v2")' trace $CLI bundle deploy -trace print_requests.py --sort --del-body raw_body '//repo_snapshots/' +trace print_requests.py --sort --del-body raw_body '//.air_snapshots/' diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml index 6e4639a532a..b49665f6fe9 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml +++ b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml @@ -2,6 +2,9 @@ RecordRequests = true Ignore = [ '.databricks', + # The mutator writes the content-addressed code snapshot here; bundle file sync + # uploads it. It is a generated build artifact, not a fixture. + '.air_snapshots', ] # The archive is content-addressed: _.tar.gz. The hash is stable diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index 97959996b9c..96ce3d8a69d 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -1,19 +1,30 @@ // Package aicode packages a local code directory referenced by an AI Runtime -// task's code_source_path and uploads it to the workspace during deploy. +// task's code_source_path so the deployed job runs against the packaged code. // -// The SDK jobs.AiRuntimeTask.code_source_path field expects a workspace or UC -// volume path to an uploaded code archive; its doc comment states that the CLI -// is responsible for packaging the user's local code directory into that -// archive. This mutator implements that contract for DABs: when a user points -// code_source_path at a local directory, it packages the directory into a -// reproducible tarball (.git and gitignored files excluded), uploads the archive -// to the user's workspace code snapshot directory, and rewrites the field to the -// resulting remote path so the deployed job runs against the uploaded code. Values -// that are already remote are left untouched. +// The SDK jobs.AiRuntimeTask.code_source_path field expects a workspace path to an +// uploaded code archive; its doc comment states that the CLI is responsible for +// packaging the user's local code directory into that archive. This mutator +// implements that contract for DABs: when a user points code_source_path at a local +// directory, it packages the directory into a reproducible tarball (.git and +// gitignored files excluded) written into the bundle's sync tree, and rewrites +// code_source_path to the workspace path the tarball will occupy once synced. The +// tarball is uploaded by the normal bundle file sync during the deploy phase — this +// mutator performs no workspace writes itself, so it is safe to run in the build +// phase (before `bundle plan`). Values that are already remote are left untouched. // -// The archive is content-addressed: its name embeds the SHA-256 of the -// (reproducible) tarball, so an unchanged code directory resolves to the same -// remote path across deploys and re-uploads are skipped (see snapshot_package.go). +// Placing the archive in the bundle (rather than a separate cache) means +// `bundle destroy` removes it like any other bundle file, and bundle file sync is +// incremental so an unchanged archive is not re-uploaded. The archive is +// content-addressed: its name embeds the SHA-256 of the (reproducible) tarball, so +// unchanged code keeps the same name and the same synced path across deploys (see +// snapshot_package.go). +// +// Note for reviewers: code_source_path could alternatively be translated to its +// synced workspace path by mutator.TranslatePaths (like command_path is). That +// runs in the initialize phase, which also runs on `bundle validate`, so the +// archive would have to be materialized there too — writing files during validate. +// Keeping materialization in the build phase (deploy-only) and computing the synced +// path here avoids that. package aicode import ( @@ -25,14 +36,12 @@ import ( "os" "path" "path/filepath" - "strings" "github.com/databricks/cli/bundle" "github.com/databricks/cli/bundle/deploy/files" "github.com/databricks/cli/bundle/libraries" "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/dyn" - "github.com/databricks/cli/libs/filer" "github.com/databricks/cli/libs/fileset" "github.com/databricks/cli/libs/log" libsync "github.com/databricks/cli/libs/sync" @@ -76,16 +85,12 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia return diags } - // The current user is resolved during the initialize phase, which runs before - // this build-phase mutator, so it is always set here. - userDir := "/Workspace/Users/" + b.Config.Workspace.CurrentUser.UserName - - // remotePaths maps each config location to the remote archive path it should - // point to after upload. Built outside the Mutate closure so upload failures - // are reported before any config is rewritten. + // remotePaths maps each config location to the synced workspace path its archive + // will occupy. Built outside the Mutate closure so packaging failures are + // reported before any config is rewritten. remotePaths := make(map[string]string, len(sources)) for _, cs := range sources { - remote, err := packageOne(ctx, b, cs, userDir) + remote, err := packageOne(ctx, b, cs) if err != nil { diags = diags.Extend(diag.FromErr(err)) return diags @@ -111,17 +116,19 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia return diags } -// repoSnapshotsSubdir is the per-user workspace location for code snapshots, -// under the user's home. It matches the Python air CLI (and PR #5897) and is -// deliberately NOT /.internal, which artifacts.CleanUp() deletes at -// the start of every deploy. -const repoSnapshotsSubdir = ".air/repo_snapshots" +// snapshotSubdir is the bundle-local directory (relative to the sync root) that +// code snapshots are written into. It is a dedicated folder — not the user's source +// tree — so a snapshot is never nested inside the directory it snapshots, and the +// generated archives are grouped in one predictable place. It is synced to the +// workspace like the rest of the bundle and removed by `bundle destroy`. +const snapshotSubdir = ".air_snapshots" // packageOne packages the local directory for a single code source into a -// reproducible tarball, uploads it to the user's repo_snapshots dir (skipping the -// upload when a content-identical archive already exists there), and returns the -// remote path the config should point to. -func packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, userDir string) (string, error) { +// reproducible, content-addressed tarball written under the bundle's snapshot +// subdir, and returns the workspace path that archive will occupy once bundle file +// sync uploads it. No workspace write happens here — the file is placed locally and +// carried by the deploy-phase file sync. +func packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource) (string, error) { localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) dirName := filepath.Base(localDir) @@ -138,40 +145,30 @@ func packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, userDir st return "", fmt.Errorf("failed to list files for code_source_path %q: %w", cs.value, err) } - uploadPath := path.Join(userDir, repoSnapshotsSubdir, dirName) - client, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath) - if err != nil { - return "", err - } - - // Build the archive in memory so its content hash can name the upload; the hash - // is computed while gzipping, so this adds no extra pass over the files. + // Build the archive in memory so its content hash can name the file; the hash is + // computed while gzipping, so this adds no extra pass over the files. var buf bytes.Buffer sha, err := buildCodeSnapshot(b.SyncRoot, relBase, files, dirName, &buf) if err != nil { return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) } - archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16]) - // The AI Runtime snapshot fetcher expects code_source_path in the legacy - // "/Users/..." form (no "/Workspace" prefix), matching the Python air CLI. The - // filer needs the "/Workspace/Users/..." form to upload, so upload to uploadPath - // but record the de-prefixed path on the task. - remotePath := strings.TrimPrefix(path.Join(uploadPath, archiveName), "/Workspace") - - // The archive is reproducible, so a matching name means identical content is - // already uploaded: skip the upload and just point the config at it. - if _, err := client.Stat(ctx, archiveName); err == nil { - log.Debugf(ctx, "code snapshot already present at %s, skipping upload", remotePath) - return remotePath, nil - } else if !errors.Is(err, fs.ErrNotExist) { - return "", fmt.Errorf("failed to check for existing code snapshot %q: %w", remotePath, err) - } - if err := client.Write(ctx, archiveName, &buf, filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { - return "", fmt.Errorf("failed to upload code snapshot %q: %w", remotePath, err) + // Write the archive into the bundle's snapshot subdir. Content-addressed name + + // incremental file sync means an unchanged archive is not re-uploaded. + relArchive := path.Join(snapshotSubdir, archiveName) + localArchive := filepath.Join(b.SyncRootPath, filepath.FromSlash(relArchive)) + if err := os.MkdirAll(filepath.Dir(localArchive), 0o755); err != nil { + return "", fmt.Errorf("failed to create snapshot dir for code_source_path %q: %w", cs.value, err) } - return remotePath, nil + if err := os.WriteFile(localArchive, buf.Bytes(), 0o644); err != nil { + return "", fmt.Errorf("failed to write code snapshot for code_source_path %q: %w", cs.value, err) + } + log.Debugf(ctx, "wrote code snapshot %s for code_source_path %q", relArchive, cs.value) + + // The workspace path the archive occupies once file sync uploads it. Matches how + // command_path is translated (workspace.file_path + sync-relative path). + return path.Join(b.Config.Workspace.FilePath, relArchive), nil } // codeSourceFiles returns the files under the code directory (relBase, relative to From d89922b06618e98a9efb2c7594c230045f6914bc Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Fri, 31 Jul 2026 14:07:06 +0000 Subject: [PATCH 08/16] aicode: overlay the code snapshot instead of writing it to disk MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The snapshot is now injected into the bundle sync root as an in-memory overlay file rather than written under the working tree. This keeps the deploy artifact inside the bundle (synced to /files/.air_snapshots/, removed by `bundle destroy`, content-addressed so unchanged code is not re-uploaded) while leaving the user's working tree clean — no generated .tar.gz appears in their checkout. Adds libs/vfs.Overlay: a vfs.Path wrapper that serves a set of in-memory files in addition to a base path, participating in Open/Stat/ReadDir/ReadFile and fs.WalkDir so bundle file sync uploads the overlaid files transparently. PackageAndUpload builds each snapshot in memory and swaps b.SyncRoot for an overlay carrying them. Co-authored-by: Isaac --- .../local_code_source/test.toml | 3 - .../config/mutator/aicode/package_upload.go | 68 +++---- libs/vfs/overlay.go | 168 ++++++++++++++++++ libs/vfs/overlay_test.go | 63 +++++++ 4 files changed, 265 insertions(+), 37 deletions(-) create mode 100644 libs/vfs/overlay.go create mode 100644 libs/vfs/overlay_test.go diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml index b49665f6fe9..6e4639a532a 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml +++ b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml @@ -2,9 +2,6 @@ RecordRequests = true Ignore = [ '.databricks', - # The mutator writes the content-addressed code snapshot here; bundle file sync - # uploads it. It is a generated build artifact, not a fixture. - '.air_snapshots', ] # The archive is content-addressed: _.tar.gz. The hash is stable diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index 96ce3d8a69d..613fefbec12 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -45,6 +45,7 @@ import ( "github.com/databricks/cli/libs/fileset" "github.com/databricks/cli/libs/log" libsync "github.com/databricks/cli/libs/sync" + "github.com/databricks/cli/libs/vfs" ) // codeSourcePattern is the config location of an AI Runtime task's @@ -86,18 +87,29 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia } // remotePaths maps each config location to the synced workspace path its archive - // will occupy. Built outside the Mutate closure so packaging failures are - // reported before any config is rewritten. + // will occupy; overlayFiles maps each archive's sync-relative path to its bytes. + // Both are built before any config mutation so packaging failures are reported + // first. The archives are added to the sync root as in-memory overlay files + // (see below) rather than written to disk, so the user's working tree is not + // dirtied by deploy. remotePaths := make(map[string]string, len(sources)) + overlayFiles := make(map[string][]byte, len(sources)) for _, cs := range sources { - remote, err := packageOne(ctx, b, cs) + relArchive, archive, err := packageOne(ctx, b, cs) if err != nil { diags = diags.Extend(diag.FromErr(err)) return diags } - remotePaths[cs.configPath.String()] = remote + overlayFiles[relArchive] = archive + // The workspace path the archive occupies once file sync uploads it. Matches + // how command_path is translated (workspace.file_path + sync-relative path). + remotePaths[cs.configPath.String()] = path.Join(b.Config.Workspace.FilePath, relArchive) } + // Overlay the archives onto the sync root: bundle file sync walks and uploads + // them like real files, but they never touch the user's working tree. + b.SyncRoot = vfs.Overlay(b.SyncRoot, overlayFiles) + err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { for _, cs := range sources { remote := remotePaths[cs.configPath.String()] @@ -116,19 +128,19 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia return diags } -// snapshotSubdir is the bundle-local directory (relative to the sync root) that -// code snapshots are written into. It is a dedicated folder — not the user's source -// tree — so a snapshot is never nested inside the directory it snapshots, and the -// generated archives are grouped in one predictable place. It is synced to the -// workspace like the rest of the bundle and removed by `bundle destroy`. +// snapshotSubdir is the bundle-local directory (sync-relative) that code snapshots +// live under. It is a dedicated folder — not the user's source tree — so a snapshot +// is never nested inside the directory it snapshots, and the archives are grouped in +// one predictable place. The archives are overlaid onto the sync root in memory, so +// this directory is synced to the workspace (and removed by `bundle destroy`) but is +// never materialized in the user's working tree. const snapshotSubdir = ".air_snapshots" // packageOne packages the local directory for a single code source into a -// reproducible, content-addressed tarball written under the bundle's snapshot -// subdir, and returns the workspace path that archive will occupy once bundle file -// sync uploads it. No workspace write happens here — the file is placed locally and -// carried by the deploy-phase file sync. -func packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource) (string, error) { +// reproducible, content-addressed tarball and returns its sync-relative path plus +// the archive bytes. It performs no disk or workspace write: the caller overlays the +// bytes onto the sync root and the deploy-phase file sync uploads them. +func packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource) (string, []byte, error) { localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) dirName := filepath.Base(localDir) @@ -136,13 +148,13 @@ func packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource) (string, e // sync file list to this directory and to re-base archive entry names under it. relBase, err := filepath.Rel(b.SyncRootPath, localDir) if err != nil { - return "", fmt.Errorf("code_source_path %q: %w", cs.value, err) + return "", nil, fmt.Errorf("code_source_path %q: %w", cs.value, err) } relBase = filepath.ToSlash(relBase) files, err := codeSourceFiles(ctx, b, relBase) if err != nil { - return "", fmt.Errorf("failed to list files for code_source_path %q: %w", cs.value, err) + return "", nil, fmt.Errorf("failed to list files for code_source_path %q: %w", cs.value, err) } // Build the archive in memory so its content hash can name the file; the hash is @@ -150,25 +162,13 @@ func packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource) (string, e var buf bytes.Buffer sha, err := buildCodeSnapshot(b.SyncRoot, relBase, files, dirName, &buf) if err != nil { - return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) - } - archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16]) - - // Write the archive into the bundle's snapshot subdir. Content-addressed name + - // incremental file sync means an unchanged archive is not re-uploaded. - relArchive := path.Join(snapshotSubdir, archiveName) - localArchive := filepath.Join(b.SyncRootPath, filepath.FromSlash(relArchive)) - if err := os.MkdirAll(filepath.Dir(localArchive), 0o755); err != nil { - return "", fmt.Errorf("failed to create snapshot dir for code_source_path %q: %w", cs.value, err) + return "", nil, fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) } - if err := os.WriteFile(localArchive, buf.Bytes(), 0o644); err != nil { - return "", fmt.Errorf("failed to write code snapshot for code_source_path %q: %w", cs.value, err) - } - log.Debugf(ctx, "wrote code snapshot %s for code_source_path %q", relArchive, cs.value) - - // The workspace path the archive occupies once file sync uploads it. Matches how - // command_path is translated (workspace.file_path + sync-relative path). - return path.Join(b.Config.Workspace.FilePath, relArchive), nil + // Content-addressed name + incremental file sync means an unchanged archive keeps + // the same synced path and is not re-uploaded. + relArchive := path.Join(snapshotSubdir, fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16])) + log.Debugf(ctx, "packaged code snapshot %s for code_source_path %q", relArchive, cs.value) + return relArchive, buf.Bytes(), nil } // codeSourceFiles returns the files under the code directory (relBase, relative to diff --git a/libs/vfs/overlay.go b/libs/vfs/overlay.go new file mode 100644 index 00000000000..c2d240ce7bd --- /dev/null +++ b/libs/vfs/overlay.go @@ -0,0 +1,168 @@ +package vfs + +import ( + "bytes" + "io" + "io/fs" + "path" + "time" +) + +// Overlay returns a Path that behaves like base but additionally serves a set of +// in-memory files. It lets a caller inject generated content into a filesystem +// tree (e.g. the bundle sync root) without writing it to disk: the overlaid files +// participate in Open/Stat/ReadDir/ReadFile and in fs.WalkDir walks, but the real +// working tree is left untouched. +// +// Overlaid names are slash-separated paths relative to the root (matching io/fs +// conventions). A name must not collide with a real entry in base; on collision the +// overlaid file wins for direct access, which callers should avoid. +func Overlay(base Path, files map[string][]byte) Path { + // Copy so later mutations of the caller's map don't leak in. + overlay := make(map[string][]byte, len(files)) + // dirs maps each ancestor directory to its overlaid children (base-name -> isFile), + // so ReadDir and fs.WalkDir surface the synthetic entries even when the parent + // directory itself exists only in the overlay. + dirs := make(map[string]map[string]bool) + addChild := func(dir, base string, isFile bool) { + entries, ok := dirs[dir] + if !ok { + entries = make(map[string]bool) + dirs[dir] = entries + } + // Don't downgrade a dir entry to file if seen both ways; files never collide. + if isFile || !entries[base] { + entries[base] = isFile + } + } + for name, data := range files { + clean := path.Clean(name) + overlay[clean] = data + // Register clean under its parent as a file, then each ancestor dir under its + // own parent as a directory, up to ".". + child := clean + isFile := true + for child != "." { + parent := path.Dir(child) + addChild(parent, path.Base(child), isFile) + child = parent + isFile = false + } + } + return &overlayPath{base: base, files: overlay, dirs: dirs} +} + +type overlayPath struct { + base Path + files map[string][]byte + // dirs maps a directory name to its overlaid child base-names (value true = the + // child is an overlaid file, false = an overlaid subdirectory). + dirs map[string]map[string]bool +} + +func (o *overlayPath) Open(name string) (fs.File, error) { + if data, ok := o.files[path.Clean(name)]; ok { + return newMemFile(path.Base(name), data), nil + } + return o.base.Open(name) +} + +func (o *overlayPath) Stat(name string) (fs.FileInfo, error) { + if data, ok := o.files[path.Clean(name)]; ok { + return memFileInfo{name: path.Base(name), size: int64(len(data))}, nil + } + return o.base.Stat(name) +} + +func (o *overlayPath) ReadFile(name string) ([]byte, error) { + if data, ok := o.files[path.Clean(name)]; ok { + return append([]byte(nil), data...), nil + } + return o.base.ReadFile(name) +} + +func (o *overlayPath) ReadDir(name string) ([]fs.DirEntry, error) { + clean := path.Clean(name) + baseEntries, err := o.base.ReadDir(name) + // A dir that exists only in the overlay (e.g. .air_snapshots) has no real + // counterpart; tolerate a not-exist error when we have overlaid children for it. + overlaid := o.dirs[clean] + if err != nil && overlaid == nil { + return nil, err + } + + seen := make(map[string]bool, len(baseEntries)) + entries := make([]fs.DirEntry, 0, len(baseEntries)+len(overlaid)) + for _, e := range baseEntries { + seen[e.Name()] = true + entries = append(entries, e) + } + for child, isFile := range overlaid { + if seen[child] { + continue + } + if isFile { + full := child + if clean != "." { + full = clean + "/" + child + } + entries = append(entries, memDirEntry{name: child, size: int64(len(o.files[full]))}) + } else { + entries = append(entries, memDirEntry{name: child, dir: true}) + } + } + return entries, nil +} + +func (o *overlayPath) Parent() Path { return o.base.Parent() } +func (o *overlayPath) Native() string { return o.base.Native() } + +// memFile is an in-memory fs.File for an overlaid file. +type memFile struct { + info memFileInfo + reader *bytes.Reader +} + +func newMemFile(name string, data []byte) *memFile { + return &memFile{ + info: memFileInfo{name: name, size: int64(len(data))}, + reader: bytes.NewReader(data), + } +} + +func (f *memFile) Stat() (fs.FileInfo, error) { return f.info, nil } +func (f *memFile) Read(p []byte) (int, error) { return f.reader.Read(p) } +func (f *memFile) Close() error { return nil } + +type memFileInfo struct { + name string + size int64 +} + +func (i memFileInfo) Name() string { return i.name } +func (i memFileInfo) Size() int64 { return i.size } +func (i memFileInfo) Mode() fs.FileMode { return 0o644 } +func (i memFileInfo) ModTime() time.Time { return time.Time{} } +func (i memFileInfo) IsDir() bool { return false } +func (i memFileInfo) Sys() any { return nil } + +type memDirEntry struct { + name string + dir bool + size int64 +} + +func (e memDirEntry) Name() string { return e.name } +func (e memDirEntry) IsDir() bool { return e.dir } +func (e memDirEntry) Type() fs.FileMode { + if e.dir { + return fs.ModeDir + } + return 0 +} + +func (e memDirEntry) Info() (fs.FileInfo, error) { + return memFileInfo{name: e.name, size: e.size}, nil +} + +var _ io.Reader = (*memFile)(nil) diff --git a/libs/vfs/overlay_test.go b/libs/vfs/overlay_test.go new file mode 100644 index 00000000000..60e0664410a --- /dev/null +++ b/libs/vfs/overlay_test.go @@ -0,0 +1,63 @@ +package vfs + +import ( + "io" + "io/fs" + "os" + "path/filepath" + "slices" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestOverlayServesVirtualFileAndRealFiles(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "real.txt"), []byte("real"), 0o644)) + + base := MustNew(dir) + ov := Overlay(base, map[string][]byte{ + ".air_snapshots/src_abc.tar.gz": []byte("SNAPSHOT-BYTES"), + }) + + // Real file still readable through the overlay. + got, err := ov.ReadFile("real.txt") + require.NoError(t, err) + assert.Equal(t, "real", string(got)) + + // Virtual file readable via Open (the path sync's applyPut uses). + f, err := ov.Open(".air_snapshots/src_abc.tar.gz") + require.NoError(t, err) + b, err := io.ReadAll(f) + require.NoError(t, err) + require.NoError(t, f.Close()) + assert.Equal(t, "SNAPSHOT-BYTES", string(b)) + + // Virtual file is NOT on disk — the working tree stays clean. + assert.NoFileExists(t, filepath.Join(dir, ".air_snapshots", "src_abc.tar.gz")) +} + +func TestOverlayWalkDirSurfacesVirtualFile(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "real.txt"), []byte("real"), 0o644)) + ov := Overlay(MustNew(dir), map[string][]byte{ + ".air_snapshots/src_abc.tar.gz": []byte("x"), + }) + + var walked []string + require.NoError(t, fs.WalkDir(ov, ".", func(name string, d fs.DirEntry, err error) error { + if err != nil { + return err + } + if !d.IsDir() { + walked = append(walked, name) + } + return nil + })) + + slices.Sort(walked) + // Both the real file and the virtual snapshot (in its virtual dir) are walked, + // so bundle file sync uploads the snapshot without it existing on disk. + assert.Equal(t, []string{".air_snapshots/src_abc.tar.gz", "real.txt"}, walked) +} From 4e4a18533ebaa6cb629beea29aea22b184661283 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 04:22:37 +0000 Subject: [PATCH 09/16] =?UTF-8?q?aicode:=20address=20review=20=E2=80=94=20?= =?UTF-8?q?validation=20guards=20+=20unfilterable=20snapshot=20dir?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fixes from PR review: - Reject a code_source_path that escapes the bundle sync root (e.g. "../shared") in Validate with a clear message, instead of failing later as an opaque io/fs "invalid argument". - Reject a local code_source_path combined with source-linked deployment: that mode makes files.Upload a no-op and ignores workspace.file_path, so the packaged snapshot would never be uploaded and the job would point at a missing path. - Reject a local code_source_path nested under a for_each_task, which the mutator does not package (it operates on direct tasks only) — previously silently skipped. - Error when the code_source directory has no files to package (all excluded by .gitignore/sync.exclude, or empty) rather than deploying a job with no code. - Force-include the .air_snapshots snapshot dir in GetSyncIncludePatterns (like .databricks) so a user ignore rule such as "*.tar.gz" can't filter the deployed job's code_source_path archive out of the sync set. Shared bundle.AiCodeSnapshotDir constant keeps the mutator and the sync include in agreement. Co-authored-by: Isaac --- bundle/bundle.go | 13 ++- .../config/mutator/aicode/package_upload.go | 11 ++- bundle/config/mutator/aicode/validate.go | 84 ++++++++++++++----- bundle/config/mutator/aicode/validate_test.go | 56 +++++++++++++ 4 files changed, 138 insertions(+), 26 deletions(-) diff --git a/bundle/bundle.go b/bundle/bundle.go index 94f55389b26..46c3d3b2cfe 100644 --- a/bundle/bundle.go +++ b/bundle/bundle.go @@ -36,6 +36,12 @@ import ( const internalFolder = ".internal" +// AiCodeSnapshotDir is the sync-root-relative directory (dot-prefixed so it does not +// clash with user files) that the aicode mutator writes AI Runtime code snapshots +// into. It is force-included in sync (see GetSyncIncludePatterns) so user ignore +// rules cannot filter the deployed job's code_source_path archives out. +const AiCodeSnapshotDir = ".air_snapshots" + // Filename where resources are stored for DATABRICKS_BUNDLE_ENGINE=direct const resourcesFilename = "resources.json" @@ -358,7 +364,12 @@ func (b *Bundle) GetSyncIncludePatterns(ctx context.Context) ([]string, error) { if err != nil { return nil, err } - return append(b.Config.Sync.Include, filepath.ToSlash(filepath.Join(internalDirRel, "*.*"))), nil + includes := append(b.Config.Sync.Include, filepath.ToSlash(filepath.Join(internalDirRel, "*.*"))) + // Force-sync generated AI Runtime code snapshots so a user ignore rule (e.g. + // "*.tar.gz" in .gitignore) can't filter them out — the deployed job's + // code_source_path points at these archives (see bundle/config/mutator/aicode). + includes = append(includes, AiCodeSnapshotDir+"/*") + return includes, nil } // AuthEnv returns a map with environment variables and their values diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index 613fefbec12..ee2d5d8f944 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -133,8 +133,9 @@ func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Dia // is never nested inside the directory it snapshots, and the archives are grouped in // one predictable place. The archives are overlaid onto the sync root in memory, so // this directory is synced to the workspace (and removed by `bundle destroy`) but is -// never materialized in the user's working tree. -const snapshotSubdir = ".air_snapshots" +// never materialized in the user's working tree. bundle.GetSyncIncludePatterns +// force-includes it so a user ignore rule can't filter the archives out of sync. +const snapshotSubdir = bundle.AiCodeSnapshotDir // packageOne packages the local directory for a single code source into a // reproducible, content-addressed tarball and returns its sync-relative path plus @@ -156,6 +157,12 @@ func packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource) (string, [ if err != nil { return "", nil, fmt.Errorf("failed to list files for code_source_path %q: %w", cs.value, err) } + // An empty file list means every file under the directory was filtered out + // (gitignore / sync.exclude) or the directory is empty. Packaging it would deploy + // a job with no code, so fail with an actionable message instead. + if len(files) == 0 { + return "", nil, fmt.Errorf("code_source_path %q has no files to package (all excluded by .gitignore or sync.exclude, or the directory is empty)", cs.value) + } // Build the archive in memory so its content hash can name the file; the hash is // computed while gzipping, so this adds no extra pass over the files. diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go index 1b2777dfee6..69e9d6300f2 100644 --- a/bundle/config/mutator/aicode/validate.go +++ b/bundle/config/mutator/aicode/validate.go @@ -6,6 +6,7 @@ import ( "path/filepath" "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/libraries" "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/dyn" @@ -34,11 +35,30 @@ func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics jobPath := jobsPath.Append(dyn.Key(name)) for i, task := range job.Tasks { + taskPath := jobPath.Append(dyn.Key("tasks"), dyn.Index(i)) + + // A local code_source_path under a for_each_task is not packaged by this + // mutator (aicode collects only direct tasks). Reject it rather than let a + // nested ai_runtime_task deploy an un-packaged local path. + if task.ForEachTask != nil && task.ForEachTask.Task.AiRuntimeTask != nil { + nestedCode := task.ForEachTask.Task.AiRuntimeTask.CodeSourcePath + if nestedCode != "" && libraries.IsLocalPath(nestedCode) { + p := taskPath.Append(dyn.Key("for_each_task"), dyn.Key("task"), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) + diags = diags.Append(diag.Diagnostic{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path is not supported inside a for_each_task", + Detail: "Set code_source_path to a workspace or volume path", + Locations: b.Config.GetLocations(p.String()), + Paths: []dyn.Path{p}, + }) + } + } + if task.AiRuntimeTask == nil { continue } - codePath := jobPath.Append(dyn.Key("tasks"), dyn.Index(i), - dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) + codePath := taskPath.Append(dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) diags = diags.Extend(v.validateTask(b, job.GitSource, task.AiRuntimeTask.CodeSourcePath, codePath)) } } @@ -53,6 +73,28 @@ func (v *validate) validateTask(b *bundle.Bundle, gitSource *jobs.GitSource, cod return nil } + locations := b.Config.GetLocations(codePath.String()) + reject := func(summary, detail string) diag.Diagnostics { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: summary, + Detail: detail, + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + // The packaged directory must live inside the bundle sync root: it is uploaded as + // part of the bundle, so a path escaping the root (e.g. "../shared") can't be + // synced. Reject it here with a clear message rather than letting it fail later as + // an opaque io/fs "invalid argument" when the file list is built. + if rel, err := filepath.Rel(b.SyncRootPath, filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath))); err != nil || !filepath.IsLocal(rel) { + return reject( + fmt.Sprintf("code_source_path %q is outside the bundle root", codeSourcePath), + "code_source_path must point at a directory inside the bundle, or at a workspace or volume path", + ) + } + // This mutator packages a local *directory*. A local path that is not an existing // directory is left alone so it flows through the standard artifact-upload path: // a pre-built tarball delivered via an `artifacts` block is produced during the @@ -63,40 +105,36 @@ func (v *validate) validateTask(b *bundle.Bundle, gitSource *jobs.GitSource, cod localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath)) isDir, err := isExistingDir(localDir) if err != nil { - return diag.Diagnostics{{ - Severity: diag.Error, - Summary: fmt.Sprintf("failed to inspect code_source_path %q: %v", codeSourcePath, err), - Locations: b.Config.GetLocations(codePath.String()), - Paths: []dyn.Path{codePath}, - }} + return reject(fmt.Sprintf("failed to inspect code_source_path %q: %v", codeSourcePath, err), "") } if !isDir { return nil } - locations := b.Config.GetLocations(codePath.String()) - // The deploy engine retrieves task files from git when git_source is set, so // packaging a local directory would be silently ignored. Reject the combination. if gitSource != nil { - return diag.Diagnostics{{ - Severity: diag.Error, - Summary: "ai_runtime_task with a local code_source_path cannot be combined with git_source", - Detail: "Remove git_source or set code_source_path to a workspace or volume path", - Locations: locations, - Paths: []dyn.Path{codePath}, - }} + return reject( + "ai_runtime_task with a local code_source_path cannot be combined with git_source", + "Remove git_source or set code_source_path to a workspace or volume path", + ) } // Immutable-folder deployments upload a single content-addressed snapshot and // do not support the per-task code packaging this mutator performs. if b.IsImmutableFolder() { - return diag.Diagnostics{{ - Severity: diag.Error, - Summary: "ai_runtime_task with a local code_source_path is not supported with experimental.immutable_folder", - Locations: locations, - Paths: []dyn.Path{codePath}, - }} + return reject("ai_runtime_task with a local code_source_path is not supported with experimental.immutable_folder", "") + } + + // Source-linked deployment runs jobs against the source files in place and does + // not copy them to the workspace file path (files.Upload is a no-op). This + // mutator relies on file sync uploading the packaged snapshot, so the two are + // incompatible: reject rather than deploy a job pointing at an un-uploaded path. + if config.IsExplicitlyEnabled(b.Config.Presets.SourceLinkedDeployment) { + return reject( + "ai_runtime_task with a local code_source_path is not supported with source-linked deployment", + "Disable source-linked deployment, or set code_source_path to a workspace or volume path", + ) } return nil diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go index bddcc8f309a..1839e8c227d 100644 --- a/bundle/config/mutator/aicode/validate_test.go +++ b/bundle/config/mutator/aicode/validate_test.go @@ -80,3 +80,59 @@ func TestValidateRemoteCodeSourceIsSkipped(t *testing.T) { diags := Validate().Apply(t.Context(), b) assert.Empty(t, diags) } + +// A code_source_path escaping the bundle sync root is rejected with a clear +// message (it can't be synced), rather than failing later as an opaque io/fs error. +func TestValidateCodeSourceOutsideBundleRoot(t *testing.T) { + b := bundleForValidate(t, "../shared", nil) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "outside the bundle root") +} + +// Source-linked deployment doesn't copy files to the workspace file path, so the +// packaged snapshot would never be uploaded; the combination is rejected. +func TestValidateSourceLinkedConflict(t *testing.T) { + b := bundleForValidate(t, "src", nil) + mkCodeDir(t, b, "src") + enabled := true + b.Config.Presets.SourceLinkedDeployment = &enabled + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "source-linked deployment") +} + +// A local code_source_path nested under a for_each_task is not packaged by the +// mutator, so it is rejected rather than silently skipped. +func TestValidateForEachTaskCodeSourceRejected(t *testing.T) { + dir := t.TempDir() + b := &bundle.Bundle{ + BundleRootPath: dir, + SyncRootPath: dir, + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "train": { + JobSettings: jobs.JobSettings{ + Tasks: []jobs.Task{ + { + TaskKey: "fanout", + ForEachTask: &jobs.ForEachTask{ + Task: jobs.Task{ + AiRuntimeTask: &jobs.AiRuntimeTask{CodeSourcePath: "src"}, + }, + }, + }, + }, + }, + }, + }, + }, + }, + } + bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "for_each_task") +} From 2fc2844ef23410404d3d57796999d50afddbabab Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 04:25:42 +0000 Subject: [PATCH 10/16] aicode: drop SynthesizeRequirements; deps ride on environments spec MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remove the SynthesizeRequirements mutator. It wrote a co-located requirements.yaml next to command_path during the build phase — which (a) made workspace writes in a phase that runs before `bundle plan` and (b) duplicated dependency delivery that the native path already handles. The AI Runtime runtime installs pip deps from the job's serverless environment (environments[].spec.dependencies) via the server-side --deps-config path; the co-located requirements.yaml is the legacy mechanism new clients no longer create. convert-to-dabs already emits deps into environments[].spec, so the sidecar was redundant. Verified end-to-end on A10: with SynthesizeRequirements disabled and no requirements.yaml uploaded, the run logs "No co-located requirements.yaml ... skipping" and installs torch from the inline environment deps to SUCCESS. Co-authored-by: Isaac --- .../local_code_source/output.txt | 18 +-- .../ai_runtime_task/local_code_source/script | 8 +- .../config/mutator/aicode/package_upload.go | 5 +- bundle/config/mutator/aicode/requirements.go | 141 ------------------ .../mutator/aicode/requirements_test.go | 36 ----- bundle/phases/build.go | 12 +- 6 files changed, 12 insertions(+), 208 deletions(-) delete mode 100644 bundle/config/mutator/aicode/requirements.go delete mode 100644 bundle/config/mutator/aicode/requirements_test.go diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt index c9d0380d7f3..df6b0f54acc 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -15,9 +15,9 @@ src/command.sh src/data/model.bin src/train.py -=== code_source_path points into the bundle (.air_snapshots), command_path is rewritten, requirements.yaml written +=== code_source_path points into the bundle (.air_snapshots), command_path rewritten, deps carried on environments spec ->>> print_requests.py --sort --del-body raw_body //.air_snapshots/ //jobs/create //files/src/requirements.yaml +>>> print_requests.py --sort --del-body raw_body //.air_snapshots/ //jobs/create { "method": "POST", "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", @@ -25,20 +25,6 @@ src/train.py "overwrite": "true" } } -{ - "method": "POST", - "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml", - "q": { - "overwrite": "true" - } -} -{ - "method": "POST", - "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml", - "q": { - "overwrite": "true" - } -} { "method": "POST", "path": "/api/2.2/jobs/create", diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script index c417b896cc8..cd704f28119 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/script +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -1,7 +1,7 @@ # A local directory code_source_path is packaged into a content-addressed tarball # written into the bundle (.air_snapshots/), uploaded by normal bundle file sync, -# and code_source_path/command_path rewritten; a requirements.yaml is written next -# to command_path. +# and code_source_path/command_path rewritten. Pip deps ride on the job's +# environments[].spec.dependencies (no requirements.yaml is synthesized). title "deploy packages and uploads the local code source\n" trace $CLI bundle deploy @@ -9,11 +9,11 @@ trace $CLI bundle deploy title "the tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded\n" trace list_code_snapshot.py -title "code_source_path points into the bundle (.air_snapshots), command_path is rewritten, requirements.yaml written\n" +title "code_source_path points into the bundle (.air_snapshots), command_path rewritten, deps carried on environments spec\n" # --del-body raw_body drops the binary tarball upload body (kept readable). Filters # use a leading // so Git Bash on Windows does not path-convert them. --keep is not # passed, so print_requests.py consumes out.requests.txt. -trace print_requests.py --sort --del-body raw_body '//.air_snapshots/' '//jobs/create' '//files/src/requirements.yaml' +trace print_requests.py --sort --del-body raw_body '//.air_snapshots/' '//jobs/create' title "editing a file changes the snapshot hash (content-addressed name changes)\n" update_file.py src/train.py 'print("training")' 'print("training v2")' diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index ee2d5d8f944..6c901a347f3 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -50,9 +50,8 @@ import ( // codeSourcePattern is the config location of an AI Runtime task's // code_source_path. It matches a direct task only — the same scope aicode.Validate -// and aicode.SynthesizeRequirements operate on. ai_runtime_task nested under a -// for_each_task is not a supported combination yet; when it is, all three mutators -// should gain it together. +// operates on. ai_runtime_task nested under a for_each_task is not a supported +// combination yet (Validate rejects it); when it is, both should gain it together. var codeSourcePattern = dyn.NewPattern( dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), dyn.Key("tasks"), dyn.AnyIndex(), diff --git a/bundle/config/mutator/aicode/requirements.go b/bundle/config/mutator/aicode/requirements.go deleted file mode 100644 index 8927500d78d..00000000000 --- a/bundle/config/mutator/aicode/requirements.go +++ /dev/null @@ -1,141 +0,0 @@ -package aicode - -import ( - "bytes" - "context" - "fmt" - "path" - - "github.com/databricks/cli/bundle" - "github.com/databricks/cli/libs/diag" - "github.com/databricks/cli/libs/filer" - "github.com/databricks/cli/libs/log" - "github.com/databricks/databricks-sdk-go/service/compute" - "github.com/databricks/databricks-sdk-go/service/jobs" - "go.yaml.in/yaml/v3" -) - -// requirementsFileName is the file the AI Runtime entry script reads from the -// command_path directory to install the workload's pip dependencies. The server -// derives its path from command_path, so it must sit next to command.sh. -const requirementsFileName = "requirements.yaml" - -// requirementsSpec is the requirements.yaml shape the AI Runtime consumes: a -// client image version and a list of pip requirement lines. It mirrors what the -// Python air CLI synthesizes. -type requirementsSpec struct { - Version string `yaml:"version,omitempty"` - Dependencies []string `yaml:"dependencies,omitempty"` -} - -// SynthesizeRequirements uploads a requirements.yaml next to each AI Runtime task's -// command_path, derived from the job-level serverless environment referenced by the -// task's environment_key. The AI Runtime entry script reads this file (resolved -// from command_path's directory) to set up the workload environment; without it the -// run fails during setup. It runs at deploy time, after command_path has been -// translated to its absolute workspace path. -func SynthesizeRequirements() bundle.Mutator { - return &synthesizeRequirements{} -} - -type synthesizeRequirements struct{} - -func (m *synthesizeRequirements) Name() string { - return "aicode.SynthesizeRequirements" -} - -func (m *synthesizeRequirements) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { - var diags diag.Diagnostics - - for name, job := range b.Config.Resources.Jobs { - envs := environmentsByKey(job.Environments) - for i := range job.Tasks { - task := &job.Tasks[i] - if task.AiRuntimeTask == nil { - continue - } - if err := synthesizeForTask(ctx, b, name, task, envs); err != nil { - diags = diags.Extend(diag.FromErr(err)) - } - } - } - - return diags -} - -// synthesizeForTask uploads requirements.yaml next to each deployment's -// command_path. It is a no-op when the task has no matching environment spec (the -// run can still supply deps another way, so this is not an error). Each deployment -// carries its own command_path and the entry script reads requirements.yaml from -// that directory, so every deployment's directory needs the file; directories are -// deduplicated so a shared command_path directory is written only once. -func synthesizeForTask(ctx context.Context, b *bundle.Bundle, jobName string, task *jobs.Task, envs map[string]*compute.Environment) error { - env := envs[task.EnvironmentKey] - if env == nil { - return nil - } - - content, err := renderRequirements(env) - if err != nil { - return err - } - - written := make(map[string]bool) - for _, dep := range task.AiRuntimeTask.Deployments { - if dep.CommandPath == "" { - continue - } - // command_path is an absolute workspace path (translated during initialize); - // the entry script reads requirements.yaml from its directory. - dir := path.Dir(dep.CommandPath) - if written[dir] { - continue - } - written[dir] = true - - client, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), dir) - if err != nil { - return err - } - log.Debugf(ctx, "writing %s for job %s task %s to %s", requirementsFileName, jobName, task.TaskKey, dir) - if err := client.Write(ctx, requirementsFileName, bytes.NewReader(content), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { - return fmt.Errorf("failed to upload %s next to command_path %q: %w", requirementsFileName, dep.CommandPath, err) - } - } - return nil -} - -// renderRequirements builds the requirements.yaml content from a serverless -// environment spec: version from environment_version (or the legacy client field), -// dependencies from its pip dependency list. -func renderRequirements(env *compute.Environment) ([]byte, error) { - version := env.EnvironmentVersion - if version == "" { - version = env.Client - } - spec := requirementsSpec{ - Version: version, - Dependencies: env.Dependencies, - } - var buf bytes.Buffer - enc := yaml.NewEncoder(&buf) - enc.SetIndent(2) - if err := enc.Encode(spec); err != nil { - return nil, err - } - if err := enc.Close(); err != nil { - return nil, err - } - return buf.Bytes(), nil -} - -// environmentsByKey indexes a job's environments by their key for task lookup. -func environmentsByKey(envs []jobs.JobEnvironment) map[string]*compute.Environment { - out := make(map[string]*compute.Environment, len(envs)) - for i := range envs { - if envs[i].Spec != nil { - out[envs[i].EnvironmentKey] = envs[i].Spec - } - } - return out -} diff --git a/bundle/config/mutator/aicode/requirements_test.go b/bundle/config/mutator/aicode/requirements_test.go deleted file mode 100644 index a8044e3e35c..00000000000 --- a/bundle/config/mutator/aicode/requirements_test.go +++ /dev/null @@ -1,36 +0,0 @@ -package aicode - -import ( - "testing" - - "github.com/databricks/databricks-sdk-go/service/compute" - "github.com/databricks/databricks-sdk-go/service/jobs" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -func TestRenderRequirementsFromEnvironmentVersion(t *testing.T) { - out, err := renderRequirements(&compute.Environment{ - EnvironmentVersion: "5", - Dependencies: []string{"numpy", "torch>=2.0.0"}, - }) - require.NoError(t, err) - assert.Equal(t, "version: \"5\"\ndependencies:\n - numpy\n - torch>=2.0.0\n", string(out)) -} - -func TestRenderRequirementsFallsBackToClient(t *testing.T) { - out, err := renderRequirements(&compute.Environment{Client: "5"}) - require.NoError(t, err) - assert.Equal(t, "version: \"5\"\n", string(out)) -} - -func TestEnvironmentsByKey(t *testing.T) { - envs := []jobs.JobEnvironment{ - {EnvironmentKey: "default", Spec: &compute.Environment{EnvironmentVersion: "5"}}, - {EnvironmentKey: "nospec"}, - } - got := environmentsByKey(envs) - require.Contains(t, got, "default") - assert.Equal(t, "5", got["default"].EnvironmentVersion) - assert.NotContains(t, got, "nospec", "environments without a spec are skipped") -} diff --git a/bundle/phases/build.go b/bundle/phases/build.go index a386093bc24..febfb66d669 100644 --- a/bundle/phases/build.go +++ b/bundle/phases/build.go @@ -29,15 +29,11 @@ func Build(ctx context.Context, b *bundle.Bundle) LibLocationMap { scripts.Execute(config.ScriptPostBuild), // Package any AI Runtime task code_source_path that points at a local - // directory into a tarball, upload it, and rewrite the field to the remote - // path. Runs before ExpandGlobReferences/ReplaceWithRemotePath below so those - // see an already-remote code_source_path (a directory would otherwise be - // collected as a local "library" and fail to upload as a file). + // directory into a content-addressed tarball overlaid on the sync root, and + // rewrite the field to the synced workspace path. No requirements.yaml is + // synthesized: the runtime installs pip deps from the job's serverless + // environment (environments[].spec.dependencies) directly. aicode.PackageAndUpload(), - // Write requirements.yaml next to the (already translated) command_path, - // derived from the job's serverless environment, so the AI Runtime harness - // can set up the workload environment. - aicode.SynthesizeRequirements(), mutator.ResolveVariableReferencesWithoutResources( "artifacts", From ac23267c984c8afc0c719dfbddb5078186d1d089 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 16:56:16 +0000 Subject: [PATCH 11/16] aicode: guard reserved snapshot dir; rename; broaden acceptance MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Second review round: - Reject a sync.exclude that matches the reserved .air_snapshots dir. The sync set is (git ∪ include) − exclude, so exclude wins over the force-include; a pattern like ".air_snapshots/*" or "**/*.tar.gz" would silently drop the generated archive and leave the job pointing at an un-uploaded path. Validate now errors. - Reject a real file/dir at .air_snapshots in the bundle: it would collide with the in-memory overlay (sync would carry the user's entry, not the generated archive). - Both guards only fire when a task actually packages a local code_source. - Rename PackageAndUpload -> PackageCodeSource: it packages and overlays the archive for file sync to upload; it no longer uploads itself. - Acceptance test: add a second AI Runtime task so two tarballs are packaged, both under the repo-root .air_snapshots; assert a no-op re-deploy; and destroy at the end so the test cleans up after itself. list_code_snapshot.py now lists every task's archive. Co-authored-by: Isaac --- acceptance/bin/list_code_snapshot.py | 54 +++++++++------ .../local_code_source/databricks.yml | 12 ++++ .../local_code_source/output.txt | 59 ++++++++++++++-- .../ai_runtime_task/local_code_source/script | 20 ++++-- .../local_code_source/src2/command.sh | 2 + .../local_code_source/src2/train.py | 1 + .../local_code_source/test.toml | 6 +- .../config/mutator/aicode/package_upload.go | 12 ++-- bundle/config/mutator/aicode/validate.go | 68 +++++++++++++++++++ bundle/config/mutator/aicode/validate_test.go | 43 ++++++++++++ bundle/phases/build.go | 2 +- 11 files changed, 235 insertions(+), 44 deletions(-) create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src2/command.sh create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src2/train.py diff --git a/acceptance/bin/list_code_snapshot.py b/acceptance/bin/list_code_snapshot.py index 088829a4456..917d2eba6ec 100755 --- a/acceptance/bin/list_code_snapshot.py +++ b/acceptance/bin/list_code_snapshot.py @@ -1,10 +1,11 @@ #!/usr/bin/env python3 """ -List the entries of the AI Runtime code snapshot tarball uploaded during deploy. +List the entries of each AI Runtime code snapshot tarball uploaded during deploy. -Reads out.requests.txt, takes code_source_path from the jobs/create request, -exports that workspace file via the CLI, and prints its sorted tar entries. Used -to assert which local files the snapshot includes (gitignore / sync rules). +Reads out.requests.txt, takes every ai_runtime_task.code_source_path from the +jobs/create request, exports each workspace archive via the CLI, and prints its +sorted tar entries (grouped per archive). Used to assert which local files each +snapshot includes (gitignore / sync rules), across one or more tasks. """ import gzip @@ -30,30 +31,39 @@ obj, pos = decoder.raw_decode(text, pos) requests.append(obj) -code_source_path = None +# Collect every task's code_source_path from the jobs/create request(s). +code_source_paths = [] for req in requests: body = req.get("body") if isinstance(body, dict) and req.get("path", "").endswith("/jobs/create"): - code_source_path = body["tasks"][0]["ai_runtime_task"]["code_source_path"] + for task in body.get("tasks", []): + art = task.get("ai_runtime_task") + if art and art.get("code_source_path"): + code_source_paths.append(art["code_source_path"]) -if not code_source_path: +if not code_source_paths: sys.exit("no jobs/create request with code_source_path in out.requests.txt") -# code_source_path is an absolute workspace path (/Workspace/Users/.../files/...). -remote = code_source_path cli = os.environ["CLI"] -local = "code_snapshot.tar.gz" # MSYS_NO_PATHCONV stops Git Bash on Windows from rewriting the /Workspace path. env = {**os.environ, "MSYS_NO_PATHCONV": "1"} -subprocess.run( - [cli, "workspace", "export", remote, "--format", "AUTO", "--file", local], - check=True, - env=env, -) -with open(local, "rb") as f: - data = gzip.decompress(f.read()) -os.remove(local) - -with tarfile.open(fileobj=io.BytesIO(data)) as tar: - for name in sorted(tar.getnames()): - print(name) + +# code_source_path is an absolute workspace path (/Workspace/Users/.../files/...). +# Sort so multi-task output is deterministic. +for remote in sorted(code_source_paths): + local = "code_snapshot.tar.gz" + subprocess.run( + [cli, "workspace", "export", remote, "--format", "AUTO", "--file", local], + check=True, + env=env, + ) + with open(local, "rb") as f: + data = gzip.decompress(f.read()) + os.remove(local) + + # Print the archive's sync-relative name (hash tokenized by test.toml repls) so + # multi-archive output is legible. + print(f"# {remote.split('/files/', 1)[-1]}") + with tarfile.open(fileobj=io.BytesIO(data)) as tar: + for name in sorted(tar.getnames()): + print(name) diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml index 97401944141..4a1cc0476c0 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml +++ b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml @@ -13,6 +13,8 @@ resources: train: name: "[${bundle.target}] AI Runtime training" tasks: + # Two AI Runtime tasks with distinct local code dirs: each is packaged into + # its own content-addressed tarball, both under the repo root's .air_snapshots. - task_key: train environment_key: default ai_runtime_task: @@ -23,6 +25,16 @@ resources: compute: accelerator_type: GPU_8xH100 accelerator_count: 8 + - task_key: train2 + environment_key: default + ai_runtime_task: + experiment: my-training-2 + code_source_path: ./src2 + deployments: + - command_path: src2/command.sh + compute: + accelerator_type: GPU_8xH100 + accelerator_count: 8 environments: - environment_key: default spec: diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt index df6b0f54acc..a76dc7382af 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -1,5 +1,5 @@ -=== deploy packages and uploads the local code source +=== deploy packages and uploads the local code sources >>> [CLI] bundle deploy Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... @@ -7,20 +7,31 @@ Deploying resources... Updating deployment state... Deployment complete! -=== the tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded +=== each task's tarball holds only synced files (both under the repo root .air_snapshots) >>> list_code_snapshot.py +# .air_snapshots/[SNAPSHOT].tar.gz +src2/command.sh +src2/train.py +# .air_snapshots/[SNAPSHOT].tar.gz src/.gitignore src/command.sh src/data/model.bin src/train.py -=== code_source_path points into the bundle (.air_snapshots), command_path rewritten, deps carried on environments spec +=== both code_source_paths point into the bundle .air_snapshots, command_paths rewritten, deps on environments spec >>> print_requests.py --sort --del-body raw_body //.air_snapshots/ //jobs/create { "method": "POST", - "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", + "q": { + "overwrite": "true" + } +} +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", "q": { "overwrite": "true" } @@ -54,7 +65,7 @@ src/train.py "tasks": [ { "ai_runtime_task": { - "code_source_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", + "code_source_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", "deployments": [ { "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh", @@ -68,11 +79,36 @@ src/train.py }, "environment_key": "default", "task_key": "train" + }, + { + "ai_runtime_task": { + "code_source_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src2/command.sh", + "compute": { + "accelerator_count": 8, + "accelerator_type": "GPU_8xH100" + } + } + ], + "experiment": "my-training-2" + }, + "environment_key": "default", + "task_key": "train2" } ] } } +=== re-deploying unchanged code is a no-op + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + === editing a file changes the snapshot hash (content-addressed name changes) >>> [CLI] bundle deploy @@ -84,8 +120,19 @@ Deployment complete! >>> print_requests.py --sort --del-body raw_body //.air_snapshots/ { "method": "POST", - "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", "q": { "overwrite": "true" } } + +=== destroy removes the deployed bundle (including the synced snapshots) + +>>> [CLI] bundle destroy --auto-approve +The following resources will be deleted: + delete resources.jobs.train + +All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default + +Deleting files... +Destroy complete! diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script index cd704f28119..574b0869cb4 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/script +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -1,21 +1,29 @@ -# A local directory code_source_path is packaged into a content-addressed tarball -# written into the bundle (.air_snapshots/), uploaded by normal bundle file sync, -# and code_source_path/command_path rewritten. Pip deps ride on the job's +# Two AI Runtime tasks each with a local directory code_source_path. Each is +# packaged into a content-addressed tarball written into the bundle +# (.air_snapshots/, at the repo root) and uploaded by normal bundle file sync; +# code_source_path/command_path are rewritten. Pip deps ride on the job's # environments[].spec.dependencies (no requirements.yaml is synthesized). -title "deploy packages and uploads the local code source\n" +title "deploy packages and uploads the local code sources\n" trace $CLI bundle deploy -title "the tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded\n" +title "each task's tarball holds only synced files (both under the repo root .air_snapshots)\n" trace list_code_snapshot.py -title "code_source_path points into the bundle (.air_snapshots), command_path rewritten, deps carried on environments spec\n" +title "both code_source_paths point into the bundle .air_snapshots, command_paths rewritten, deps on environments spec\n" # --del-body raw_body drops the binary tarball upload body (kept readable). Filters # use a leading // so Git Bash on Windows does not path-convert them. --keep is not # passed, so print_requests.py consumes out.requests.txt. trace print_requests.py --sort --del-body raw_body '//.air_snapshots/' '//jobs/create' +title "re-deploying unchanged code is a no-op\n" +trace $CLI bundle deploy + title "editing a file changes the snapshot hash (content-addressed name changes)\n" update_file.py src/train.py 'print("training")' 'print("training v2")' trace $CLI bundle deploy trace print_requests.py --sort --del-body raw_body '//.air_snapshots/' + +title "destroy removes the deployed bundle (including the synced snapshots)\n" +trace $CLI bundle destroy --auto-approve +rm out.requests.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src2/command.sh b/acceptance/bundle/ai_runtime_task/local_code_source/src2/command.sh new file mode 100644 index 00000000000..7ce3949767c --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src2/command.sh @@ -0,0 +1,2 @@ +cd $CODE_SOURCE_PATH +python train.py diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src2/train.py b/acceptance/bundle/ai_runtime_task/local_code_source/src2/train.py new file mode 100644 index 00000000000..4ff3f662781 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src2/train.py @@ -0,0 +1 @@ +print("train2") diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml index 6e4639a532a..0aa3f208009 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml +++ b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml @@ -6,7 +6,7 @@ Ignore = [ # The archive is content-addressed: _.tar.gz. The hash is stable # given the committed inputs, but collapse it to a token so the test does not pin a -# specific digest. +# specific digest. Matches both src_ and src2_ archives. [[Repls]] -Old = 'src_[0-9a-f]{16}\.tar\.gz' -New = 'src_[SNAPSHOT].tar.gz' +Old = '(src2?)_[0-9a-f]{16}\.tar\.gz' +New = '$1_[SNAPSHOT].tar.gz' diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index 6c901a347f3..5c523ee39dd 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -66,17 +66,17 @@ type codeSource struct { value string } -func PackageAndUpload() bundle.Mutator { - return &packageAndUpload{} +func PackageCodeSource() bundle.Mutator { + return &packageCodeSource{} } -type packageAndUpload struct{} +type packageCodeSource struct{} -func (m *packageAndUpload) Name() string { - return "aicode.PackageAndUpload" +func (m *packageCodeSource) Name() string { + return "aicode.PackageCodeSource" } -func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { +func (m *packageCodeSource) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { sources, diags := collectLocalCodeSources(b) if diags.HasError() { return diags diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go index 69e9d6300f2..dd9fbf0b970 100644 --- a/bundle/config/mutator/aicode/validate.go +++ b/bundle/config/mutator/aicode/validate.go @@ -3,6 +3,7 @@ package aicode import ( "context" "fmt" + "os" "path/filepath" "github.com/databricks/cli/bundle" @@ -11,6 +12,7 @@ import ( "github.com/databricks/cli/libs/diag" "github.com/databricks/cli/libs/dyn" "github.com/databricks/databricks-sdk-go/service/jobs" + ignore "github.com/sabhiram/go-gitignore" ) // Validate checks AI Runtime tasks that reference a local code_source_path so @@ -31,11 +33,21 @@ func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics jobsPath := dyn.NewPath(dyn.Key("resources"), dyn.Key("jobs")) + // packagesCode records whether any task actually has a local code_source_path + // this mutator will package. The bundle-level guards below (which reject configs + // that would drop the generated snapshot from sync) only matter in that case, so + // they're gated on it to avoid spurious errors on unrelated bundles. + packagesCode := false + for name, job := range b.Config.Resources.Jobs { jobPath := jobsPath.Append(dyn.Key(name)) for i, task := range job.Tasks { taskPath := jobPath.Append(dyn.Key("tasks"), dyn.Index(i)) + if task.AiRuntimeTask != nil && task.AiRuntimeTask.CodeSourcePath != "" && + libraries.IsLocalPath(task.AiRuntimeTask.CodeSourcePath) { + packagesCode = true + } // A local code_source_path under a for_each_task is not packaged by this // mutator (aicode collects only direct tasks). Reject it rather than let a @@ -63,9 +75,65 @@ func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics } } + if packagesCode { + diags = diags.Extend(validateSnapshotDir(b)) + } + + return diags +} + +// validateSnapshotDir guards the reserved snapshot directory. The mutator writes +// generated code archives into bundle.AiCodeSnapshotDir (overlaid on the sync root) +// and relies on file sync uploading them. Two config situations would silently break +// that — surface both at validate time: +// +// - A real file/dir at that path in the bundle would collide with the overlay +// (sync would carry the user's entry, not the generated archive). +// - A sync.exclude pattern matching that path removes the archive from the sync +// set (exclude is applied after include, so it wins over the force-include), +// leaving the deployed job pointing at an un-uploaded path. +func validateSnapshotDir(b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + syncExcludePath := dyn.NewPath(dyn.Key("sync"), dyn.Key("exclude")) + + // A user-owned file or directory at the reserved path collides with the overlay. + local := filepath.Join(b.SyncRootPath, bundle.AiCodeSnapshotDir) + if _, err := os.Stat(local); err == nil { + diags = diags.Append(diag.Diagnostic{ + Severity: diag.Error, + Summary: fmt.Sprintf("%q is reserved for AI Runtime code snapshots and must not exist in the bundle", bundle.AiCodeSnapshotDir), + Detail: "Remove it; the deploy generates code archives under this path.", + }) + } + + // A sync.exclude matching the reserved path would drop the generated archive from + // the upload (exclude wins over the force-include). + if matchesSnapshotDir(b.Config.Sync.Exclude) { + diags = diags.Append(diag.Diagnostic{ + Severity: diag.Error, + Summary: fmt.Sprintf("sync.exclude must not match %q, which holds the AI Runtime code snapshot", bundle.AiCodeSnapshotDir), + Detail: "Remove the pattern that excludes it; otherwise the deployed job's code_source_path would not be uploaded.", + Locations: b.Config.GetLocations(syncExcludePath.String()), + Paths: []dyn.Path{syncExcludePath}, + }) + } + return diags } +// matchesSnapshotDir reports whether any sync.exclude pattern would remove a file +// under the reserved snapshot directory, using the same gitignore-style matcher the +// sync engine applies to exclude patterns (see libs/fileset). +func matchesSnapshotDir(exclude []string) bool { + if len(exclude) == 0 { + return false + } + matcher := ignore.CompileIgnoreLines(exclude...) + // A representative archive path; the mutator names archives + // /_.tar.gz. + return matcher.MatchesPath(bundle.AiCodeSnapshotDir + "/probe.tar.gz") +} + func (v *validate) validateTask(b *bundle.Bundle, gitSource *jobs.GitSource, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { // Only local code_source_path values are packaged at deploy; remote values // are used as-is and need no validation here. diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go index 1839e8c227d..8ab48060af1 100644 --- a/bundle/config/mutator/aicode/validate_test.go +++ b/bundle/config/mutator/aicode/validate_test.go @@ -136,3 +136,46 @@ func TestValidateForEachTaskCodeSourceRejected(t *testing.T) { require.Len(t, diags, 1) assert.Contains(t, diags[0].Summary, "for_each_task") } + +// A sync.exclude pattern matching the reserved snapshot dir would drop the +// generated archive from the upload (exclude wins over include), so it is rejected. +func TestValidateSyncExcludeMatchingSnapshotDir(t *testing.T) { + for _, pattern := range []string{".air_snapshots/*", "**/*.tar.gz", ".air_snapshots"} { + b := bundleForValidate(t, "src", nil) + mkCodeDir(t, b, "src") + b.Config.Sync.Exclude = []string{pattern} + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1, "pattern %q should be rejected", pattern) + assert.Contains(t, diags[0].Summary, "sync.exclude") + } +} + +// An unrelated sync.exclude is fine — the guard only fires for patterns that would +// filter the snapshot dir. +func TestValidateSyncExcludeUnrelatedIsAllowed(t *testing.T) { + b := bundleForValidate(t, "src", nil) + mkCodeDir(t, b, "src") + b.Config.Sync.Exclude = []string{"*.log", "build/**"} + assert.Empty(t, Validate().Apply(t.Context(), b)) +} + +// A real file/dir at the reserved snapshot path collides with the overlay, so it is +// rejected. +func TestValidateReservedSnapshotDirCollision(t *testing.T) { + b := bundleForValidate(t, "src", nil) + mkCodeDir(t, b, "src") + require.NoError(t, os.MkdirAll(filepath.Join(b.SyncRootPath, ".air_snapshots"), 0o700)) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "reserved") +} + +// The snapshot-dir guards only fire when a task actually packages a local +// code_source — an unrelated bundle with a stray .air_snapshots is not this +// mutator's concern. +func TestValidateSnapshotGuardsSkippedWithoutLocalCodeSource(t *testing.T) { + b := bundleForValidate(t, "/Volumes/main/default/code/x.tar.gz", nil) + require.NoError(t, os.MkdirAll(filepath.Join(b.SyncRootPath, ".air_snapshots"), 0o700)) + b.Config.Sync.Exclude = []string{".air_snapshots/*"} + assert.Empty(t, Validate().Apply(t.Context(), b)) +} diff --git a/bundle/phases/build.go b/bundle/phases/build.go index febfb66d669..580a18f7ab6 100644 --- a/bundle/phases/build.go +++ b/bundle/phases/build.go @@ -33,7 +33,7 @@ func Build(ctx context.Context, b *bundle.Bundle) LibLocationMap { // rewrite the field to the synced workspace path. No requirements.yaml is // synthesized: the runtime installs pip deps from the job's serverless // environment (environments[].spec.dependencies) directly. - aicode.PackageAndUpload(), + aicode.PackageCodeSource(), mutator.ResolveVariableReferencesWithoutResources( "artifacts", From 91e10693120c8c1e59ca95a0e761619ab06d7671 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 17:10:27 +0000 Subject: [PATCH 12/16] acceptance/ai_runtime: literal no-op plan + empty code_source coverage Close two review gaps: - Assert a literal `bundle plan` no-op ("0 to add, 0 to change, 0 to delete") after deploy, instead of a no-op re-deploy proxy. - Add an empty_code_source acceptance test: a code_source dir whose contents are all filtered out (src/.gitignore of "*") makes `bundle deploy` fail with the "no files to package" error, covering packageOne's empty-source unhappy path end-to-end (the pipeline needs a workspace client, so this is exercised via acceptance rather than a mocked unit test). Co-authored-by: Isaac --- .../empty_code_source/databricks.yml | 22 +++++++++++++++++++ .../empty_code_source/out.test.toml | 3 +++ .../empty_code_source/output.txt | 6 +++++ .../ai_runtime_task/empty_code_source/script | 4 ++++ .../empty_code_source/src/.gitignore | 1 + .../empty_code_source/src/command.sh | 2 ++ .../empty_code_source/src/train.py | 1 + .../local_code_source/output.txt | 9 +++----- .../ai_runtime_task/local_code_source/script | 4 ++-- 9 files changed, 44 insertions(+), 8 deletions(-) create mode 100644 acceptance/bundle/ai_runtime_task/empty_code_source/databricks.yml create mode 100644 acceptance/bundle/ai_runtime_task/empty_code_source/out.test.toml create mode 100644 acceptance/bundle/ai_runtime_task/empty_code_source/output.txt create mode 100644 acceptance/bundle/ai_runtime_task/empty_code_source/script create mode 100644 acceptance/bundle/ai_runtime_task/empty_code_source/src/.gitignore create mode 100644 acceptance/bundle/ai_runtime_task/empty_code_source/src/command.sh create mode 100644 acceptance/bundle/ai_runtime_task/empty_code_source/src/train.py diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/empty_code_source/databricks.yml new file mode 100644 index 00000000000..8dd3b66fc95 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/databricks.yml @@ -0,0 +1,22 @@ +bundle: + name: ai-runtime-empty + +resources: + jobs: + train: + name: "[${bundle.target}] AI Runtime training" + tasks: + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source_path: ./src + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_8xH100 + accelerator_count: 8 + environments: + - environment_key: default + spec: + environment_version: "5" diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/out.test.toml b/acceptance/bundle/ai_runtime_task/empty_code_source/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/output.txt b/acceptance/bundle/ai_runtime_task/empty_code_source/output.txt new file mode 100644 index 00000000000..6ad352ed70b --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/output.txt @@ -0,0 +1,6 @@ + +>>> errcode [CLI] bundle deploy +Error: code_source_path "./src" has no files to package (all excluded by .gitignore or sync.exclude, or the directory is empty) + + +Exit code: 1 diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/script b/acceptance/bundle/ai_runtime_task/empty_code_source/script new file mode 100644 index 00000000000..3705b196b87 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/script @@ -0,0 +1,4 @@ +# When code_source_path resolves to a directory whose contents are all filtered out +# (here: a src/.gitignore of "*"), there is nothing to package. Deploy must fail +# with an actionable message rather than shipping an empty code archive. +trace errcode $CLI bundle deploy diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/src/.gitignore b/acceptance/bundle/ai_runtime_task/empty_code_source/src/.gitignore new file mode 100644 index 00000000000..72e8ffc0db8 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/src/.gitignore @@ -0,0 +1 @@ +* diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/src/command.sh b/acceptance/bundle/ai_runtime_task/empty_code_source/src/command.sh new file mode 100644 index 00000000000..7ce3949767c --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/src/command.sh @@ -0,0 +1,2 @@ +cd $CODE_SOURCE_PATH +python train.py diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/src/train.py b/acceptance/bundle/ai_runtime_task/empty_code_source/src/train.py new file mode 100644 index 00000000000..399ce99aa9f --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/src/train.py @@ -0,0 +1 @@ +print("x") diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt index a76dc7382af..4fb56dfe113 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -101,13 +101,10 @@ src/train.py } } -=== re-deploying unchanged code is a no-op +=== re-planning unchanged code is a no-op (no changes) ->>> [CLI] bundle deploy -Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... -Deploying resources... -Updating deployment state... -Deployment complete! +>>> [CLI] bundle plan +Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged === editing a file changes the snapshot hash (content-addressed name changes) diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script index 574b0869cb4..7d4e68d09aa 100644 --- a/acceptance/bundle/ai_runtime_task/local_code_source/script +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -16,8 +16,8 @@ title "both code_source_paths point into the bundle .air_snapshots, command_path # passed, so print_requests.py consumes out.requests.txt. trace print_requests.py --sort --del-body raw_body '//.air_snapshots/' '//jobs/create' -title "re-deploying unchanged code is a no-op\n" -trace $CLI bundle deploy +title "re-planning unchanged code is a no-op (no changes)\n" +trace $CLI bundle plan title "editing a file changes the snapshot hash (content-addressed name changes)\n" update_file.py src/train.py 'print("training")' 'print("training v2")' From 80d10b0222faa3fc10263b09f459c9ba22fbca31 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 17:23:18 +0000 Subject: [PATCH 13/16] aicode: scope the .air_snapshots force-include to AI Runtime bundles Only force-sync .air_snapshots for bundles that actually package an AI Runtime code_source, rather than for every bundle. The aicode mutator sets a new Bundle.HasAiRuntimeCodeSnapshot flag when it packages one, and GetSyncIncludePatterns gates the force-include on it. Keeps the generic sync include list from carrying a feature-specific entry for unrelated bundles. Also tighten the comments added across this feature to be concise. Co-authored-by: Isaac --- bundle/bundle.go | 18 +++++-- bundle/bundle_test.go | 23 +++++++++ .../config/mutator/aicode/package_upload.go | 50 +++++++------------ bundle/config/mutator/aicode/validate.go | 15 +++--- libs/vfs/overlay.go | 14 +++--- 5 files changed, 66 insertions(+), 54 deletions(-) diff --git a/bundle/bundle.go b/bundle/bundle.go index 46c3d3b2cfe..942a2ba9dde 100644 --- a/bundle/bundle.go +++ b/bundle/bundle.go @@ -36,10 +36,9 @@ import ( const internalFolder = ".internal" -// AiCodeSnapshotDir is the sync-root-relative directory (dot-prefixed so it does not -// clash with user files) that the aicode mutator writes AI Runtime code snapshots -// into. It is force-included in sync (see GetSyncIncludePatterns) so user ignore -// rules cannot filter the deployed job's code_source_path archives out. +// AiCodeSnapshotDir is the sync-relative dir the aicode mutator writes AI Runtime +// code snapshots into. Force-included in sync (see GetSyncIncludePatterns) so user +// ignore rules can't filter the deployed job's code_source_path archives out. const AiCodeSnapshotDir = ".air_snapshots" // Filename where resources are stored for DATABRICKS_BUNDLE_ENGINE=direct @@ -176,6 +175,12 @@ type Bundle struct { // comparison with remote state, but local file validation would incorrectly fail. SkipLocalFileValidation bool + // HasAiRuntimeCodeSnapshot is set by the aicode.PackageCodeSource build-phase + // mutator when it packages a local AI Runtime code_source into the bundle's + // snapshot dir. GetSyncIncludePatterns reads it to force-sync that dir only for + // bundles that actually use the feature, rather than for every bundle. + HasAiRuntimeCodeSnapshot bool + // Tagging is used to normalize tag keys and values. // The implementation depends on the cloud being targeted. Tagging tags.Cloud @@ -368,7 +373,10 @@ func (b *Bundle) GetSyncIncludePatterns(ctx context.Context) ([]string, error) { // Force-sync generated AI Runtime code snapshots so a user ignore rule (e.g. // "*.tar.gz" in .gitignore) can't filter them out — the deployed job's // code_source_path points at these archives (see bundle/config/mutator/aicode). - includes = append(includes, AiCodeSnapshotDir+"/*") + // Scoped to bundles that actually package one, so it's not a global include. + if b.HasAiRuntimeCodeSnapshot { + includes = append(includes, AiCodeSnapshotDir+"/*") + } return includes, nil } diff --git a/bundle/bundle_test.go b/bundle/bundle_test.go index 9bd667afe62..09c3636758f 100644 --- a/bundle/bundle_test.go +++ b/bundle/bundle_test.go @@ -71,6 +71,29 @@ func TestBundleLocalStateDir(t *testing.T) { assert.Equal(t, filepath.Join(projectDir, ".databricks", "bundle", "default"), cacheDir) } +func TestGetSyncIncludePatternsScopesSnapshotDir(t *testing.T) { + ctx := t.Context() + projectDir := t.TempDir() + f, err := os.Create(filepath.Join(projectDir, "databricks.yml")) + require.NoError(t, err) + f.Close() + + b, err := Load(ctx, projectDir) + require.NoError(t, err) + b.Config.Bundle.Target = "default" + + // Without an AI Runtime code snapshot, the dir is not force-included. + includes, err := b.GetSyncIncludePatterns(ctx) + require.NoError(t, err) + assert.NotContains(t, includes, AiCodeSnapshotDir+"/*") + + // Once the aicode mutator sets the flag, it is. + b.HasAiRuntimeCodeSnapshot = true + includes, err = b.GetSyncIncludePatterns(ctx) + require.NoError(t, err) + assert.Contains(t, includes, AiCodeSnapshotDir+"/*") +} + func TestBundleLocalStateDirOverride(t *testing.T) { ctx := t.Context() projectDir := t.TempDir() diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index 5c523ee39dd..0fc8cddae65 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -1,30 +1,17 @@ -// Package aicode packages a local code directory referenced by an AI Runtime -// task's code_source_path so the deployed job runs against the packaged code. +// Package aicode packages a local directory referenced by an AI Runtime task's +// code_source_path into a content-addressed tarball inside the bundle, and rewrites +// code_source_path to the workspace path that tarball occupies once synced. Remote +// values are left untouched. // -// The SDK jobs.AiRuntimeTask.code_source_path field expects a workspace path to an -// uploaded code archive; its doc comment states that the CLI is responsible for -// packaging the user's local code directory into that archive. This mutator -// implements that contract for DABs: when a user points code_source_path at a local -// directory, it packages the directory into a reproducible tarball (.git and -// gitignored files excluded) written into the bundle's sync tree, and rewrites -// code_source_path to the workspace path the tarball will occupy once synced. The -// tarball is uploaded by the normal bundle file sync during the deploy phase — this -// mutator performs no workspace writes itself, so it is safe to run in the build -// phase (before `bundle plan`). Values that are already remote are left untouched. +// The archive is overlaid on the sync tree and uploaded by normal bundle file sync +// in the deploy phase; the mutator performs no workspace writes, so it is safe in +// the build phase (which runs before `bundle plan`). Living in the bundle means +// `bundle destroy` cleans it, and the content-addressed name lets incremental sync +// skip re-uploading unchanged code. // -// Placing the archive in the bundle (rather than a separate cache) means -// `bundle destroy` removes it like any other bundle file, and bundle file sync is -// incremental so an unchanged archive is not re-uploaded. The archive is -// content-addressed: its name embeds the SHA-256 of the (reproducible) tarball, so -// unchanged code keeps the same name and the same synced path across deploys (see -// snapshot_package.go). -// -// Note for reviewers: code_source_path could alternatively be translated to its -// synced workspace path by mutator.TranslatePaths (like command_path is). That -// runs in the initialize phase, which also runs on `bundle validate`, so the -// archive would have to be materialized there too — writing files during validate. -// Keeping materialization in the build phase (deploy-only) and computing the synced -// path here avoids that. +// Not done via mutator.TranslatePaths (which handles command_path): that runs in +// initialize, which also runs on `bundle validate`, so the archive would be +// materialized during validate. Build phase is deploy-only. package aicode import ( @@ -109,6 +96,10 @@ func (m *packageCodeSource) Apply(ctx context.Context, b *bundle.Bundle) diag.Di // them like real files, but they never touch the user's working tree. b.SyncRoot = vfs.Overlay(b.SyncRoot, overlayFiles) + // Signal GetSyncIncludePatterns to force-sync the snapshot dir for this bundle, + // so a user ignore rule can't filter the archives out of the upload set. + b.HasAiRuntimeCodeSnapshot = true + err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { for _, cs := range sources { remote := remotePaths[cs.configPath.String()] @@ -127,13 +118,8 @@ func (m *packageCodeSource) Apply(ctx context.Context, b *bundle.Bundle) diag.Di return diags } -// snapshotSubdir is the bundle-local directory (sync-relative) that code snapshots -// live under. It is a dedicated folder — not the user's source tree — so a snapshot -// is never nested inside the directory it snapshots, and the archives are grouped in -// one predictable place. The archives are overlaid onto the sync root in memory, so -// this directory is synced to the workspace (and removed by `bundle destroy`) but is -// never materialized in the user's working tree. bundle.GetSyncIncludePatterns -// force-includes it so a user ignore rule can't filter the archives out of sync. +// snapshotSubdir is the sync-relative dir the archives are placed under (dedicated +// so a snapshot is never nested in the dir it snapshots). See bundle.AiCodeSnapshotDir. const snapshotSubdir = bundle.AiCodeSnapshotDir // packageOne packages the local directory for a single code source into a diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go index dd9fbf0b970..a3de250b5b3 100644 --- a/bundle/config/mutator/aicode/validate.go +++ b/bundle/config/mutator/aicode/validate.go @@ -82,16 +82,13 @@ func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics return diags } -// validateSnapshotDir guards the reserved snapshot directory. The mutator writes -// generated code archives into bundle.AiCodeSnapshotDir (overlaid on the sync root) -// and relies on file sync uploading them. Two config situations would silently break -// that — surface both at validate time: +// validateSnapshotDir rejects two configs that would silently drop the generated +// code archive from sync (leaving the job pointing at an un-uploaded path): // -// - A real file/dir at that path in the bundle would collide with the overlay -// (sync would carry the user's entry, not the generated archive). -// - A sync.exclude pattern matching that path removes the archive from the sync -// set (exclude is applied after include, so it wins over the force-include), -// leaving the deployed job pointing at an un-uploaded path. +// - A real file/dir at bundle.AiCodeSnapshotDir collides with the overlay (sync +// carries the user's entry, not the archive). +// - A sync.exclude matching that path removes the archive (exclude is applied +// after include, so it beats the force-include). func validateSnapshotDir(b *bundle.Bundle) diag.Diagnostics { var diags diag.Diagnostics syncExcludePath := dyn.NewPath(dyn.Key("sync"), dyn.Key("exclude")) diff --git a/libs/vfs/overlay.go b/libs/vfs/overlay.go index c2d240ce7bd..5dbdc013781 100644 --- a/libs/vfs/overlay.go +++ b/libs/vfs/overlay.go @@ -8,15 +8,13 @@ import ( "time" ) -// Overlay returns a Path that behaves like base but additionally serves a set of -// in-memory files. It lets a caller inject generated content into a filesystem -// tree (e.g. the bundle sync root) without writing it to disk: the overlaid files -// participate in Open/Stat/ReadDir/ReadFile and in fs.WalkDir walks, but the real -// working tree is left untouched. +// Overlay returns a Path that behaves like base but also serves a set of in-memory +// files, letting a caller inject generated content into a tree (e.g. the bundle sync +// root) without writing to disk. Overlaid files participate in Open/Stat/ReadDir/ +// ReadFile and fs.WalkDir; the real tree is untouched. // -// Overlaid names are slash-separated paths relative to the root (matching io/fs -// conventions). A name must not collide with a real entry in base; on collision the -// overlaid file wins for direct access, which callers should avoid. +// Names are slash-separated, relative to the root. A name must not collide with a +// real entry in base; on collision the overlaid file wins for direct access. func Overlay(base Path, files map[string][]byte) Path { // Copy so later mutations of the caller's map don't leak in. overlay := make(map[string][]byte, len(files)) From 2d06955af7c39618559878debdc83d5c1b9c1ca4 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 17:26:53 +0000 Subject: [PATCH 14/16] aicode: fix changelog fragment (no requirements.yaml; in-bundle upload) The fragment described the pre-review behavior. requirements.yaml is no longer synthesized (deps come from the environments spec), and the archive is packaged into the bundle and uploaded by file sync, not written to a separate workspace cache dir. Co-authored-by: Isaac --- .nextchanges/bundles/ai-runtime-code-source-dir.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.nextchanges/bundles/ai-runtime-code-source-dir.md b/.nextchanges/bundles/ai-runtime-code-source-dir.md index d3c5d0aac89..1ed59820a09 100644 --- a/.nextchanges/bundles/ai-runtime-code-source-dir.md +++ b/.nextchanges/bundles/ai-runtime-code-source-dir.md @@ -1 +1 @@ -An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy` (honoring `.gitignore` and `sync.include`/`sync.exclude`, content-addressed so unchanged code skips re-upload), uploaded to the user's workspace code snapshot directory, and rewritten to the uploaded path. A `requirements.yaml` derived from the job's serverless `environments` is written next to the task's `command_path`. Complements the existing pre-built-tarball path by letting users point at a source directory directly. +An `ai_runtime_task.code_source_path` that points at a local directory is now packaged into a tarball during `bundle deploy` (honoring `.gitignore` and `sync.include`/`sync.exclude`, content-addressed so unchanged code skips re-upload), placed in the bundle and uploaded by bundle file sync, and rewritten to the uploaded workspace path. Pip dependencies are taken from the job's serverless `environments` spec. Complements the existing pre-built-tarball path by letting users point at a source directory directly. From 07cd81b2aba60ed8af747f0cb54ede7c10320a7a Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 21:45:57 +0000 Subject: [PATCH 15/16] aicode: preserve the execute bit when packaging code MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The snapshot builder hardcoded 0o644, stripping the execute bit — a bundled helper the user invokes (e.g. a ./run.sh called from command.sh) would arrive non-executable and fail at runtime. Preserve the owner execute bit (0o755 when set, else 0o644) while still normalizing the rest of the mode. Deriving the mode from the file's own bits keeps the archive reproducible. Co-authored-by: Isaac --- .../config/mutator/aicode/snapshot_package.go | 15 +++++--- .../mutator/aicode/snapshot_package_test.go | 37 +++++++++++++++++++ 2 files changed, 47 insertions(+), 5 deletions(-) diff --git a/bundle/config/mutator/aicode/snapshot_package.go b/bundle/config/mutator/aicode/snapshot_package.go index ee545df1d6c..0bd4771c758 100644 --- a/bundle/config/mutator/aicode/snapshot_package.go +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -100,15 +100,20 @@ func addFileToArchive(tw *tar.Writer, syncRoot vfs.Path, relBase string, f files return nil } + // Preserve the owner execute bit so a bundled helper the user invokes (e.g. a + // ./run.sh called from command.sh) stays executable, but normalize the rest to a + // canonical mode. Deriving the mode from the file's own bits keeps the archive + // reproducible (same input file -> same mode); mtime is zeroed for the same reason. + mode := int64(0o644) + if info.Mode().Perm()&0o100 != 0 { + mode = 0o755 + } hdr := &tar.Header{ Typeflag: tar.TypeReg, Name: path.Join(prefix, rel), Size: info.Size(), - // Normalize permissions and zero the mtime so the archive is reproducible - // across machines. The runtime invokes code via an interpreter, not by - // executing files from the snapshot, so execute bits are not preserved. - Mode: 0o644, - ModTime: tarEpoch, + Mode: mode, + ModTime: tarEpoch, } if err := tw.WriteHeader(hdr); err != nil { return fmt.Errorf("tar header for %s: %w", rel, err) diff --git a/bundle/config/mutator/aicode/snapshot_package_test.go b/bundle/config/mutator/aicode/snapshot_package_test.go index feb2e10c1e7..726cf28b01c 100644 --- a/bundle/config/mutator/aicode/snapshot_package_test.go +++ b/bundle/config/mutator/aicode/snapshot_package_test.go @@ -51,6 +51,43 @@ func tarEntries(t *testing.T, b []byte) map[string]string { return out } +// tarModes reads a gzipped tarball and returns entry name -> permission bits. +func tarModes(t *testing.T, b []byte) map[string]int64 { + t.Helper() + gzr, err := gzip.NewReader(bytes.NewReader(b)) + require.NoError(t, err) + tr := tar.NewReader(gzr) + out := map[string]int64{} + for { + hdr, err := tr.Next() + if err == io.EOF { + break + } + require.NoError(t, err) + out[hdr.Name] = hdr.Mode & 0o777 + } + return out +} + +// An executable file keeps the execute bit (0o755) so a bundled helper the user +// invokes still runs; a non-executable file is normalized to 0o644. +func TestBuildCodeSnapshotPreservesExecuteBit(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "run.sh"), []byte("#!/bin/sh\n"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "train.py"), []byte("x"), 0o644)) + root := vfs.MustNew(dir) + files, err := fileset.New(root).Files() + require.NoError(t, err) + + var buf bytes.Buffer + _, err = buildCodeSnapshot(root, ".", files, "code", &buf) + require.NoError(t, err) + + modes := tarModes(t, buf.Bytes()) + assert.Equal(t, int64(0o755), modes["code/run.sh"], "executable helper must keep its execute bit") + assert.Equal(t, int64(0o644), modes["code/train.py"], "non-executable file is normalized to 0o644") +} + func TestBuildCodeSnapshotPrefixesEntries(t *testing.T) { root, files := writeTree(t, map[string]string{ "train.py": "print('train')", From bfcf976b7d7360ac4e352cb9a15c1da23dbdfa62 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 3 Aug 2026 22:03:07 +0000 Subject: [PATCH 16/16] aicode: skip execute-bit test on Windows Windows file modes don't carry a Unix execute bit, so writing 0o755 doesn't make the file report executable and the archived-mode assertion fails there. The bit only matters on the Unix hosts that run the workload; skip the test on Windows. The production path is already Windows-safe (files report non-exec -> 0o644). Co-authored-by: Isaac --- bundle/config/mutator/aicode/snapshot_package_test.go | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/bundle/config/mutator/aicode/snapshot_package_test.go b/bundle/config/mutator/aicode/snapshot_package_test.go index 726cf28b01c..5342924fd58 100644 --- a/bundle/config/mutator/aicode/snapshot_package_test.go +++ b/bundle/config/mutator/aicode/snapshot_package_test.go @@ -7,6 +7,7 @@ import ( "io" "os" "path/filepath" + "runtime" "testing" "github.com/databricks/cli/libs/fileset" @@ -72,6 +73,12 @@ func tarModes(t *testing.T, b []byte) map[string]int64 { // An executable file keeps the execute bit (0o755) so a bundled helper the user // invokes still runs; a non-executable file is normalized to 0o644. func TestBuildCodeSnapshotPreservesExecuteBit(t *testing.T) { + if runtime.GOOS == "windows" { + // Windows file modes don't carry a Unix execute bit, so there's nothing to + // preserve; files simply archive as 0o644. The bit only matters on the + // Unix hosts that run the deployed workload. + t.Skip("execute bit is a Unix concept; not represented on Windows") + } dir := t.TempDir() require.NoError(t, os.WriteFile(filepath.Join(dir, "run.sh"), []byte("#!/bin/sh\n"), 0o755)) require.NoError(t, os.WriteFile(filepath.Join(dir, "train.py"), []byte("x"), 0o644))