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..4590fb6be8f --- /dev/null +++ b/.nextchanges/bundles/ai-runtime-code-source-dir.md @@ -0,0 +1 @@ +For jobs where `ai_runtime_task.code_source_path` is a relative path to a local directory, the directory is now packaged into a tarball (honoring `.gitignore` and `sync.include`/`sync.exclude`), uploaded during deployment, and `code_source_path` is rewritten to the uploaded workspace path. diff --git a/acceptance/bin/list_code_snapshot.py b/acceptance/bin/list_code_snapshot.py new file mode 100755 index 00000000000..b8f1263cae7 --- /dev/null +++ b/acceptance/bin/list_code_snapshot.py @@ -0,0 +1,72 @@ +#!/usr/bin/env python3 +""" +List the entries of each AI Runtime code snapshot tarball uploaded during deploy. + +Reads out.requests.txt, takes every ai_runtime_task.code_source_path from the +jobs/create request, exports each workspace archive via the CLI, and prints its +sorted tar entries (grouped per archive). Used to assert which local files each +snapshot includes (gitignore / sync rules), across one or more tasks. +""" + +import gzip +import io +import os +import subprocess +import sys +import tarfile + +from print_requests import read_json_many + + +def code_source_paths(requests): + """Every task's code_source_path from the jobs/create request(s).""" + result = [] + for req in requests: + body = req.get("body") + if isinstance(body, dict) and req.get("path", "").endswith("/jobs/create"): + for task in body.get("tasks", []): + art = task.get("ai_runtime_task") + if art and art.get("code_source_path"): + result.append(art["code_source_path"]) + return result + + +def print_entries(cli, env, remote): + local = "code_snapshot.tar.gz" + subprocess.run( + [cli, "workspace", "export", remote, "--format", "AUTO", "--file", local], + check=True, + env=env, + ) + with open(local, "rb") as f: + data = gzip.decompress(f.read()) + os.remove(local) + + # Print the archive's sync-relative name (hash tokenized by test.toml repls) so + # multi-archive output is legible. + print(f"# {remote.split('/files/', 1)[-1]}") + with tarfile.open(fileobj=io.BytesIO(data)) as tar: + for name in sorted(tar.getnames()): + print(name) + + +def main(): + with open("out.requests.txt") as f: + requests = read_json_many(f.read()) + + paths = code_source_paths(requests) + if not paths: + sys.exit("no jobs/create request with code_source_path in out.requests.txt") + + cli = os.environ["CLI"] + # MSYS_NO_PATHCONV stops Git Bash on Windows from rewriting the /Workspace path. + env = {**os.environ, "MSYS_NO_PATHCONV": "1"} + + # code_source_path is an absolute workspace path (/Workspace/Users/.../files/...). + # Sort so multi-task output is deterministic. + for remote in sorted(paths): + print_entries(cli, env, remote) + + +if __name__ == "__main__": + main() diff --git a/acceptance/bin/print_requests.py b/acceptance/bin/print_requests.py index 82720ac5b23..0e8a74c0759 100755 --- a/acceptance/bin/print_requests.py +++ b/acceptance/bin/print_requests.py @@ -201,10 +201,20 @@ def main(): "body fields that diverge between deployment engines, e.g. identity fields the " "terraform provider serializes into the body but the direct engine sends as query params.", ) + parser.add_argument( + "--del-field", + action="append", + default=[], + metavar="FIELDS", + help="Comma-separated top-level request fields to delete (repeatable). Unlike " + "--del-body, which edits the parsed JSON body, this drops a field of the request " + "record itself, e.g. raw_body for a binary upload payload.", + ) parser.add_argument("--fname", default="out.requests.txt") args = parser.parse_args() del_body_fields = [field for group in args.del_body for field in group.split(",")] + del_fields = [field for group in args.del_field for field in group.split(",")] test_tmp_dir = os.environ.get("TEST_TMP_DIR") if test_tmp_dir: @@ -229,6 +239,8 @@ def main(): if isinstance(body, dict): for field in del_body_fields: body.pop(field, None) + for field in del_fields: + 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/empty_code_source/assets/model.bin b/acceptance/bundle/ai_runtime_task/empty_code_source/assets/model.bin new file mode 100644 index 00000000000..e69de29bb2d diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/empty_code_source/databricks.yml new file mode 100644 index 00000000000..562924ddc06 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/databricks.yml @@ -0,0 +1,29 @@ +bundle: + name: ai-runtime-empty + +# An unrelated force-include elsewhere in the bundle. sync.include is added to the +# file list regardless of the scoped walk, so it must not make an all-filtered code +# directory look non-empty. +sync: + include: + - assets/*.bin + +resources: + jobs: + train: + name: "[${bundle.target}] AI Runtime training" + tasks: + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source_path: ./src + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_8xH100 + accelerator_count: 8 + environments: + - environment_key: default + spec: + environment_version: "5" diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/out.test.toml b/acceptance/bundle/ai_runtime_task/empty_code_source/out.test.toml new file mode 100644 index 00000000000..f784a183258 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/out.test.toml @@ -0,0 +1,3 @@ +Local = true +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/output.txt b/acceptance/bundle/ai_runtime_task/empty_code_source/output.txt new file mode 100644 index 00000000000..dbfdc962645 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/output.txt @@ -0,0 +1,4 @@ + +>>> [CLI] bundle deploy +Error: code_source_path "./src" has no files to package (all excluded by .gitignore or sync.exclude, or the directory is empty) + diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/script b/acceptance/bundle/ai_runtime_task/empty_code_source/script new file mode 100644 index 00000000000..6c9858c4368 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/script @@ -0,0 +1,5 @@ +# When code_source_path resolves to a directory whose contents are all filtered out +# (here: a src/.gitignore of "*"), there is nothing to package. Deploy must fail +# with an actionable message rather than shipping an empty code archive. +# A sync.include elsewhere in the bundle must not defeat that guard (see databricks.yml). +musterr trace $CLI bundle deploy diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/src/.gitignore b/acceptance/bundle/ai_runtime_task/empty_code_source/src/.gitignore new file mode 100644 index 00000000000..72e8ffc0db8 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/src/.gitignore @@ -0,0 +1 @@ +* diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/src/command.sh b/acceptance/bundle/ai_runtime_task/empty_code_source/src/command.sh new file mode 100644 index 00000000000..7ce3949767c --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/src/command.sh @@ -0,0 +1,2 @@ +cd $CODE_SOURCE_PATH +python train.py diff --git a/acceptance/bundle/ai_runtime_task/empty_code_source/src/train.py b/acceptance/bundle/ai_runtime_task/empty_code_source/src/train.py new file mode 100644 index 00000000000..399ce99aa9f --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/empty_code_source/src/train.py @@ -0,0 +1 @@ +print("x") diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml new file mode 100644 index 00000000000..4a1cc0476c0 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml @@ -0,0 +1,43 @@ +bundle: + name: ai-runtime-test + +sync: + # *.log excluded; data/*.bin force-included despite .gitignore. + exclude: + - "**/*.log" + include: + - src/data/*.bin + +resources: + jobs: + train: + name: "[${bundle.target}] AI Runtime training" + tasks: + # Two AI Runtime tasks with distinct local code dirs: each is packaged into + # its own content-addressed tarball, both under the repo root's .air_snapshots. + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source_path: ./src + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_8xH100 + accelerator_count: 8 + - task_key: train2 + environment_key: default + ai_runtime_task: + experiment: my-training-2 + code_source_path: ./src2 + deployments: + - command_path: src2/command.sh + compute: + accelerator_type: GPU_8xH100 + accelerator_count: 8 + environments: + - environment_key: default + spec: + 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..06cdf841c55 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -0,0 +1,135 @@ + +=== deploy packages and uploads the local code sources + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +=== each task's tarball holds only synced files (both under the repo root .air_snapshots) + +>>> list_code_snapshot.py +# .air_snapshots/[SNAPSHOT].tar.gz +src2/command.sh +src2/train.py +# .air_snapshots/[SNAPSHOT].tar.gz +src/.gitignore +src/command.sh +src/data/model.bin +src/train.py + +=== both code_source_paths point into the bundle .air_snapshots, command_paths rewritten, deps on environments spec + +>>> print_requests.py --sort --del-field raw_body //.air_snapshots/ //jobs/create +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", + "q": { + "overwrite": "true" + } +} +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", + "q": { + "overwrite": "true" + } +} +{ + "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": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh", + "compute": { + "accelerator_count": 8, + "accelerator_type": "GPU_8xH100" + } + } + ], + "experiment": "my-training" + }, + "environment_key": "default", + "task_key": "train" + }, + { + "ai_runtime_task": { + "code_source_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src2/command.sh", + "compute": { + "accelerator_count": 8, + "accelerator_type": "GPU_8xH100" + } + } + ], + "experiment": "my-training-2" + }, + "environment_key": "default", + "task_key": "train2" + } + ] + } +} + +=== re-planning unchanged code is a no-op (no changes) + +>>> [CLI] bundle plan +Plan: 0 to add, 0 to change, 0 to delete, 1 unchanged + +=== editing a file changes the snapshot hash (content-addressed name changes) + +>>> [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-field raw_body //.air_snapshots/ +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/[SNAPSHOT].tar.gz", + "q": { + "overwrite": "true" + } +} + +=== destroy removes the deployed bundle (including the synced snapshots) + +>>> [CLI] bundle destroy --auto-approve +The following resources will be deleted: + delete resources.jobs.train + +All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default + +Deleting files... +Destroy complete! diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script new file mode 100644 index 00000000000..fe5e707ce01 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -0,0 +1,29 @@ +# Two AI Runtime tasks each with a local directory code_source_path. Each is +# packaged into a content-addressed tarball written into the bundle +# (.air_snapshots/, at the repo root) and uploaded by normal bundle file sync; +# code_source_path/command_path are rewritten. Pip deps ride on the job's +# environments[].spec.dependencies (no requirements.yaml is synthesized). + +title "deploy packages and uploads the local code sources\n" +trace $CLI bundle deploy + +title "each task's tarball holds only synced files (both under the repo root .air_snapshots)\n" +trace list_code_snapshot.py + +title "both code_source_paths point into the bundle .air_snapshots, command_paths rewritten, deps on environments spec\n" +# --del-field raw_body drops the binary tarball upload payload (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-field raw_body '//.air_snapshots/' '//jobs/create' + +title "re-planning unchanged code is a no-op (no changes)\n" +trace $CLI bundle plan + +title "editing a file changes the snapshot hash (content-addressed name changes)\n" +update_file.py src/train.py 'print("training")' 'print("training v2")' +trace $CLI bundle deploy +trace print_requests.py --sort --del-field raw_body '//.air_snapshots/' + +title "destroy removes the deployed bundle (including the synced snapshots)\n" +trace $CLI bundle destroy --auto-approve +rm out.requests.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore new file mode 100644 index 00000000000..54553edaac1 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore @@ -0,0 +1,2 @@ +ignored_by_git.txt +data/ 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/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/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/src2/command.sh b/acceptance/bundle/ai_runtime_task/local_code_source/src2/command.sh new file mode 100644 index 00000000000..7ce3949767c --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src2/command.sh @@ -0,0 +1,2 @@ +cd $CODE_SOURCE_PATH +python train.py diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src2/train.py b/acceptance/bundle/ai_runtime_task/local_code_source/src2/train.py new file mode 100644 index 00000000000..4ff3f662781 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src2/train.py @@ -0,0 +1 @@ +print("train2") diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml new file mode 100644 index 00000000000..0aa3f208009 --- /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. Matches both src_ and src2_ archives. +[[Repls]] +Old = '(src2?)_[0-9a-f]{16}\.tar\.gz' +New = '$1_[SNAPSHOT].tar.gz' diff --git a/acceptance/bundle/artifacts/ai_runtime_code_source/databricks.yml b/acceptance/bundle/artifacts/ai_runtime_code_source/databricks.yml index c2108f4988e..96783455893 100644 --- a/acceptance/bundle/artifacts/ai_runtime_code_source/databricks.yml +++ b/acceptance/bundle/artifacts/ai_runtime_code_source/databricks.yml @@ -5,11 +5,15 @@ workspace: artifact_path: /foo/bar/artifacts # The tarball is delivered as an artifact; don't also sync it as a bundle file. +# The archive glob is natural here and must not be mistaken for an exclude of the +# generated code-snapshot dir: this task's code_source_path is a local FILE, so no +# snapshot is packaged and the snapshot-dir guards must stay silent. sync: exclude: - build/** - dist/** - code.tgz + - "**/*.tar.gz" # The AI Runtime task ships user code as a tarball built locally by an artifact. artifacts: diff --git a/bundle/bundle.go b/bundle/bundle.go index 94f55389b26..942a2ba9dde 100644 --- a/bundle/bundle.go +++ b/bundle/bundle.go @@ -36,6 +36,11 @@ import ( const internalFolder = ".internal" +// AiCodeSnapshotDir is the sync-relative dir the aicode mutator writes AI Runtime +// code snapshots into. Force-included in sync (see GetSyncIncludePatterns) so user +// ignore rules can't filter the deployed job's code_source_path archives out. +const AiCodeSnapshotDir = ".air_snapshots" + // Filename where resources are stored for DATABRICKS_BUNDLE_ENGINE=direct const resourcesFilename = "resources.json" @@ -170,6 +175,12 @@ type Bundle struct { // comparison with remote state, but local file validation would incorrectly fail. SkipLocalFileValidation bool + // HasAiRuntimeCodeSnapshot is set by the aicode.PackageCodeSource build-phase + // mutator when it packages a local AI Runtime code_source into the bundle's + // snapshot dir. GetSyncIncludePatterns reads it to force-sync that dir only for + // bundles that actually use the feature, rather than for every bundle. + HasAiRuntimeCodeSnapshot bool + // Tagging is used to normalize tag keys and values. // The implementation depends on the cloud being targeted. Tagging tags.Cloud @@ -358,7 +369,15 @@ func (b *Bundle) GetSyncIncludePatterns(ctx context.Context) ([]string, error) { if err != nil { return nil, err } - return append(b.Config.Sync.Include, filepath.ToSlash(filepath.Join(internalDirRel, "*.*"))), nil + includes := append(b.Config.Sync.Include, filepath.ToSlash(filepath.Join(internalDirRel, "*.*"))) + // Force-sync generated AI Runtime code snapshots so a user ignore rule (e.g. + // "*.tar.gz" in .gitignore) can't filter them out — the deployed job's + // code_source_path points at these archives (see bundle/config/mutator/aicode). + // Scoped to bundles that actually package one, so it's not a global include. + if b.HasAiRuntimeCodeSnapshot { + includes = append(includes, AiCodeSnapshotDir+"/*") + } + return includes, nil } // AuthEnv returns a map with environment variables and their values diff --git a/bundle/bundle_test.go b/bundle/bundle_test.go index 9bd667afe62..09c3636758f 100644 --- a/bundle/bundle_test.go +++ b/bundle/bundle_test.go @@ -71,6 +71,29 @@ func TestBundleLocalStateDir(t *testing.T) { assert.Equal(t, filepath.Join(projectDir, ".databricks", "bundle", "default"), cacheDir) } +func TestGetSyncIncludePatternsScopesSnapshotDir(t *testing.T) { + ctx := t.Context() + projectDir := t.TempDir() + f, err := os.Create(filepath.Join(projectDir, "databricks.yml")) + require.NoError(t, err) + f.Close() + + b, err := Load(ctx, projectDir) + require.NoError(t, err) + b.Config.Bundle.Target = "default" + + // Without an AI Runtime code snapshot, the dir is not force-included. + includes, err := b.GetSyncIncludePatterns(ctx) + require.NoError(t, err) + assert.NotContains(t, includes, AiCodeSnapshotDir+"/*") + + // Once the aicode mutator sets the flag, it is. + b.HasAiRuntimeCodeSnapshot = true + includes, err = b.GetSyncIncludePatterns(ctx) + require.NoError(t, err) + assert.Contains(t, includes, AiCodeSnapshotDir+"/*") +} + func TestBundleLocalStateDirOverride(t *testing.T) { ctx := t.Context() projectDir := t.TempDir() diff --git a/bundle/config/mutator/aicode/package_code_source.go b/bundle/config/mutator/aicode/package_code_source.go new file mode 100644 index 00000000000..b2f78a9aebd --- /dev/null +++ b/bundle/config/mutator/aicode/package_code_source.go @@ -0,0 +1,261 @@ +// Package aicode packages a local directory referenced by an AI Runtime task's +// code_source_path into a content-addressed tarball inside the bundle, and rewrites +// code_source_path to the workspace path that tarball occupies once synced. Remote +// values are left untouched. +// +// The archive is overlaid on the sync tree and uploaded by normal bundle file sync +// in the deploy phase; the mutator performs no workspace writes, so it is safe in +// the build phase (which runs before `bundle plan`). Living in the bundle means +// `bundle destroy` cleans it, and the content-addressed name lets incremental sync +// skip re-uploading unchanged code. +// +// Not done via mutator.TranslatePaths (which handles command_path): that runs in +// initialize, which also runs on `bundle validate`, so the archive would be +// materialized during validate. Build phase is deploy-only. +package aicode + +import ( + "bytes" + "context" + "errors" + "fmt" + "io/fs" + "os" + "path" + "path/filepath" + "slices" + "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/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 +// code_source_path. It matches a direct task only — the same scope aicode.Validate +// operates on. ai_runtime_task nested under a for_each_task is not a supported +// combination yet (Validate rejects it); when it is, both should gain it together. +var codeSourcePattern = dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + 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 PackageCodeSource() bundle.Mutator { + return &packageCodeSource{} +} + +type packageCodeSource struct{} + +func (m *packageCodeSource) Name() string { + return "aicode.PackageCodeSource" +} + +func (m *packageCodeSource) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + sources, diags := collectLocalCodeSources(b) + if diags.HasError() { + return diags + } + if len(sources) == 0 { + return diags + } + + // remotePaths maps each config location to the synced workspace path its archive + // 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 { + relArchive, archive, err := packageOne(ctx, b, cs) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + return diags + } + 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. + syncRoot, err := vfs.Overlay(b.SyncRoot, overlayFiles) + if err != nil { + return diags.Extend(diag.FromErr(err)) + } + b.SyncRoot = syncRoot + + // Signal GetSyncIncludePatterns to force-sync the snapshot dir for this bundle, + // so a user ignore rule can't filter the archives out of the upload set. + b.HasAiRuntimeCodeSnapshot = true + + err = b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { + for _, cs := range sources { + remote := remotePaths[cs.configPath.String()] + 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) + } + } + return root, nil + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + + return diags +} + +// snapshotSubdir is the sync-relative dir the archives are placed under (dedicated +// so a snapshot is never nested in the dir it snapshots). See bundle.AiCodeSnapshotDir. +const snapshotSubdir = bundle.AiCodeSnapshotDir + +// packageOne packages the local directory for a single code source into a +// 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) + + // 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 "", 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 "", nil, fmt.Errorf("failed to list files for code_source_path %q: %w", cs.value, err) + } + // An empty file list means every file under the directory was filtered out + // (gitignore / sync.exclude) or the directory is empty. Packaging it would deploy + // a job with no code, so fail with an actionable message instead. + if len(files) == 0 { + return "", nil, fmt.Errorf("code_source_path %q has no files to package (all excluded by .gitignore or sync.exclude, or the directory is empty)", cs.value) + } + + // Build the archive in memory so its content hash can name the file; the hash is + // computed while gzipping, so this adds no extra pass over the files. + var buf bytes.Buffer + sha, err := buildCodeSnapshot(b.SyncRoot, relBase, files, dirName, &buf) + if err != nil { + return "", nil, fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) + } + // 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 +// 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 + } + all, err := fl.Files(ctx) + if err != nil { + return nil, err + } + + // sync.include is force-added regardless of the scoped walk, so the list can + // contain files outside the code directory. Keep only what is under relBase, or + // those strays make an all-filtered directory look non-empty and an empty archive + // ships. + if relBase == "." { + return all, nil + } + prefix := relBase + "/" + return slices.DeleteFunc(all, func(f fileset.File) bool { + return !strings.HasPrefix(f.Relative, prefix) + }), 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 + + 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)) + } + + 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/package_code_source_test.go b/bundle/config/mutator/aicode/package_code_source_test.go new file mode 100644 index 00000000000..fca33a0a7bf --- /dev/null +++ b/bundle/config/mutator/aicode/package_code_source_test.go @@ -0,0 +1,89 @@ +package aicode + +import ( + "os" + "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 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) + 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) + } +} + +// 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/snapshot_package.go b/bundle/config/mutator/aicode/snapshot_package.go new file mode 100644 index 00000000000..af1d5d583e0 --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -0,0 +1,125 @@ +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). +// Reproducible per platform, not across them (see addFileToArchive). +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 macOS archives carry no extra entries. +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 + } + + // Preserve the owner execute bit so a bundled helper stays executable, and + // normalize the rest to a canonical mode. Windows has no execute bit, so files are + // archived 0644 there and the archive hash differs from a Unix-built one. + mode := int64(0o644) + if info.Mode().Perm()&0o100 != 0 { + mode = 0o755 + } + hdr := &tar.Header{ + Typeflag: tar.TypeReg, + Name: path.Join(prefix, rel), + Size: info.Size(), + Mode: mode, + 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..5342924fd58 --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package_test.go @@ -0,0 +1,165 @@ +package aicode + +import ( + "archive/tar" + "bytes" + "compress/gzip" + "io" + "os" + "path/filepath" + "runtime" + "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 +} + +// tarModes reads a gzipped tarball and returns entry name -> permission bits. +func tarModes(t *testing.T, b []byte) map[string]int64 { + t.Helper() + gzr, err := gzip.NewReader(bytes.NewReader(b)) + require.NoError(t, err) + tr := tar.NewReader(gzr) + out := map[string]int64{} + for { + hdr, err := tr.Next() + if err == io.EOF { + break + } + require.NoError(t, err) + out[hdr.Name] = hdr.Mode & 0o777 + } + return out +} + +// An executable file keeps the execute bit (0o755) so a bundled helper the user +// invokes still runs; a non-executable file is normalized to 0o644. +func TestBuildCodeSnapshotPreservesExecuteBit(t *testing.T) { + if runtime.GOOS == "windows" { + // Windows file modes don't carry a Unix execute bit, so there's nothing to + // preserve; files simply archive as 0o644. The bit only matters on the + // Unix hosts that run the deployed workload. + t.Skip("execute bit is a Unix concept; not represented on Windows") + } + dir := t.TempDir() + require.NoError(t, os.WriteFile(filepath.Join(dir, "run.sh"), []byte("#!/bin/sh\n"), 0o755)) + require.NoError(t, os.WriteFile(filepath.Join(dir, "train.py"), []byte("x"), 0o644)) + root := vfs.MustNew(dir) + files, err := fileset.New(root).Files() + require.NoError(t, err) + + var buf bytes.Buffer + _, err = buildCodeSnapshot(root, ".", files, "code", &buf) + require.NoError(t, err) + + modes := tarModes(t, buf.Bytes()) + assert.Equal(t, int64(0o755), modes["code/run.sh"], "executable helper must keep its execute bit") + assert.Equal(t, int64(0o644), modes["code/train.py"], "non-executable file is normalized to 0o644") +} + +func TestBuildCodeSnapshotPrefixesEntries(t *testing.T) { + root, files := writeTree(t, map[string]string{ + "train.py": "print('train')", + "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..492203bc9f4 --- /dev/null +++ b/bundle/config/mutator/aicode/validate.go @@ -0,0 +1,218 @@ +package aicode + +import ( + "context" + "fmt" + "os" + "path/filepath" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/libraries" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/databricks-sdk-go/service/jobs" + ignore "github.com/sabhiram/go-gitignore" +) + +// Validate checks AI Runtime tasks that reference a local code_source_path so +// 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")) + + // packagesCode records whether any task actually has a code_source_path this + // mutator will package. The bundle-level guards below (which reject configs that + // would drop the generated snapshot from sync) only matter in that case, so + // they're gated on it to avoid spurious errors on unrelated bundles. + packagesCode := false + + for name, job := range b.Config.Resources.Jobs { + jobPath := jobsPath.Append(dyn.Key(name)) + + for i, task := range job.Tasks { + taskPath := jobPath.Append(dyn.Key("tasks"), dyn.Index(i)) + if task.AiRuntimeTask != nil && v.packagesLocalDir(b, task.AiRuntimeTask.CodeSourcePath) { + packagesCode = true + } + + // A local code_source_path under a for_each_task is not packaged by this + // mutator (aicode collects only direct tasks). Reject it rather than let a + // nested ai_runtime_task deploy an un-packaged local path. + if task.ForEachTask != nil && task.ForEachTask.Task.AiRuntimeTask != nil { + nestedCode := task.ForEachTask.Task.AiRuntimeTask.CodeSourcePath + if nestedCode != "" && libraries.IsLocalPath(nestedCode) { + p := taskPath.Append(dyn.Key("for_each_task"), dyn.Key("task"), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) + diags = diags.Append(diag.Diagnostic{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path is not supported inside a for_each_task", + Detail: "Set code_source_path to a workspace or volume path", + Locations: b.Config.GetLocations(p.String()), + Paths: []dyn.Path{p}, + }) + } + } + + if task.AiRuntimeTask == nil { + continue + } + codePath := taskPath.Append(dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) + diags = diags.Extend(v.validateTask(b, job.GitSource, task.AiRuntimeTask.CodeSourcePath, codePath)) + } + } + + if packagesCode { + diags = diags.Extend(validateSnapshotDir(b)) + } + + return diags +} + +// packagesLocalDir reports whether codeSourcePath is one this mutator packages: a +// local path that is an existing directory. A local *file* (a pre-built tarball from +// an `artifacts` block) is uploaded by the artifact path instead, so it packages no +// snapshot and the snapshot-directory guards must not apply to it. +func (v *validate) packagesLocalDir(b *bundle.Bundle, codeSourcePath string) bool { + if codeSourcePath == "" || !libraries.IsLocalPath(codeSourcePath) { + return false + } + // A stat error is reported by validateTask; treat it as "not a directory" here. + isDir, err := isExistingDir(filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath))) + return err == nil && isDir +} + +// validateSnapshotDir rejects two configs that would silently drop the generated +// code archive from sync (leaving the job pointing at an un-uploaded path): +// +// - A real file/dir at bundle.AiCodeSnapshotDir collides with the overlay (sync +// carries the user's entry, not the archive). +// - A sync.exclude matching that path removes the archive (exclude is applied +// after include, so it beats the force-include). +func validateSnapshotDir(b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + syncExcludePath := dyn.NewPath(dyn.Key("sync"), dyn.Key("exclude")) + + // A user-owned file or directory at the reserved path collides with the overlay. + local := filepath.Join(b.SyncRootPath, bundle.AiCodeSnapshotDir) + if _, err := os.Stat(local); err == nil { + diags = diags.Append(diag.Diagnostic{ + Severity: diag.Error, + Summary: fmt.Sprintf("%q is reserved for AI Runtime code snapshots and must not exist in the bundle", bundle.AiCodeSnapshotDir), + Detail: "Remove it; the deploy generates code archives under this path.", + }) + } + + // A sync.exclude matching the reserved path would drop the generated archive from + // the upload (exclude wins over the force-include). + if matchesSnapshotDir(b.Config.Sync.Exclude) { + diags = diags.Append(diag.Diagnostic{ + Severity: diag.Error, + Summary: fmt.Sprintf("sync.exclude must not match %q, which holds the AI Runtime code snapshot", bundle.AiCodeSnapshotDir), + Detail: "Remove the pattern that excludes it; otherwise the deployed job's code_source_path would not be uploaded.", + Locations: b.Config.GetLocations(syncExcludePath.String()), + Paths: []dyn.Path{syncExcludePath}, + }) + } + + return diags +} + +// matchesSnapshotDir reports whether any sync.exclude pattern would remove a file +// under the reserved snapshot directory, using the same gitignore-style matcher the +// sync engine applies to exclude patterns (see libs/fileset). +func matchesSnapshotDir(exclude []string) bool { + if len(exclude) == 0 { + return false + } + matcher := ignore.CompileIgnoreLines(exclude...) + // A representative archive path; the mutator names archives + // /_.tar.gz. + return matcher.MatchesPath(bundle.AiCodeSnapshotDir + "/probe.tar.gz") +} + +func (v *validate) validateTask(b *bundle.Bundle, gitSource *jobs.GitSource, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { + // Only local code_source_path values are packaged at deploy; remote values + // are used as-is and need no validation here. + if codeSourcePath == "" || !libraries.IsLocalPath(codeSourcePath) { + return nil + } + + locations := b.Config.GetLocations(codePath.String()) + reject := func(summary, detail string) diag.Diagnostics { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: summary, + Detail: detail, + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + // The packaged directory must live inside the bundle sync root: it is uploaded as + // part of the bundle, so a path escaping the root (e.g. "../shared") can't be + // synced. Reject it here with a clear message rather than letting it fail later as + // an opaque io/fs "invalid argument" when the file list is built. + if rel, err := filepath.Rel(b.SyncRootPath, filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath))); err != nil || !filepath.IsLocal(rel) { + return reject( + fmt.Sprintf("code_source_path %q is outside the bundle root", codeSourcePath), + "code_source_path must point at a directory inside the bundle, or at a workspace or volume path", + ) + } + + // This mutator packages a local *directory*. A local path that is not an existing + // directory is left alone so it flows through the standard artifact-upload path: + // a pre-built tarball delivered via an `artifacts` block is produced during the + // 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. 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)) + isDir, err := isExistingDir(localDir) + if err != nil { + return reject(fmt.Sprintf("failed to inspect code_source_path %q: %v", codeSourcePath, err), "") + } + if !isDir { + return nil + } + + // The deploy engine retrieves task files from git when git_source is set, so + // packaging a local directory would be silently ignored. Reject the combination. + if gitSource != nil { + return reject( + "ai_runtime_task with a local code_source_path cannot be combined with git_source", + "Remove git_source or set code_source_path to a workspace or volume path", + ) + } + + // Immutable-folder deployments upload a single content-addressed snapshot and + // do not support the per-task code packaging this mutator performs. + if b.IsImmutableFolder() { + return reject("ai_runtime_task with a local code_source_path is not supported with experimental.immutable_folder", "") + } + + // Source-linked deployment runs jobs against the source files in place and does + // not copy them to the workspace file path (files.Upload is a no-op). This + // mutator relies on file sync uploading the packaged snapshot, so the two are + // incompatible: reject rather than deploy a job pointing at an un-uploaded path. + if config.IsExplicitlyEnabled(b.Config.Presets.SourceLinkedDeployment) { + return reject( + "ai_runtime_task with a local code_source_path is not supported with source-linked deployment", + "Disable source-linked deployment, or set code_source_path to a workspace or volume path", + ) + } + + return nil +} diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go new file mode 100644 index 00000000000..8ab48060af1 --- /dev/null +++ b/bundle/config/mutator/aicode/validate_test.go @@ -0,0 +1,181 @@ +package aicode + +import ( + "os" + "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 +} + +// 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) + 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") +} + +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) +} + +// A code_source_path escaping the bundle sync root is rejected with a clear +// message (it can't be synced), rather than failing later as an opaque io/fs error. +func TestValidateCodeSourceOutsideBundleRoot(t *testing.T) { + b := bundleForValidate(t, "../shared", nil) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "outside the bundle root") +} + +// Source-linked deployment doesn't copy files to the workspace file path, so the +// packaged snapshot would never be uploaded; the combination is rejected. +func TestValidateSourceLinkedConflict(t *testing.T) { + b := bundleForValidate(t, "src", nil) + mkCodeDir(t, b, "src") + enabled := true + b.Config.Presets.SourceLinkedDeployment = &enabled + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "source-linked deployment") +} + +// A local code_source_path nested under a for_each_task is not packaged by the +// mutator, so it is rejected rather than silently skipped. +func TestValidateForEachTaskCodeSourceRejected(t *testing.T) { + dir := t.TempDir() + b := &bundle.Bundle{ + BundleRootPath: dir, + SyncRootPath: dir, + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "train": { + JobSettings: jobs.JobSettings{ + Tasks: []jobs.Task{ + { + TaskKey: "fanout", + ForEachTask: &jobs.ForEachTask{ + Task: jobs.Task{ + AiRuntimeTask: &jobs.AiRuntimeTask{CodeSourcePath: "src"}, + }, + }, + }, + }, + }, + }, + }, + }, + }, + } + bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "for_each_task") +} + +// A sync.exclude pattern matching the reserved snapshot dir would drop the +// generated archive from the upload (exclude wins over include), so it is rejected. +func TestValidateSyncExcludeMatchingSnapshotDir(t *testing.T) { + for _, pattern := range []string{".air_snapshots/*", "**/*.tar.gz", ".air_snapshots"} { + b := bundleForValidate(t, "src", nil) + mkCodeDir(t, b, "src") + b.Config.Sync.Exclude = []string{pattern} + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1, "pattern %q should be rejected", pattern) + assert.Contains(t, diags[0].Summary, "sync.exclude") + } +} + +// An unrelated sync.exclude is fine — the guard only fires for patterns that would +// filter the snapshot dir. +func TestValidateSyncExcludeUnrelatedIsAllowed(t *testing.T) { + b := bundleForValidate(t, "src", nil) + mkCodeDir(t, b, "src") + b.Config.Sync.Exclude = []string{"*.log", "build/**"} + assert.Empty(t, Validate().Apply(t.Context(), b)) +} + +// A real file/dir at the reserved snapshot path collides with the overlay, so it is +// rejected. +func TestValidateReservedSnapshotDirCollision(t *testing.T) { + b := bundleForValidate(t, "src", nil) + mkCodeDir(t, b, "src") + require.NoError(t, os.MkdirAll(filepath.Join(b.SyncRootPath, ".air_snapshots"), 0o700)) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "reserved") +} + +// The snapshot-dir guards only fire when a task actually packages a local +// code_source — an unrelated bundle with a stray .air_snapshots is not this +// mutator's concern. +func TestValidateSnapshotGuardsSkippedWithoutLocalCodeSource(t *testing.T) { + b := bundleForValidate(t, "/Volumes/main/default/code/x.tar.gz", nil) + require.NoError(t, os.MkdirAll(filepath.Join(b.SyncRootPath, ".air_snapshots"), 0o700)) + b.Config.Sync.Exclude = []string{".air_snapshots/*"} + assert.Empty(t, Validate().Apply(t.Context(), b)) +} diff --git a/bundle/phases/build.go b/bundle/phases/build.go index db376e07e28..580a18f7ab6 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,13 @@ 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 content-addressed tarball overlaid on the sync root, and + // rewrite the field to the synced workspace path. No requirements.yaml is + // synthesized: the runtime installs pip deps from the job's serverless + // environment (environments[].spec.dependencies) directly. + aicode.PackageCodeSource(), + 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 diff --git a/libs/vfs/overlay.go b/libs/vfs/overlay.go new file mode 100644 index 00000000000..5d22c37919a --- /dev/null +++ b/libs/vfs/overlay.go @@ -0,0 +1,178 @@ +package vfs + +import ( + "bytes" + "fmt" + "io" + "io/fs" + "path" + "time" +) + +// Overlay returns a Path that behaves like base but also serves a set of in-memory +// files, letting a caller inject generated content into a tree (e.g. the bundle sync +// root) without writing to disk. Overlaid files participate in Open/Stat/ReadDir/ +// ReadFile and fs.WalkDir; the real tree is untouched. +// +// Names are slash-separated and relative to the root, per fs.ValidPath: an absolute +// name, an empty name, or one escaping the root with ".." is rejected. A name must +// not collide with a real entry in base; on collision the overlaid file wins for +// direct access. +func Overlay(base Path, files map[string][]byte) (Path, error) { + // 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) + + for name, data := range files { + clean := path.Clean(name) + // Reject anything that isn't a file rooted in the tree. An absolute path would + // also make the ancestor walk below spin forever: path.Dir("/") is "/", never + // ".". Clean("") is "." (the root itself), which is a directory, not a file. + if clean == "." || !fs.ValidPath(clean) { + return nil, fmt.Errorf("overlay: invalid file name %q: must be a relative slash-separated path inside the root", 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) + addOverlayChild(dirs, parent, path.Base(child), isFile) + child = parent + isFile = false + } + } + return &overlayPath{base: base, files: overlay, dirs: dirs}, nil +} + +// addOverlayChild records child under dir in the ancestor-directory index. +func addOverlayChild(dirs map[string]map[string]bool, dir, child 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[child] { + entries[child] = isFile + } +} + +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..df53cef4cda --- /dev/null +++ b/libs/vfs/overlay_test.go @@ -0,0 +1,95 @@ +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, err := Overlay(base, map[string][]byte{ + ".air_snapshots/src_abc.tar.gz": []byte("SNAPSHOT-BYTES"), + }) + require.NoError(t, err) + + // 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, err := Overlay(MustNew(dir), map[string][]byte{ + ".air_snapshots/src_abc.tar.gz": []byte("x"), + }) + require.NoError(t, err) + + 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) +} + +// Names outside the root are rejected rather than silently mis-registered. An +// absolute name additionally used to hang the ancestor walk (path.Dir("/") == "/"). +func TestOverlayRejectsNamesOutsideRoot(t *testing.T) { + for _, name := range []string{ + "/foo/bar.txt", + "../foo.txt", + "a/../../b", + "", + } { + t.Run(name, func(t *testing.T) { + _, err := Overlay(MustNew(t.TempDir()), map[string][]byte{name: []byte("x")}) + require.ErrorContains(t, err, "invalid file name") + }) + } +} + +// A relative name that stays inside the root is accepted, including one that only +// normalizes to an in-root path after cleaning. +func TestOverlayAcceptsInRootNames(t *testing.T) { + for _, name := range []string{"a.txt", "./b/c.txt", "d/../e.txt"} { + t.Run(name, func(t *testing.T) { + ov, err := Overlay(MustNew(t.TempDir()), map[string][]byte{name: []byte("x")}) + require.NoError(t, err) + got, err := ov.ReadFile(name) + require.NoError(t, err) + assert.Equal(t, "x", string(got)) + }) + } +}