From 470ac256338edd516b2c4eb2b173a01512eab688 Mon Sep 17 00:00:00 2001 From: "ben.hansen" Date: Fri, 17 Jul 2026 21:37:38 +0000 Subject: [PATCH 1/8] 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 2/8] 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 3/8] =?UTF-8?q?aicode:=20address=20review=20=E2=80=94=20dr?= =?UTF-8?q?op=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 4/8] 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 5/8] 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 6/8] 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 7/8] 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 8/8] 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) +}