From c25fc2512ad9982829504dbda7472b17f08fe3a3 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 20 Jul 2026 20:57:25 +0000 Subject: [PATCH 1/3] Add aicode mutator: package local-dir code_source_path at deploy AI Runtime tasks can point code_source_path at a local directory. The aicode mutator packages that directory into a reproducible, content- addressed tarball (.git + gitignored files excluded, sync globs honored), uploads it to the user's ~/.air/repo_snapshots, and rewrites the field to the remote path. The content-addressed name lets an unchanged directory skip re-upload across deploys. SynthesizeRequirements writes requirements.yaml next to command_path from the serverless environment (no-op when it has none); Validate surfaces local-dir misconfigurations at validate time. Based on the commit before the SDK v0.160.0 bump, which temporarily removed jobs.AiRuntimeTask.code_source_path (returns next week). The acceptance fixture runs the direct engine only until the field is back in the generated terraform provider schema. Co-authored-by: Isaac --- .../local_code_source/databricks.yml | 29 +++ .../local_code_source/out.test.toml | 3 + .../local_code_source/output.txt | 29 +++ .../ai_runtime_task/local_code_source/script | 30 +++ .../local_code_source/src/.gitignore | 1 + .../local_code_source/src/command.sh | 2 + .../local_code_source/src/train.py | 1 + .../local_code_source/test.toml | 17 ++ .../config/mutator/aicode/package_upload.go | 215 ++++++++++++++++++ .../mutator/aicode/package_upload_test.go | 74 ++++++ bundle/config/mutator/aicode/requirements.go | 137 +++++++++++ .../mutator/aicode/requirements_test.go | 36 +++ .../config/mutator/aicode/snapshot_package.go | 100 ++++++++ .../mutator/aicode/snapshot_package_test.go | 121 ++++++++++ bundle/config/mutator/aicode/validate.go | 97 ++++++++ bundle/config/mutator/aicode/validate_test.go | 64 ++++++ bundle/phases/build.go | 12 + bundle/phases/initialize.go | 7 + 18 files changed, 975 insertions(+) create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/output.txt create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/script create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/src/train.py create mode 100644 acceptance/bundle/ai_runtime_task/local_code_source/test.toml create mode 100644 bundle/config/mutator/aicode/package_upload.go create mode 100644 bundle/config/mutator/aicode/package_upload_test.go create mode 100644 bundle/config/mutator/aicode/requirements.go create mode 100644 bundle/config/mutator/aicode/requirements_test.go create mode 100644 bundle/config/mutator/aicode/snapshot_package.go create mode 100644 bundle/config/mutator/aicode/snapshot_package_test.go create mode 100644 bundle/config/mutator/aicode/validate.go create mode 100644 bundle/config/mutator/aicode/validate_test.go diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml new file mode 100644 index 00000000000..ff4a7b5dc1b --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/databricks.yml @@ -0,0 +1,29 @@ +bundle: + name: ai-runtime-test + +sync: + # Exclude a file from both bundle sync and the code snapshot. + exclude: + - "**/*.log" + +resources: + jobs: + train: + name: "[${bundle.target}] AI Runtime training" + tasks: + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source_path: ./src + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_8xH100 + accelerator_count: 8 + environments: + - environment_key: default + spec: + environment_version: "5" + dependencies: + - torch>=2.0.0 diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/out.test.toml new file mode 100644 index 00000000000..e90b6d5d1ba --- /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 = ["direct"] diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/output.txt b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt new file mode 100644 index 00000000000..33fb18eba49 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/output.txt @@ -0,0 +1,29 @@ + +=== deploy packages the local code source + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +=== the snapshot tarball is uploaded to the user's repo_snapshots directory +/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== code_source_path points at the uploaded archive (no /Workspace prefix) +/Users/[USERNAME]/.air/repo_snapshots/src/src_[SNAPSHOT].tar.gz + +=== command_path is translated to its absolute synced workspace path +/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/command.sh + +=== a requirements.yaml is written next to command_path +/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files/src/requirements.yaml + +=== re-deploying unchanged code is a cache hit: the tarball is not re-uploaded + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-test/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! +tarball uploads on redeploy: 0 diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/script b/acceptance/bundle/ai_runtime_task/local_code_source/script new file mode 100644 index 00000000000..d0a9e2ea965 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/script @@ -0,0 +1,30 @@ +# A local code_source_path is packaged into a content-addressed tarball +# (_.tar.gz) and uploaded to the user's repo_snapshots directory. The +# command_path is translated to its synced workspace path, and a requirements.yaml +# derived from the job's serverless environment is written beside it. + +title "deploy packages the local code source\n" +trace $CLI bundle deploy + +title "the snapshot tarball is uploaded to the user's repo_snapshots directory\n" +# The upload appears twice: WSFS import-file 404s until the parent dir exists, then +# the filer mkdirs and retries. Collapse to the distinct path. +jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u + +title "code_source_path points at the uploaded archive (no /Workspace prefix)\n" +jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.code_source_path' out.requests.txt + +title "command_path is translated to its absolute synced workspace path\n" +jq -r 'select(.path == "/api/2.2/jobs/create") | .body.tasks[0].ai_runtime_task.deployments[0].command_path' out.requests.txt + +title "a requirements.yaml is written next to command_path\n" +jq -r 'select(.path | test("/import-file/.*/files/src/requirements.yaml$")) | .path' out.requests.txt | sort -u + +rm out.requests.txt + +title "re-deploying unchanged code is a cache hit: the tarball is not re-uploaded\n" +trace $CLI bundle deploy +echo -n "tarball uploads on redeploy: " +jq -r 'select(.path | test("/import-file/.*/.air/repo_snapshots/src/src_")) | .path' out.requests.txt | sort -u | wc -l | tr -d ' ' + +rm out.requests.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore new file mode 100644 index 00000000000..7a79c3b103e --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/.gitignore @@ -0,0 +1 @@ +ignored_by_git.txt diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh b/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh new file mode 100644 index 00000000000..7ce3949767c --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/command.sh @@ -0,0 +1,2 @@ +cd $CODE_SOURCE_PATH +python train.py diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py b/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py new file mode 100644 index 00000000000..d8b062e5dfe --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/src/train.py @@ -0,0 +1 @@ +print("training") diff --git a/acceptance/bundle/ai_runtime_task/local_code_source/test.toml b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml new file mode 100644 index 00000000000..894a7d77b11 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/local_code_source/test.toml @@ -0,0 +1,17 @@ +RecordRequests = true + +# Direct engine only: the terraform engine needs code_source_path in the generated +# terraform provider schema, absent while the field is temporarily removed from the +# jobs API/provider (returns next week). Restore ["terraform", "direct"] then. +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] + +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..4a59792cefb --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload.go @@ -0,0 +1,215 @@ +// Package aicode packages a local code_source_path directory into a +// content-addressed tarball at deploy, uploads it to the user's workspace, and +// rewrites code_source_path to the remote archive. Already-remote values are left +// untouched. The archive name embeds its SHA-256, so unchanged code skips re-upload. +package aicode + +import ( + "bytes" + "context" + "errors" + "fmt" + "io/fs" + "path" + "path/filepath" + "strings" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/deploy/files" + "github.com/databricks/cli/bundle/libraries" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/filer" + "github.com/databricks/cli/libs/fileset" + "github.com/databricks/cli/libs/log" + libsync "github.com/databricks/cli/libs/sync" +) + +// codeSourcePatterns locate code_source_path on a direct task and under a for_each_task. +var codeSourcePatterns = []dyn.Pattern{ + dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"), + ), + dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + dyn.Key("for_each_task"), dyn.Key("task"), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path"), + ), +} + +// codeSource is a single local code_source_path occurrence to package. +type codeSource struct { + configPath dyn.Path + location dyn.Location + value string // raw code_source_path string as written in config +} + +func PackageAndUpload() bundle.Mutator { + return &packageAndUpload{} +} + +type packageAndUpload struct { + client filer.Filer // nil in normal use (a filer is built per code source); set only in tests +} + +func (m *packageAndUpload) Name() string { + return "aicode.PackageAndUpload" +} + +func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + sources, diags := collectLocalCodeSources(b) + if diags.HasError() { + return diags + } + if len(sources) == 0 { + return diags + } + + userDir, err := userWorkspaceHome(b) + if err != nil { + return diags.Extend(diag.FromErr(err)) + } + + // Upload all sources before rewriting any config, so an upload failure leaves + // the config untouched. + remotePaths := make(map[string]string, len(sources)) + for _, cs := range sources { + remote, err := m.packageOne(ctx, b, cs, userDir) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + return diags + } + remotePaths[cs.configPath.String()] = remote + } + + err = b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { + for _, cs := range sources { + remote := remotePaths[cs.configPath.String()] + root, err = dyn.SetByPath(root, cs.configPath, dyn.NewValue(remote, []dyn.Location{cs.location})) + if err != nil { + return root, fmt.Errorf("failed to update code_source_path %q to %q: %w", cs.value, remote, err) + } + } + return root, nil + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + + return diags +} + +// repoSnapshotsSubdir holds code snapshots under the user's home. Deliberately not +// /.internal, which artifacts.CleanUp() wipes each deploy (matches the Python air CLI). +const repoSnapshotsSubdir = ".air/repo_snapshots" + +// packageOne tarballs one code source, uploads it to repo_snapshots (skipping the +// upload when a content-identical archive is already there), and returns its remote path. +func (m *packageAndUpload) packageOne(ctx context.Context, b *bundle.Bundle, cs codeSource, userDir string) (string, error) { + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.value)) + dirName := filepath.Base(localDir) + + // relBase scopes the sync file list to the code dir and re-bases archive entries under it. + relBase, err := filepath.Rel(b.SyncRootPath, localDir) + if err != nil { + return "", fmt.Errorf("code_source_path %q: %w", cs.value, err) + } + relBase = filepath.ToSlash(relBase) + + files, err := codeSourceFiles(ctx, b, relBase) + if err != nil { + return "", fmt.Errorf("failed to list files for code_source_path %q: %w", cs.value, err) + } + + uploadPath := path.Join(userDir, repoSnapshotsSubdir, dirName) + client := m.client + if client == nil { + client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), uploadPath) + if err != nil { + return "", err + } + } + + // Build in memory so the content hash (computed while gzipping) can name the upload. + var buf bytes.Buffer + sha, err := buildCodeSnapshot(b.SyncRoot, relBase, files, dirName, &buf) + if err != nil { + return "", fmt.Errorf("failed to package code_source_path %q: %w", cs.value, err) + } + + archiveName := fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16]) + // The AI Runtime fetcher wants the legacy "/Users/..." form (no "/Workspace"), while + // the filer uploads via uploadPath; record the de-prefixed path on the task. + remotePath := strings.TrimPrefix(path.Join(uploadPath, archiveName), "/Workspace") + + // A matching name means identical content is already uploaded (the archive is reproducible). + if _, err := client.Stat(ctx, archiveName); err == nil { + log.Debugf(ctx, "code snapshot already present at %s, skipping upload", remotePath) + return remotePath, nil + } else if !errors.Is(err, fs.ErrNotExist) { + return "", fmt.Errorf("failed to check for existing code snapshot %q: %w", remotePath, err) + } + + if err := client.Write(ctx, archiveName, &buf, filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return "", fmt.Errorf("failed to upload code snapshot %q: %w", remotePath, err) + } + return remotePath, nil +} + +// codeSourceFiles lists the files under relBase (relative to the sync root) for the +// snapshot, filtered like bundle file sync: .gitignore-aware plus sync.include/exclude. +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 + } + fl, err := libsync.NewFileList(ctx, opts.WorktreeRoot, opts.LocalRoot, []string{relBase}, opts.Include, opts.Exclude) + if err != nil { + return nil, err + } + return fl.Files(ctx) +} + +// userWorkspaceHome returns /Workspace/Users/, the root for code snapshots. +func userWorkspaceHome(b *bundle.Bundle) (string, error) { + u := b.Config.Workspace.CurrentUser + if u == nil || u.User == nil || u.UserName == "" { + return "", errors.New("unable to resolve code snapshot location: current user not set") + } + return "/Workspace/Users/" + u.UserName, nil +} + +// collectLocalCodeSources returns every AI Runtime task code_source_path that +// points at a local directory. Already-remote values are skipped. +func collectLocalCodeSources(b *bundle.Bundle) ([]codeSource, diag.Diagnostics) { + var sources []codeSource + var diags diag.Diagnostics + + for _, pattern := range codeSourcePatterns { + err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { + return dyn.MapByPattern(root, pattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) { + value, ok := v.AsString() + if !ok { + return v, fmt.Errorf("expected string, got %s", v.Kind()) + } + if !libraries.IsLocalPath(value) { + return v, nil + } + sources = append(sources, codeSource{ + configPath: p, + location: v.Location(), + value: value, + }) + return v, nil + }) + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + } + + return sources, diags +} diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go new file mode 100644 index 00000000000..25192332f81 --- /dev/null +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -0,0 +1,74 @@ +package aicode + +import ( + "path/filepath" + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/bundle/internal/bundletest" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/databricks-sdk-go/service/iam" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// bundleWithCodeSource builds a bundle rooted at dir whose single AI Runtime task +// points at codeSourcePath. +// +// The end-to-end package/upload/rewrite behavior (local dir -> tarball -> upload -> +// rewritten code_source_path) runs the full mutator pipeline (sync file list, +// workspace filer) and is covered by acceptance tests under +// acceptance/bundle/ai_runtime_task. This unit test covers only the pure +// config-collection seam that does not touch the pipeline. +func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bundle { + t.Helper() + b := &bundle.Bundle{ + BundleRootPath: dir, + SyncRootPath: dir, + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Workspace: config.Workspace{ + CurrentUser: &config.User{User: &iam.User{UserName: "me@databricks.com"}}, + }, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "train": { + JobSettings: jobs.JobSettings{ + Tasks: []jobs.Task{ + { + TaskKey: "train", + AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp", CodeSourcePath: codeSourcePath}, + }, + }, + }, + }, + }, + }, + }, + } + bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + return b +} + +func TestCollectLocalCodeSourcesFindsLocalPath(t *testing.T) { + b := bundleWithCodeSource(t, t.TempDir(), "./src") + sources, diags := collectLocalCodeSources(b) + require.Empty(t, diags) + require.Len(t, sources, 1) + assert.Equal(t, "./src", sources[0].value) +} + +func TestCollectLocalCodeSourcesSkipsRemotePaths(t *testing.T) { + for _, remote := range []string{ + "/Workspace/Users/me/code.tar.gz", + "/Volumes/main/default/code/existing.tar.gz", + } { + b := bundleWithCodeSource(t, t.TempDir(), remote) + sources, diags := collectLocalCodeSources(b) + require.Empty(t, diags) + assert.Empty(t, sources, "remote code_source_path %q must not be collected", remote) + } +} diff --git a/bundle/config/mutator/aicode/requirements.go b/bundle/config/mutator/aicode/requirements.go new file mode 100644 index 00000000000..6e55d3ea282 --- /dev/null +++ b/bundle/config/mutator/aicode/requirements.go @@ -0,0 +1,137 @@ +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 read by the AI Runtime entry script from the command_path +// directory, so it must sit next to command.sh. +const requirementsFileName = "requirements.yaml" + +// requirementsSpec is the requirements.yaml shape the AI Runtime consumes. +type requirementsSpec struct { + Version string `yaml:"version,omitempty"` + Dependencies []string `yaml:"dependencies,omitempty"` +} + +// SynthesizeRequirements uploads requirements.yaml next to each AI Runtime task's +// command_path, derived from the task's serverless environment. Runs after +// command_path has been translated to its absolute workspace path. +func SynthesizeRequirements() bundle.Mutator { + return &synthesizeRequirements{} +} + +type synthesizeRequirements struct { + client filer.Filer // nil in normal use (a filer is built per directory); set only in tests +} + +func (m *synthesizeRequirements) Name() string { + return "aicode.SynthesizeRequirements" +} + +func (m *synthesizeRequirements) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + + for name, job := range b.Config.Resources.Jobs { + envs := environmentsByKey(job.Environments) + for i := range job.Tasks { + task := &job.Tasks[i] + if task.AiRuntimeTask == nil { + continue + } + if err := m.synthesizeForTask(ctx, b, name, task, envs); err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + } + } + + return diags +} + +// synthesizeForTask uploads requirements.yaml next to task's command_path; a no-op +// when there's no command_path or matching environment (deps can come from elsewhere). +func (m *synthesizeRequirements) synthesizeForTask(ctx context.Context, b *bundle.Bundle, jobName string, task *jobs.Task, envs map[string]*compute.Environment) error { + if len(task.AiRuntimeTask.Deployments) == 0 { + return nil + } + commandPath := task.AiRuntimeTask.Deployments[0].CommandPath + if commandPath == "" { + return nil + } + + env := envs[task.EnvironmentKey] + if env == nil { + return nil + } + + // Skip when the environment has no dependencies: a deps-less file would clobber a + // requirements.yaml the workload supplies another way (e.g. the air CLI's own). + if len(env.Dependencies) == 0 { + return nil + } + + content, err := renderRequirements(env) + if err != nil { + return err + } + + // The entry script reads requirements.yaml from command_path's directory. + dir := path.Dir(commandPath) + client := m.client + if client == nil { + client, err = filer.NewWorkspaceFilesClient(b.WorkspaceClient(ctx), dir) + if err != nil { + return err + } + } + + log.Debugf(ctx, "writing %s for job %s task %s to %s", requirementsFileName, jobName, task.TaskKey, dir) + if err := client.Write(ctx, requirementsFileName, bytes.NewReader(content), filer.OverwriteIfExists, filer.CreateParentDirectories); err != nil { + return fmt.Errorf("failed to upload %s next to command_path %q: %w", requirementsFileName, commandPath, err) + } + return nil +} + +// renderRequirements builds requirements.yaml from a serverless environment spec. +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..6c3684ec132 --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_package.go @@ -0,0 +1,100 @@ +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 stamped on every entry so identical content yields identical bytes +// (and SHA-256) regardless of mtimes, which is what makes the archive content-addressed. +var tarEpoch = time.Date(2000, 1, 1, 0, 0, 0, 0, time.UTC) + +// appleDoublePrefix marks macOS AppleDouble metadata files, excluded to match the AIR CLI. +const appleDoublePrefix = "._" + +// buildCodeSnapshot writes a reproducible gzipped tarball of files to out and returns +// its SHA-256. Each file at "/" becomes "/", so the archive +// expands to /... matching the runtime's /databricks/code_source/ contract. +func buildCodeSnapshot(syncRoot vfs.Path, relBase string, files []fileset.File, prefix string, out io.Writer) (string, error) { + // Sort by relative path so the byte stream (and hash) is order-independent. + 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 { + // Re-base f.Relative (relative to syncRoot) under the code dir so it nests under prefix. + rel := f.Relative + if relBase != "." { + trimmed, ok := strings.CutPrefix(rel, relBase+"/") + if !ok { + return nil // outside the code dir; the file list is scoped to it, so skip defensively + } + 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 (symlinks are out of scope). + if !info.Mode().IsRegular() { + return nil + } + + hdr := &tar.Header{ + Typeflag: tar.TypeReg, + Name: path.Join(prefix, rel), + Size: info.Size(), + // Fixed mode + mtime for reproducibility; code runs via an interpreter, so exec bits don't matter. + 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..eeafb9aa7e4 --- /dev/null +++ b/bundle/config/mutator/aicode/validate.go @@ -0,0 +1,97 @@ +package aicode + +import ( + "context" + "fmt" + "os" + "path/filepath" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/libraries" + "github.com/databricks/cli/libs/diag" + "github.com/databricks/cli/libs/dyn" +) + +// Validate surfaces local code_source_path misconfigurations at `bundle validate` +// time rather than mid-deploy. It performs no uploads. +func Validate() bundle.ReadOnlyMutator { + return &validate{} +} + +type validate struct{ bundle.RO } + +func (v *validate) Name() string { + return "aicode.Validate" +} + +func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + var diags diag.Diagnostics + + jobsPath := dyn.NewPath(dyn.Key("resources"), dyn.Key("jobs")) + + for name, job := range b.Config.Resources.Jobs { + jobPath := jobsPath.Append(dyn.Key(name)) + + for i, task := range job.Tasks { + if task.AiRuntimeTask == nil { + continue + } + codePath := jobPath.Append(dyn.Key("tasks"), dyn.Index(i), + dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) + diags = diags.Extend(v.validateTask(b, job.GitSource != nil, task.AiRuntimeTask.CodeSourcePath, codePath)) + } + } + + return diags +} + +func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { + // Only local values are packaged; remote ones are used as-is. + if codeSourcePath == "" || !libraries.IsLocalPath(codeSourcePath) { + return nil + } + + locations := b.Config.GetLocations(codePath.String()) + + // git_source retrieves task files from git, so a local directory would be ignored. + if jobHasGitSource { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path cannot be combined with git_source", + Detail: "Remove git_source or set code_source_path to a workspace or volume path", + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + // Immutable-folder uploads one snapshot and doesn't support per-task packaging. + if b.IsImmutableFolder() { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: "ai_runtime_task with a local code_source_path is not supported with experimental.immutable_folder", + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath)) + info, err := os.Stat(localDir) + if err != nil { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: fmt.Sprintf("code_source_path %q not found", codeSourcePath), + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + if !info.IsDir() { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: fmt.Sprintf("code_source_path %q must be a directory", codeSourcePath), + Locations: locations, + Paths: []dyn.Path{codePath}, + }} + } + + return nil +} diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go new file mode 100644 index 00000000000..cd3548372e1 --- /dev/null +++ b/bundle/config/mutator/aicode/validate_test.go @@ -0,0 +1,64 @@ +package aicode + +import ( + "path/filepath" + "testing" + + "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" + "github.com/databricks/cli/bundle/config/resources" + "github.com/databricks/cli/bundle/internal/bundletest" + "github.com/databricks/cli/libs/dyn" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func bundleForValidate(t *testing.T, codeSourcePath string, gitSource *jobs.GitSource) *bundle.Bundle { + t.Helper() + dir := t.TempDir() + b := &bundle.Bundle{ + BundleRootPath: dir, + SyncRootPath: dir, + Config: config.Root{ + Bundle: config.Bundle{Target: "default"}, + Resources: config.Resources{ + Jobs: map[string]*resources.Job{ + "train": { + JobSettings: jobs.JobSettings{ + GitSource: gitSource, + Tasks: []jobs.Task{ + { + TaskKey: "train", + AiRuntimeTask: &jobs.AiRuntimeTask{CodeSourcePath: codeSourcePath}, + }, + }, + }, + }, + }, + }, + }, + } + bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + return b +} + +func TestValidateMissingCodeSourceDir(t *testing.T) { + b := bundleForValidate(t, "does-not-exist", nil) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Equal(t, `code_source_path "does-not-exist" not found`, diags[0].Summary) +} + +func TestValidateGitSourceConflict(t *testing.T) { + b := bundleForValidate(t, "src", &jobs.GitSource{GitUrl: "https://example.invalid/repo"}) + diags := Validate().Apply(t.Context(), b) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "cannot be combined with git_source") +} + +func TestValidateRemoteCodeSourceIsSkipped(t *testing.T) { + b := bundleForValidate(t, "/Volumes/main/default/code/x.tar.gz", nil) + diags := Validate().Apply(t.Context(), b) + assert.Empty(t, diags) +} diff --git a/bundle/phases/build.go b/bundle/phases/build.go index db376e07e28..a386093bc24 100644 --- a/bundle/phases/build.go +++ b/bundle/phases/build.go @@ -7,6 +7,7 @@ import ( "github.com/databricks/cli/bundle/artifacts" "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/config/mutator" + "github.com/databricks/cli/bundle/config/mutator/aicode" "github.com/databricks/cli/bundle/libraries" "github.com/databricks/cli/bundle/scripts" "github.com/databricks/cli/bundle/trampoline" @@ -27,6 +28,17 @@ func Build(ctx context.Context, b *bundle.Bundle) LibLocationMap { artifacts.Build(), scripts.Execute(config.ScriptPostBuild), + // Package any AI Runtime task code_source_path that points at a local + // directory into a tarball, upload it, and rewrite the field to the remote + // path. Runs before ExpandGlobReferences/ReplaceWithRemotePath below so those + // see an already-remote code_source_path (a directory would otherwise be + // collected as a local "library" and fail to upload as a file). + aicode.PackageAndUpload(), + // Write requirements.yaml next to the (already translated) command_path, + // derived from the job's serverless environment, so the AI Runtime harness + // can set up the workload environment. + aicode.SynthesizeRequirements(), + mutator.ResolveVariableReferencesWithoutResources( "artifacts", ), diff --git a/bundle/phases/initialize.go b/bundle/phases/initialize.go index bfa2af4124b..f1554a4b7a6 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" @@ -189,6 +190,12 @@ func Initialize(ctx context.Context, b *bundle.Bundle) { mutator.TranslatePaths(), + // Reads (typed): resources.jobs.*.tasks[*].ai_runtime_task.code_source_path, job git_source + // Validates that AI Runtime tasks referencing a local code_source_path point at an existing + // directory and are not combined with git_source or immutable-folder deployments, so these + // misconfigurations are caught at validate time rather than mid-deploy. + aicode.Validate(), + // Reads (typed): b.Config.Experimental.PythonWheelWrapper, b.Config.Presets.SourceLinkedDeployment (checks Python wheel wrapper and deployment mode settings) // Reads (dynamic): resources.jobs.*.tasks (checks for tasks with local libraries and incompatible DBR versions) // Provides warnings when Python wheel tasks are used with DBR < 13.3 or when wheel wrapper is incompatible with source-linked deployment From 5f04cd4293799088e5ae5a40fec9645d3fd5c371 Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 20 Jul 2026 22:06:52 +0000 Subject: [PATCH 2/3] aicode: read code_source_path via dyn, not the typed SDK field jobs.AiRuntimeTask.CodeSourcePath is a private-preview SDK field that is added/removed across releases, so referencing it directly makes the package fail to compile whenever the field is absent. Read code_source_path through dyn.MapByPattern instead (the same mechanism PackageAndUpload and #5922 already use), so the mutator compiles regardless of the SDK's typed struct. Tests set the field on the dyn value rather than the struct literal. Co-authored-by: Isaac --- .../mutator/aicode/package_upload_test.go | 6 +++- bundle/config/mutator/aicode/validate.go | 31 +++++++++++-------- bundle/config/mutator/aicode/validate_test.go | 6 +++- 3 files changed, 28 insertions(+), 15 deletions(-) diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go index 25192332f81..1d33f07f1dc 100644 --- a/bundle/config/mutator/aicode/package_upload_test.go +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -40,7 +40,7 @@ func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bund Tasks: []jobs.Task{ { TaskKey: "train", - AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp", CodeSourcePath: codeSourcePath}, + AiRuntimeTask: &jobs.AiRuntimeTask{Experiment: "exp"}, }, }, }, @@ -50,6 +50,10 @@ func bundleWithCodeSource(t *testing.T, dir, codeSourcePath string) *bundle.Bund }, } bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + // code_source_path is read via dyn, not the typed SDK field, so set it on the value. + bundletest.Mutate(t, b, func(v dyn.Value) (dyn.Value, error) { + return dyn.Set(v, "resources.jobs.train.tasks[0].ai_runtime_task.code_source_path", dyn.V(codeSourcePath)) + }) return b } diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go index eeafb9aa7e4..9139cedb0dc 100644 --- a/bundle/config/mutator/aicode/validate.go +++ b/bundle/config/mutator/aicode/validate.go @@ -27,25 +27,27 @@ func (v *validate) Name() string { 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 + // Read code_source_path via dyn (not the typed SDK field, which is a churny + // private-preview field): the same patterns PackageAndUpload collects. + root := b.Config.Value() + for _, pattern := range codeSourcePatterns { + _, err := dyn.MapByPattern(root, pattern, func(codePath dyn.Path, cv dyn.Value) (dyn.Value, error) { + codeSourcePath, ok := cv.AsString() + if !ok { + return cv, nil } - codePath := jobPath.Append(dyn.Key("tasks"), dyn.Index(i), - dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) - diags = diags.Extend(v.validateTask(b, job.GitSource != nil, task.AiRuntimeTask.CodeSourcePath, codePath)) + diags = diags.Extend(v.validateTask(b, root, codeSourcePath, codePath)) + return cv, nil + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) } } return diags } -func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { +func (v *validate) validateTask(b *bundle.Bundle, root dyn.Value, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { // Only local values are packaged; remote ones are used as-is. if codeSourcePath == "" || !libraries.IsLocalPath(codeSourcePath) { return nil @@ -54,7 +56,10 @@ func (v *validate) validateTask(b *bundle.Bundle, jobHasGitSource bool, codeSour locations := b.Config.GetLocations(codePath.String()) // git_source retrieves task files from git, so a local directory would be ignored. - if jobHasGitSource { + // The job path is code_source_path's path truncated to resources.jobs.. + jobPath := codePath[:3] + _, gitSourceErr := dyn.GetByPath(root, jobPath.Append(dyn.Key("git_source"))) + if gitSourceErr == nil { return diag.Diagnostics{{ Severity: diag.Error, Summary: "ai_runtime_task with a local code_source_path cannot be combined with git_source", diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go index cd3548372e1..09790d6b4dc 100644 --- a/bundle/config/mutator/aicode/validate_test.go +++ b/bundle/config/mutator/aicode/validate_test.go @@ -30,7 +30,7 @@ func bundleForValidate(t *testing.T, codeSourcePath string, gitSource *jobs.GitS Tasks: []jobs.Task{ { TaskKey: "train", - AiRuntimeTask: &jobs.AiRuntimeTask{CodeSourcePath: codeSourcePath}, + AiRuntimeTask: &jobs.AiRuntimeTask{}, }, }, }, @@ -40,6 +40,10 @@ func bundleForValidate(t *testing.T, codeSourcePath string, gitSource *jobs.GitS }, } bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) + // code_source_path is read via dyn, not the typed SDK field, so set it on the value. + bundletest.Mutate(t, b, func(v dyn.Value) (dyn.Value, error) { + return dyn.Set(v, "resources.jobs.train.tasks[0].ai_runtime_task.code_source_path", dyn.V(codeSourcePath)) + }) return b } From 705349094c032f2ec833d8355bc26094448ef61e Mon Sep 17 00:00:00 2001 From: vinchenzo-db Date: Mon, 20 Jul 2026 22:20:51 +0000 Subject: [PATCH 3/3] aicode: make tests independent of the typed SDK field The unit tests built config as typed structs and set code_source_path on the struct/dyn value, which does not survive when the SDK's typed AiRuntimeTask lacks the field (it is a churny private-preview field): Config's typed<->dyn round-trip drops the unknown key, so the tests found zero sources on a field-less SDK. Parse the test config from YAML via yamlloader instead, so code_source_path persists regardless of the typed struct. collectLocalCodeSources now takes a dyn.Value (read-only) and validate's core is a pure validateCodeSources function; both are exercised directly against the parsed value. Verified the package builds and all tests pass against both SDK v0.154.0 (field present) and v0.160.0 (field absent). Co-authored-by: Isaac --- .../config/mutator/aicode/package_upload.go | 35 ++++---- .../mutator/aicode/package_upload_test.go | 65 +++++--------- bundle/config/mutator/aicode/validate.go | 32 ++++--- bundle/config/mutator/aicode/validate_test.go | 88 ++++++++++--------- 4 files changed, 104 insertions(+), 116 deletions(-) diff --git a/bundle/config/mutator/aicode/package_upload.go b/bundle/config/mutator/aicode/package_upload.go index 4a59792cefb..449b6fd9712 100644 --- a/bundle/config/mutator/aicode/package_upload.go +++ b/bundle/config/mutator/aicode/package_upload.go @@ -60,7 +60,7 @@ func (m *packageAndUpload) Name() string { } func (m *packageAndUpload) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { - sources, diags := collectLocalCodeSources(b) + sources, diags := collectLocalCodeSources(b.Config.Value()) if diags.HasError() { return diags } @@ -182,29 +182,28 @@ func userWorkspaceHome(b *bundle.Bundle) (string, error) { return "/Workspace/Users/" + u.UserName, nil } -// collectLocalCodeSources returns every AI Runtime task code_source_path that -// points at a local directory. Already-remote values are skipped. -func collectLocalCodeSources(b *bundle.Bundle) ([]codeSource, diag.Diagnostics) { +// collectLocalCodeSources returns every AI Runtime task code_source_path in root +// that points at a local directory. Already-remote values are skipped. It reads the +// dyn value directly (read-only), so it does not depend on the typed SDK field. +func collectLocalCodeSources(root dyn.Value) ([]codeSource, diag.Diagnostics) { var sources []codeSource var diags diag.Diagnostics for _, pattern := range codeSourcePatterns { - err := b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { - return dyn.MapByPattern(root, pattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) { - value, ok := v.AsString() - if !ok { - return v, fmt.Errorf("expected string, got %s", v.Kind()) - } - if !libraries.IsLocalPath(value) { - return v, nil - } - sources = append(sources, codeSource{ - configPath: p, - location: v.Location(), - value: value, - }) + _, err := dyn.MapByPattern(root, pattern, func(p dyn.Path, v dyn.Value) (dyn.Value, error) { + value, ok := v.AsString() + if !ok { + return v, fmt.Errorf("expected string, got %s", v.Kind()) + } + if !libraries.IsLocalPath(value) { return v, nil + } + sources = append(sources, codeSource{ + configPath: p, + location: v.Location(), + value: value, }) + return v, nil }) if err != nil { diags = diags.Extend(diag.FromErr(err)) diff --git a/bundle/config/mutator/aicode/package_upload_test.go b/bundle/config/mutator/aicode/package_upload_test.go index 1d33f07f1dc..26fdc5f3368 100644 --- a/bundle/config/mutator/aicode/package_upload_test.go +++ b/bundle/config/mutator/aicode/package_upload_test.go @@ -1,65 +1,43 @@ package aicode import ( - "path/filepath" + "strings" "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/databricks/cli/libs/dyn/yamlloader" "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. +// configWithCodeSource parses a minimal bundle config whose single AI Runtime task +// points at codeSourcePath, as a dyn value. It reads YAML directly so the config +// retains code_source_path regardless of whether the SDK's typed struct has it. // // 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 { +func configWithCodeSource(t *testing.T, codeSourcePath string) dyn.Value { 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"}, - }, - }, - }, - }, - }, - }, - }, - } - bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) - // code_source_path is read via dyn, not the typed SDK field, so set it on the value. - bundletest.Mutate(t, b, func(v dyn.Value) (dyn.Value, error) { - return dyn.Set(v, "resources.jobs.train.tasks[0].ai_runtime_task.code_source_path", dyn.V(codeSourcePath)) - }) - return b + yml := ` +resources: + jobs: + train: + tasks: + - task_key: train + ai_runtime_task: + experiment: exp + code_source_path: ` + codeSourcePath + ` +` + v, err := yamlloader.LoadYAML("test.yml", strings.NewReader(yml)) + require.NoError(t, err) + return v } func TestCollectLocalCodeSourcesFindsLocalPath(t *testing.T) { - b := bundleWithCodeSource(t, t.TempDir(), "./src") - sources, diags := collectLocalCodeSources(b) + sources, diags := collectLocalCodeSources(configWithCodeSource(t, "./src")) require.Empty(t, diags) require.Len(t, sources, 1) assert.Equal(t, "./src", sources[0].value) @@ -70,8 +48,7 @@ func TestCollectLocalCodeSourcesSkipsRemotePaths(t *testing.T) { "/Workspace/Users/me/code.tar.gz", "/Volumes/main/default/code/existing.tar.gz", } { - b := bundleWithCodeSource(t, t.TempDir(), remote) - sources, diags := collectLocalCodeSources(b) + sources, diags := collectLocalCodeSources(configWithCodeSource(t, remote)) require.Empty(t, diags) assert.Empty(t, sources, "remote code_source_path %q must not be collected", remote) } diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go index 9139cedb0dc..72946fa99ef 100644 --- a/bundle/config/mutator/aicode/validate.go +++ b/bundle/config/mutator/aicode/validate.go @@ -25,18 +25,23 @@ func (v *validate) Name() string { } func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { + return validateCodeSources(b.Config.Value(), b.SyncRootPath, b.IsImmutableFolder(), b.Config.GetLocations) +} + +// validateCodeSources checks every local code_source_path in root. It reads the dyn +// value directly (not the typed SDK field, which is a churny private-preview field), +// so it works regardless of the SDK's typed struct. locations resolves a config path +// to its source locations for diagnostics. +func validateCodeSources(root dyn.Value, syncRootPath string, isImmutableFolder bool, locations func(string) []dyn.Location) diag.Diagnostics { var diags diag.Diagnostics - // Read code_source_path via dyn (not the typed SDK field, which is a churny - // private-preview field): the same patterns PackageAndUpload collects. - root := b.Config.Value() for _, pattern := range codeSourcePatterns { _, err := dyn.MapByPattern(root, pattern, func(codePath dyn.Path, cv dyn.Value) (dyn.Value, error) { codeSourcePath, ok := cv.AsString() if !ok { return cv, nil } - diags = diags.Extend(v.validateTask(b, root, codeSourcePath, codePath)) + diags = diags.Extend(validateTask(root, syncRootPath, isImmutableFolder, locations, codeSourcePath, codePath)) return cv, nil }) if err != nil { @@ -47,45 +52,44 @@ func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics return diags } -func (v *validate) validateTask(b *bundle.Bundle, root dyn.Value, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { +func validateTask(root dyn.Value, syncRootPath string, isImmutableFolder bool, locations func(string) []dyn.Location, codeSourcePath string, codePath dyn.Path) diag.Diagnostics { // Only local values are packaged; remote ones are used as-is. if codeSourcePath == "" || !libraries.IsLocalPath(codeSourcePath) { return nil } - locations := b.Config.GetLocations(codePath.String()) + locs := locations(codePath.String()) // git_source retrieves task files from git, so a local directory would be ignored. // The job path is code_source_path's path truncated to resources.jobs.. jobPath := codePath[:3] - _, gitSourceErr := dyn.GetByPath(root, jobPath.Append(dyn.Key("git_source"))) - if gitSourceErr == nil { + if _, err := dyn.GetByPath(root, jobPath.Append(dyn.Key("git_source"))); err == 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, + Locations: locs, Paths: []dyn.Path{codePath}, }} } // Immutable-folder uploads one snapshot and doesn't support per-task packaging. - if b.IsImmutableFolder() { + if 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, + Locations: locs, Paths: []dyn.Path{codePath}, }} } - localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(codeSourcePath)) + localDir := filepath.Join(syncRootPath, filepath.FromSlash(codeSourcePath)) info, err := os.Stat(localDir) if err != nil { return diag.Diagnostics{{ Severity: diag.Error, Summary: fmt.Sprintf("code_source_path %q not found", codeSourcePath), - Locations: locations, + Locations: locs, Paths: []dyn.Path{codePath}, }} } @@ -93,7 +97,7 @@ func (v *validate) validateTask(b *bundle.Bundle, root dyn.Value, codeSourcePath return diag.Diagnostics{{ Severity: diag.Error, Summary: fmt.Sprintf("code_source_path %q must be a directory", codeSourcePath), - Locations: locations, + Locations: locs, Paths: []dyn.Path{codePath}, }} } diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go index 09790d6b4dc..8a971254e0b 100644 --- a/bundle/config/mutator/aicode/validate_test.go +++ b/bundle/config/mutator/aicode/validate_test.go @@ -1,68 +1,76 @@ package aicode import ( + "os" "path/filepath" + "strings" "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/databricks/cli/libs/dyn/yamlloader" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) -func bundleForValidate(t *testing.T, codeSourcePath string, gitSource *jobs.GitSource) *bundle.Bundle { +// configForValidate parses a bundle config with one AI Runtime task pointing at +// codeSourcePath (and optionally a job git_source), as a dyn value. Reading YAML +// directly keeps code_source_path regardless of the SDK's typed struct. +func configForValidate(t *testing.T, codeSourcePath string, withGitSource bool) dyn.Value { 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{}, - }, - }, - }, - }, - }, - }, - }, + git := "" + if withGitSource { + git = "\n git_source:\n git_url: https://example.invalid/repo\n git_provider: gitHub" } - bundletest.SetLocation(b, ".", []dyn.Location{{File: filepath.Join(dir, "databricks.yml")}}) - // code_source_path is read via dyn, not the typed SDK field, so set it on the value. - bundletest.Mutate(t, b, func(v dyn.Value) (dyn.Value, error) { - return dyn.Set(v, "resources.jobs.train.tasks[0].ai_runtime_task.code_source_path", dyn.V(codeSourcePath)) - }) - return b + yml := ` +resources: + jobs: + train:` + git + ` + tasks: + - task_key: train + ai_runtime_task: + experiment: exp + code_source_path: ` + codeSourcePath + ` +` + v, err := yamlloader.LoadYAML("test.yml", strings.NewReader(yml)) + require.NoError(t, err) + return v } +// noLocations is a stub locations resolver for tests (diagnostics content, not +// locations, is what these assert). +func noLocations(string) []dyn.Location { return nil } + func TestValidateMissingCodeSourceDir(t *testing.T) { - b := bundleForValidate(t, "does-not-exist", nil) - diags := Validate().Apply(t.Context(), b) + root := configForValidate(t, "does-not-exist", false) + diags := validateCodeSources(root, t.TempDir(), false, noLocations) require.Len(t, diags, 1) assert.Equal(t, `code_source_path "does-not-exist" not found`, diags[0].Summary) } func TestValidateGitSourceConflict(t *testing.T) { - b := bundleForValidate(t, "src", &jobs.GitSource{GitUrl: "https://example.invalid/repo"}) - diags := Validate().Apply(t.Context(), b) + root := configForValidate(t, "src", true) + diags := validateCodeSources(root, t.TempDir(), false, noLocations) require.Len(t, diags, 1) assert.Contains(t, diags[0].Summary, "cannot be combined with git_source") } +func TestValidateImmutableFolderConflict(t *testing.T) { + root := configForValidate(t, "src", false) + diags := validateCodeSources(root, t.TempDir(), true, noLocations) + require.Len(t, diags, 1) + assert.Contains(t, diags[0].Summary, "experimental.immutable_folder") +} + +func TestValidateAcceptsExistingDir(t *testing.T) { + dir := t.TempDir() + require.NoError(t, os.MkdirAll(filepath.Join(dir, "src"), 0o755)) + root := configForValidate(t, "src", false) + diags := validateCodeSources(root, dir, false, noLocations) + assert.Empty(t, diags) +} + func TestValidateRemoteCodeSourceIsSkipped(t *testing.T) { - b := bundleForValidate(t, "/Volumes/main/default/code/x.tar.gz", nil) - diags := Validate().Apply(t.Context(), b) + root := configForValidate(t, "/Volumes/main/default/code/x.tar.gz", false) + diags := validateCodeSources(root, t.TempDir(), false, noLocations) assert.Empty(t, diags) }