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/bin/list_code_snapshot.py b/acceptance/bin/list_code_snapshot.py new file mode 100755 index 00000000000..088829a4456 --- /dev/null +++ b/acceptance/bin/list_code_snapshot.py @@ -0,0 +1,59 @@ +#!/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 an absolute workspace path (/Workspace/Users/.../files/...). +remote = code_source_path +cli = os.environ["CLI"] +local = "code_snapshot.tar.gz" +# MSYS_NO_PATHCONV stops Git Bash on Windows from rewriting the /Workspace path. +env = {**os.environ, "MSYS_NO_PATHCONV": "1"} +subprocess.run( + [cli, "workspace", "export", remote, "--format", "AUTO", "--file", local], + check=True, + env=env, +) +with open(local, "rb") as f: + data = gzip.decompress(f.read()) +os.remove(local) + +with tarfile.open(fileobj=io.BytesIO(data)) as tar: + for name in sorted(tar.getnames()): + print(name) 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 new file mode 100644 index 00000000000..97401944141 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml @@ -0,0 +1,31 @@ +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: + - 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..c9d0380d7f3 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -0,0 +1,105 @@ + +=== deploy packages and uploads 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 tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded + +>>> list_code_snapshot.py +src/.gitignore +src/command.sh +src/data/model.bin +src/train.py + +=== code_source_path points into the bundle (.air_snapshots), command_path is rewritten, requirements.yaml written + +>>> 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]/.bundle/ai-runtime-test/default/files/.air_snapshots/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": "/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", + "compute": { + "accelerator_count": 8, + "accelerator_type": "GPU_8xH100" + } + } + ], + "experiment": "my-training" + }, + "environment_key": "default", + "task_key": "train" + } + ] + } +} + +=== 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-body raw_body //.air_snapshots/ +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", + "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 new file mode 100644 index 00000000000..c417b896cc8 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -0,0 +1,21 @@ +# 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 + +title "the tarball contains only synced files: .gitignore'd excluded, sync.include kept, *.log excluded\n" +trace list_code_snapshot.py + +title "code_source_path points into the bundle (.air_snapshots), command_path is rewritten, requirements.yaml written\n" +# --del-body raw_body drops the binary tarball upload body (kept readable). Filters +# use a leading // so Git Bash on Windows does not path-convert them. --keep is not +# passed, so print_requests.py consumes out.requests.txt. +trace print_requests.py --sort --del-body raw_body '//.air_snapshots/' '//jobs/create' '//files/src/requirements.yaml' + +title "editing a file changes the snapshot hash (content-addressed name changes)\n" +update_file.py src/train.py 'print("training")' 'print("training v2")' +trace $CLI bundle deploy +trace print_requests.py --sort --del-body raw_body '//.air_snapshots/' 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/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..613fefbec12 --- /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 so the deployed job runs against the packaged code. +// +// 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. +// +// 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 ( + "bytes" + "context" + "errors" + "fmt" + "io/fs" + "os" + "path" + "path/filepath" + + "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 +// 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 { + 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{} + +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 + } + + // 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. + 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()] + 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 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 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) + } + + // 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 + } + return fl.Files(ctx) +} + +// 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_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go new file mode 100644 index 00000000000..fca33a0a7bf --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload_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/requirements.go b/bundle/config/mutator/aicode/requirements.go new file mode 100644 index 00000000000..8927500d78d --- /dev/null +++ b/bundle/config/mutator/aicode/requirements.go @@ -0,0 +1,141 @@ +package aicode + +import ( + "bytes" + "context" + "fmt" + "path" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go/service/compute" + "github.com/databricks/databricks-sdk-go/service/jobs" + "go.yaml.in/yaml/v3" +) + +// requirementsFileName is the file the AI Runtime entry script reads from the +// command_path directory to install the workload's pip dependencies. The server +// derives its path from command_path, so it must sit next to command.sh. +const requirementsFileName = "requirements.yaml" + +// requirementsSpec is the requirements.yaml shape the AI Runtime consumes: a +// client image version and a list of pip requirement lines. It mirrors what the +// Python air CLI synthesizes. +type requirementsSpec struct { + Version string `yaml:"version,omitempty"` + Dependencies []string `yaml:"dependencies,omitempty"` +} + +// SynthesizeRequirements uploads a requirements.yaml next to each AI Runtime task's +// command_path, derived from the job-level serverless environment referenced by the +// task's environment_key. The AI Runtime entry script reads this file (resolved +// from command_path's directory) to set up the workload environment; without it the +// run fails during setup. It runs at deploy time, after command_path has been +// translated to its absolute workspace path. +func SynthesizeRequirements() bundle.Mutator { + return &synthesizeRequirements{} +} + +type synthesizeRequirements struct{} + +func (m *synthesizeRequirements) Name() string { + return "aicode.SynthesizeRequirements" +} + +func (m *synthesizeRequirements) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + + for name, job := range b.Config.Resources.Jobs { + envs := environmentsByKey(job.Environments) + for i := range job.Tasks { + task := &job.Tasks[i] + if task.AiRuntimeTask == nil { + continue + } + if err := synthesizeForTask(ctx, b, name, task, envs); err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + } + } + + return diags +} + +// synthesizeForTask uploads requirements.yaml next to each deployment's +// command_path. It is a no-op when the task has no matching environment spec (the +// run can still supply deps another way, so this is not an error). Each deployment +// carries its own command_path and the entry script reads requirements.yaml from +// that directory, so every deployment's directory needs the file; directories are +// deduplicated so a shared command_path directory is written only once. +func synthesizeForTask(ctx context.Context, b *bundle.Bundle, jobName string, task *jobs.Task, envs map[string]*compute.Environment) error { + env := envs[task.EnvironmentKey] + if env == nil { + return nil + } + + content, err := renderRequirements(env) + if err != nil { + return err + } + + written := make(map[string]bool) + for _, dep := range task.AiRuntimeTask.Deployments { + if dep.CommandPath == "" { + continue + } + // command_path is an absolute workspace path (translated during initialize); + // the entry script reads requirements.yaml from its directory. + dir := path.Dir(dep.CommandPath) + if written[dir] { + continue + } + written[dir] = true + + client, err := filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), dir) + if err != nil { + return err + } + log.Debugf(ctx, "writing %s for job %s task %s to %s", requirementsFileName, jobName, task.TaskKey, dir) + if err := client.Write(ctx, requirementsFileName, bytes.NewReader(content), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return fmt.Errorf("failed to upload %s next to command_path %q: %w", requirementsFileName, dep.CommandPath, err) + } + } + return nil +} + +// renderRequirements builds the requirements.yaml content from a serverless +// environment spec: version from environment_version (or the legacy client field), +// dependencies from its pip dependency list. +func renderRequirements(env *compute.Environment) ([]byte, error) { + version := env.EnvironmentVersion + if version == "" { + version = env.Client + } + spec := requirementsSpec{ + Version: version, + Dependencies: env.Dependencies, + } + var buf bytes.Buffer + enc := yaml.NewEncoder(&buf) + enc.SetIndent(2) + if err := enc.Encode(spec); err != nil { + return nil, err + } + if err := enc.Close(); err != nil { + return nil, err + } + return buf.Bytes(), nil +} + +// environmentsByKey indexes a job's environments by their key for task lookup. +func environmentsByKey(envs []jobs.JobEnvironment) map[string]*compute.Environment { + out := make(map[string]*compute.Environment, len(envs)) + for i := range envs { + if envs[i].Spec != nil { + out[envs[i].EnvironmentKey] = envs[i].Spec + } + } + return out +} diff --git a/bundle/config/mutator/aicode/requirements_test.go b/bundle/config/mutator/aicode/requirements_test.go 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..1b2777dfee6 --- /dev/null +++ b/bundle/config/mutator/aicode/validate.go @@ -0,0 +1,103 @@ +package aicode + +import ( + "context" + "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 +// 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, task.AiRuntimeTask.CodeSourcePath, codePath)) + } + } + + return diags +} + +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 + } + + // 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 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 + } + + locations := b.Config.GetLocations(codePath.String()) + + // The deploy engine retrieves task files from git when git_source is set, so + // packaging a local directory would be silently ignored. Reject the combination. + if gitSource != nil { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path cannot be combined with git_source", + Detail: "Remove git_source or set code_source_path to a workspace or volume path", + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + // Immutable-folder deployments upload a single content-addressed snapshot and + // do not support the per-task code packaging this mutator performs. + if b.IsImmutableFolder() { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path is not supported with experimental.immutable_folder", + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + return 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..bddcc8f309a --- /dev/null +++ b/bundle/config/mutator/aicode/validate_test.go @@ -0,0 +1,82 @@ +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) +} 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 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) +}