diff --git a/.nextchanges/bundles/ai-runtime-code-source-block.md b/.nextchanges/bundles/ai-runtime-code-source-block.md new file mode 100644 index 00000000000..b9a5779c395 --- /dev/null +++ b/.nextchanges/bundles/ai-runtime-code-source-block.md @@ -0,0 +1 @@ +Add a DABs-native `code_source` block to `ai_runtime_task`, matching the AIR CLI's `train.yaml`. Set `code_source: { root_path, include_paths, git }` on the task and `bundle deploy` packages the directory into a content-addressed tarball (honoring `.gitignore` and sync include/exclude, narrowing to `include_paths`, or archiving a pinned `git` branch/commit), uploads it, and sets `code_source_path` for you. The existing string `code_source_path: ./dir` form is unchanged; the two are mutually exclusive on a task. diff --git a/acceptance/bundle/ai_runtime_task/code_source/block.yml b/acceptance/bundle/ai_runtime_task/code_source/block.yml new file mode 100644 index 00000000000..a6c5c4af67f --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/block.yml @@ -0,0 +1,28 @@ +bundle: + name: ai-runtime-block + +resources: + jobs: + train: + name: "[${bundle.target}] AI Runtime training" + tasks: + # code_source block, root_path form: packages ./src exactly like a bare + # code_source_path (content-addressed tarball under .air_snapshots, uploaded by + # file sync, code_source_path rewritten to the uploaded workspace path). + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source: + root_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/code_source/git.yml b/acceptance/bundle/ai_runtime_task/code_source/git.yml new file mode 100644 index 00000000000..b119811289b --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/git.yml @@ -0,0 +1,28 @@ +bundle: + name: ai-runtime-git + +resources: + jobs: + train: + name: "[${bundle.target}] AI Runtime training" + tasks: + # code_source.git pins the snapshot to a committed revision: the archive is + # produced with `git archive` of that commit, not the working tree. COMMIT_SHA + # is substituted with the repo HEAD by the script. + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source: + root_path: ./src + git: + commit: COMMIT_SHA + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_1xA10 + accelerator_count: 1 + environments: + - environment_key: default + spec: + environment_version: "5" diff --git a/acceptance/bundle/ai_runtime_task/code_source/include_paths.yml b/acceptance/bundle/ai_runtime_task/code_source/include_paths.yml new file mode 100644 index 00000000000..3b1daabe63b --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/include_paths.yml @@ -0,0 +1,28 @@ +bundle: + name: ai-runtime-includes + +resources: + jobs: + train: + name: "[${bundle.target}] AI Runtime training" + tasks: + # include_paths narrows the packaged archive to a subset of root_path: only + # keep/ and configs/ are archived; drop/ is left out. + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source: + root_path: ./src + include_paths: + - keep + - configs + deployments: + - command_path: src/keep/command.sh + compute: + accelerator_type: GPU_1xA10 + accelerator_count: 1 + environments: + - environment_key: default + spec: + environment_version: "5" diff --git a/acceptance/bundle/ai_runtime_task/code_source/out.test.toml b/acceptance/bundle/ai_runtime_task/code_source/out.test.toml new file mode 100644 index 00000000000..98ea5040486 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["terraform", "direct"] diff --git a/acceptance/bundle/ai_runtime_task/code_source/output.txt b/acceptance/bundle/ai_runtime_task/code_source/output.txt new file mode 100644 index 00000000000..4a8bca84274 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/output.txt @@ -0,0 +1,181 @@ + +=== root_path form packages ./src like a bare code_source_path + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-block/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> list_code_snapshot.py +# .air_snapshots/src_[SNAPSHOT].tar.gz +src/command.sh +src/configs/model.yaml +src/drop/ignored.py +src/keep/command.sh +src/keep/train.py +src/train.py + +>>> print_requests.py --sort --del-field raw_body //.air_snapshots/ //jobs/create +{ + "method": "POST", + "path": "/api/2.0/workspace-files/import-file/Workspace/Users/[USERNAME]/.bundle/ai-runtime-block/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", + "q": { + "overwrite": "true" + } +} +{ + "method": "POST", + "path": "/api/2.2/jobs/create", + "body": { + "deployment": { + "kind": "BUNDLE", + "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-block/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-block/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-block/default/files/src/command.sh", + "compute": { + "accelerator_count": 8, + "accelerator_type": "GPU_8xH100" + } + } + ], + "experiment": "my-training" + }, + "environment_key": "default", + "task_key": "train" + } + ] + } +} + +>>> [CLI] bundle destroy --auto-approve +The following resources will be deleted: + delete resources.jobs.train + +All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/ai-runtime-block/default + +Deleting files... +Destroy complete! + +=== include_paths narrows the archive to keep/ and configs/ (drop/ excluded) + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-includes/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> list_code_snapshot.py +# .air_snapshots/src_[SNAPSHOT].tar.gz +src/configs/model.yaml +src/keep/command.sh +src/keep/train.py + +>>> [CLI] bundle destroy --auto-approve +The following resources will be deleted: + delete resources.jobs.train + +All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/ai-runtime-includes/default + +Deleting files... +Destroy complete! + +=== git.commit archives the pinned revision via git archive + +>>> [CLI] bundle deploy +Uploading bundle files to /Workspace/Users/[USERNAME]/.bundle/ai-runtime-git/default/files... +Deploying resources... +Updating deployment state... +Deployment complete! + +>>> list_code_snapshot.py +# .air_snapshots/src_[SNAPSHOT].tar.gz +src +src/command.sh +src/configs +src/configs/model.yaml +src/drop +src/drop/ignored.py +src/keep +src/keep/command.sh +src/keep/train.py +src/train.py + +>>> print_requests.py --sort --del-field raw_body //jobs/create +{ + "method": "POST", + "path": "/api/2.2/jobs/create", + "body": { + "deployment": { + "kind": "BUNDLE", + "metadata_file_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-git/default/state/metadata.json" + }, + "edit_mode": "UI_LOCKED", + "environments": [ + { + "environment_key": "default", + "spec": { + "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-git/default/files/.air_snapshots/src_[SNAPSHOT].tar.gz", + "deployments": [ + { + "command_path": "/Workspace/Users/[USERNAME]/.bundle/ai-runtime-git/default/files/src/command.sh", + "compute": { + "accelerator_count": 1, + "accelerator_type": "GPU_1xA10" + } + } + ], + "experiment": "my-training" + }, + "environment_key": "default", + "task_key": "train" + } + ] + } +} + +>>> [CLI] bundle destroy --auto-approve +The following resources will be deleted: + delete resources.jobs.train + +All files and directories at the following location will be deleted: /Workspace/Users/[USERNAME]/.bundle/ai-runtime-git/default + +Deleting files... +Destroy complete! diff --git a/acceptance/bundle/ai_runtime_task/code_source/script b/acceptance/bundle/ai_runtime_task/code_source/script new file mode 100644 index 00000000000..56fc1add560 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/script @@ -0,0 +1,38 @@ +# The DABs-native ai_runtime_task.code_source block, exercised end-to-end in three +# forms against a shared ./src tree. Each fixture is copied to databricks.yml, deployed +# (packaging + upload + code_source_path rewrite), inspected, then destroyed. The git +# form runs last because it needs a git repo; the working-tree forms must run before +# git-repo-init so they exercise the non-git filesystem walk. + +title "root_path form packages ./src like a bare code_source_path\n" +cp block.yml databricks.yml +trace $CLI bundle deploy +trace list_code_snapshot.py +trace print_requests.py --sort --del-field raw_body '//.air_snapshots/' '//jobs/create' +trace $CLI bundle destroy --auto-approve +rm out.requests.txt + +title "include_paths narrows the archive to keep/ and configs/ (drop/ excluded)\n" +cp include_paths.yml databricks.yml +trace $CLI bundle deploy +trace list_code_snapshot.py +trace $CLI bundle destroy --auto-approve +rm out.requests.txt + +title "git.commit archives the pinned revision via git archive\n" +# Pinned commit dates keep the resolved HEAD SHA stable across runs. +export GIT_AUTHOR_DATE="2020-01-01T00:00:00Z" +export GIT_COMMITTER_DATE="2020-01-01T00:00:00Z" +git-repo-init +git add src +git commit -qm 'Add src' +# git.commit tolerates a dirty working tree (the pinned revision is archived), which +# matters here because the harness writes output.txt into the working dir as it runs. +cp git.yml databricks.yml +update_file.py databricks.yml COMMIT_SHA "$(git rev-parse HEAD)" +trace $CLI bundle deploy +trace list_code_snapshot.py +trace print_requests.py --sort --del-field raw_body '//jobs/create' +trace $CLI bundle destroy --auto-approve + +rm -fr .git out.requests.txt databricks.yml diff --git a/acceptance/bundle/ai_runtime_task/code_source/src/command.sh b/acceptance/bundle/ai_runtime_task/code_source/src/command.sh new file mode 100644 index 00000000000..12e21df9dc7 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/src/command.sh @@ -0,0 +1,2 @@ +#!/bin/bash +python train.py diff --git a/acceptance/bundle/ai_runtime_task/code_source/src/configs/model.yaml b/acceptance/bundle/ai_runtime_task/code_source/src/configs/model.yaml new file mode 100644 index 00000000000..2330fe243c2 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/src/configs/model.yaml @@ -0,0 +1 @@ +lr: 0.001 diff --git a/acceptance/bundle/ai_runtime_task/code_source/src/drop/ignored.py b/acceptance/bundle/ai_runtime_task/code_source/src/drop/ignored.py new file mode 100644 index 00000000000..680ce262a5d --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/src/drop/ignored.py @@ -0,0 +1 @@ +print("this file must not be packaged") diff --git a/acceptance/bundle/ai_runtime_task/code_source/src/keep/command.sh b/acceptance/bundle/ai_runtime_task/code_source/src/keep/command.sh new file mode 100644 index 00000000000..12e21df9dc7 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/src/keep/command.sh @@ -0,0 +1,2 @@ +#!/bin/bash +python train.py diff --git a/acceptance/bundle/ai_runtime_task/code_source/src/keep/train.py b/acceptance/bundle/ai_runtime_task/code_source/src/keep/train.py new file mode 100644 index 00000000000..d8b062e5dfe --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/src/keep/train.py @@ -0,0 +1 @@ +print("training") diff --git a/acceptance/bundle/ai_runtime_task/code_source/src/train.py b/acceptance/bundle/ai_runtime_task/code_source/src/train.py new file mode 100644 index 00000000000..d8b062e5dfe --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/src/train.py @@ -0,0 +1 @@ +print("training") diff --git a/acceptance/bundle/ai_runtime_task/code_source/test.toml b/acceptance/bundle/ai_runtime_task/code_source/test.toml new file mode 100644 index 00000000000..e02bc4afbcc --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source/test.toml @@ -0,0 +1,13 @@ +RecordRequests = true + +Ignore = [ + '.databricks', + 'databricks.yml', +] + +# The archive is content-addressed: _.tar.gz. Collapse the hash to a +# token so the test does not pin a specific digest (git's gzip output is also not +# byte-identical across git versions). +[[Repls]] +Old = 'src_[0-9a-f]{16}\.tar\.gz' +New = 'src_[SNAPSHOT].tar.gz' diff --git a/acceptance/bundle/ai_runtime_task/code_source_errors/conflict.yml b/acceptance/bundle/ai_runtime_task/code_source_errors/conflict.yml new file mode 100644 index 00000000000..762fc3848ed --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source_errors/conflict.yml @@ -0,0 +1,24 @@ +bundle: + name: ai-runtime-conflict + +resources: + jobs: + train: + tasks: + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + # Both forms set on one task: mutually exclusive, rejected at validate. + code_source_path: ./src + code_source: + root_path: ./src + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_1xA10 + accelerator_count: 1 + environments: + - environment_key: default + spec: + environment_version: "5" diff --git a/acceptance/bundle/ai_runtime_task/code_source_errors/out.test.toml b/acceptance/bundle/ai_runtime_task/code_source_errors/out.test.toml new file mode 100644 index 00000000000..0938e678987 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source_errors/out.test.toml @@ -0,0 +1,2 @@ +Cloud = false +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] diff --git a/acceptance/bundle/ai_runtime_task/code_source_errors/output.txt b/acceptance/bundle/ai_runtime_task/code_source_errors/output.txt new file mode 100644 index 00000000000..148f9bf2eea --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source_errors/output.txt @@ -0,0 +1,40 @@ + +=== code_source and code_source_path on the same task are mutually exclusive + +>>> [CLI] bundle validate +Error: ai_runtime_task sets both code_source and code_source_path + at resources.jobs.train.tasks[0].ai_runtime_task.code_source + +Set either the code_source block or code_source_path, not both + +Name: ai-runtime-conflict +Target: default +Workspace: + User: [USERNAME] + Path: /Workspace/Users/[USERNAME]/.bundle/ai-runtime-conflict/default + +Found 1 error + +=== remote_volume is not yet supported + +>>> [CLI] bundle validate +Error: code_source.remote_volume is not yet supported + at resources.jobs.train.tasks[0].ai_runtime_task.code_source + +Remove remote_volume; the archive is uploaded to the bundle's workspace file path + +Name: ai-runtime-remote-volume +Target: default +Workspace: + User: [USERNAME] + Path: /Workspace/Users/[USERNAME]/.bundle/ai-runtime-remote-volume/default + +Found 1 error + +=== code_source under a target override is rejected at load + +>>> [CLI] bundle validate -t dev +Error: failed to rewrite [TEST_TMP_DIR]/databricks.yml: ai_runtime_task.code_source at [TEST_TMP_DIR]/databricks.yml:33:19 is not supported inside a target override; move it to the task's base definition under resources.jobs + + +Found 1 error diff --git a/acceptance/bundle/ai_runtime_task/code_source_errors/remote_volume.yml b/acceptance/bundle/ai_runtime_task/code_source_errors/remote_volume.yml new file mode 100644 index 00000000000..6167da52a62 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source_errors/remote_volume.yml @@ -0,0 +1,25 @@ +bundle: + name: ai-runtime-remote-volume + +resources: + jobs: + train: + tasks: + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + code_source: + root_path: ./src + # remote_volume is not yet supported (deploy-phase Volume upload + no + # bundle-destroy cleanup are unresolved). Rejected at validate. + remote_volume: /Volumes/main/default/code + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_1xA10 + accelerator_count: 1 + environments: + - environment_key: default + spec: + environment_version: "5" diff --git a/acceptance/bundle/ai_runtime_task/code_source_errors/script b/acceptance/bundle/ai_runtime_task/code_source_errors/script new file mode 100644 index 00000000000..fdf3d0b4ec3 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source_errors/script @@ -0,0 +1,17 @@ +# Error cases for the DABs-native ai_runtime_task.code_source block. Each fixture is +# copied to databricks.yml and validated; every one must fail with an actionable +# message (musterr fails the test if validate unexpectedly succeeds). + +title "code_source and code_source_path on the same task are mutually exclusive\n" +cp conflict.yml databricks.yml +musterr trace $CLI bundle validate + +title "remote_volume is not yet supported\n" +cp remote_volume.yml databricks.yml +musterr trace $CLI bundle validate + +title "code_source under a target override is rejected at load\n" +cp target_override.yml databricks.yml +musterr trace $CLI bundle validate -t dev + +rm databricks.yml diff --git a/acceptance/bundle/ai_runtime_task/code_source_errors/src/command.sh b/acceptance/bundle/ai_runtime_task/code_source_errors/src/command.sh new file mode 100644 index 00000000000..12e21df9dc7 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source_errors/src/command.sh @@ -0,0 +1,2 @@ +#!/bin/bash +python train.py diff --git a/acceptance/bundle/ai_runtime_task/code_source_errors/src/train.py b/acceptance/bundle/ai_runtime_task/code_source_errors/src/train.py new file mode 100644 index 00000000000..d8b062e5dfe --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source_errors/src/train.py @@ -0,0 +1 @@ +print("training") diff --git a/acceptance/bundle/ai_runtime_task/code_source_errors/target_override.yml b/acceptance/bundle/ai_runtime_task/code_source_errors/target_override.yml new file mode 100644 index 00000000000..eee82efc9e3 --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source_errors/target_override.yml @@ -0,0 +1,33 @@ +bundle: + name: ai-runtime-target-override + +resources: + jobs: + train: + tasks: + - task_key: train + environment_key: default + ai_runtime_task: + experiment: my-training + deployments: + - command_path: src/command.sh + compute: + accelerator_type: GPU_1xA10 + accelerator_count: 1 + environments: + - environment_key: default + spec: + environment_version: "5" + +targets: + dev: + resources: + jobs: + train: + tasks: + # A code_source under a target override is not supported: the (job, + # task_key) stash can't represent per-target overrides. Rejected at load. + - task_key: train + ai_runtime_task: + code_source: + root_path: ./src diff --git a/acceptance/bundle/ai_runtime_task/code_source_errors/test.toml b/acceptance/bundle/ai_runtime_task/code_source_errors/test.toml new file mode 100644 index 00000000000..572e612ae9c --- /dev/null +++ b/acceptance/bundle/ai_runtime_task/code_source_errors/test.toml @@ -0,0 +1,7 @@ +# Load/validate-time errors only; no deploy, so run a single engine. +EnvMatrix.DATABRICKS_BUNDLE_ENGINE = ["direct"] + +Ignore = [ + '.databricks', + 'databricks.yml', +] diff --git a/bundle/config/ai_runtime.go b/bundle/config/ai_runtime.go new file mode 100644 index 00000000000..ae8d7bdea85 --- /dev/null +++ b/bundle/config/ai_runtime.go @@ -0,0 +1,39 @@ +package config + +// AiRuntimeTaskExtras holds the DABs-native, non-proto authoring sugar for one +// ai_runtime_task. These are not fields on the SDK jobs.AiRuntimeTask: they are +// extracted from the config tree by rewriteAiRuntimeCodeSource before normalization +// (which drops unknown struct fields) and lowered into the task by the aicode +// build-phase mutator. Today only code_source; parameters (lowered to a +// hyperparameters.yaml sidecar) is the intended next field. +type AiRuntimeTaskExtras struct { + CodeSource *CodeSourceOptions `json:"code_source,omitempty"` +} + +// CodeSourceOptions is the DABs-native `ai_runtime_task.code_source` block. It gives +// train.yaml parity for code delivery: at deploy the aicode mutator packages the +// directory (honoring .gitignore + sync include/exclude, plus git ref and +// include_paths when set), uploads it, and sets the real SDK field +// ai_runtime_task.code_source_path to the uploaded path. +type CodeSourceOptions struct { + // RootPath is the local directory to package, relative to the bundle sync root. + RootPath string `json:"root_path"` + + // IncludePaths narrows the archive to these subtrees of RootPath (relative, no + // "..") instead of packaging the whole directory. + IncludePaths []string `json:"include_paths,omitempty"` + + // Git pins the archive to a committed revision instead of the working tree. + Git *CodeSourceGit `json:"git,omitempty"` + + // RemoteVolume uploads the archive to this UC Volume path (starts with + // "/Volumes/") instead of the bundle's workspace file path. + RemoteVolume string `json:"remote_volume,omitempty"` +} + +// CodeSourceGit pins a code_source snapshot to a git ref. Branch and commit are +// mutually exclusive. +type CodeSourceGit struct { + Branch string `json:"branch,omitempty"` + Commit string `json:"commit,omitempty"` +} diff --git a/bundle/config/mutator/aicode/package_code_source.go b/bundle/config/mutator/aicode/package_code_source.go index b2f78a9aebd..e93e3c577e5 100644 --- a/bundle/config/mutator/aicode/package_code_source.go +++ b/bundle/config/mutator/aicode/package_code_source.go @@ -27,6 +27,7 @@ import ( "strings" "github.com/databricks/cli/bundle" + "github.com/databricks/cli/bundle/config" "github.com/databricks/cli/bundle/deploy/files" "github.com/databricks/cli/bundle/libraries" "github.com/databricks/cli/libs/diag" @@ -66,32 +67,65 @@ func (m *packageCodeSource) Name() string { } func (m *packageCodeSource) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics { - sources, diags := collectLocalCodeSources(b) + // Package two source kinds into a single list of archives to overlay: + // - a local-directory code_source_path string (collectLocalCodeSources), and + // - a DABs-native code_source block (collectCodeSourceBlocks), which additionally + // supports include_paths narrowing and git-ref pinning. + // Both lower to the same shape: a config path (…ai_runtime_task.code_source_path) to + // set to the uploaded workspace path, plus the archive bytes to overlay. + var packaged []packagedSource + + strSources, diags := collectLocalCodeSources(b) if diags.HasError() { return diags } - if len(sources) == 0 { - return diags + for _, cs := range strSources { + relArchive, archive, err := packageOne(ctx, b, cs) + if err != nil { + return diags.Extend(diag.FromErr(err)) + } + packaged = append(packaged, packagedSource{ + configPath: cs.configPath, + location: cs.location, + relArchive: relArchive, + archive: archive, + }) } - // 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) + blockSources, blockDiags := collectCodeSourceBlocks(b) + diags = diags.Extend(blockDiags) + if diags.HasError() { + return diags + } + for _, cs := range blockSources { + relArchive, archive, err := packageBlock(ctx, b, cs) if err != nil { - diags = diags.Extend(diag.FromErr(err)) - return diags + return diags.Extend(diag.FromErr(err)) } - overlayFiles[relArchive] = archive + packaged = append(packaged, packagedSource{ + configPath: cs.configPath, + location: cs.location, + relArchive: relArchive, + archive: archive, + }) + } + + if len(packaged) == 0 { + return diags + } + + // overlayFiles maps each archive's sync-relative path to its bytes; remotePaths maps + // each config location to the synced workspace path its archive will occupy. 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. + overlayFiles := make(map[string][]byte, len(packaged)) + remotePaths := make(map[string]string, len(packaged)) + for _, p := range packaged { + overlayFiles[p.relArchive] = p.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) + remotePaths[p.configPath.String()] = path.Join(b.Config.Workspace.FilePath, p.relArchive) } // Overlay the archives onto the sync root: bundle file sync walks and uploads @@ -107,12 +141,12 @@ func (m *packageCodeSource) Apply(ctx context.Context, b *bundle.Bundle) diag.Di b.HasAiRuntimeCodeSnapshot = true err = b.Config.Mutate(func(root dyn.Value) (dyn.Value, error) { - for _, cs := range sources { - remote := remotePaths[cs.configPath.String()] + for _, p := range packaged { + remote := remotePaths[p.configPath.String()] var err error - root, err = dyn.SetByPath(root, cs.configPath, dyn.NewValue(remote, []dyn.Location{cs.location})) + root, err = dyn.SetByPath(root, p.configPath, dyn.NewValue(remote, []dyn.Location{p.location})) if err != nil { - return root, fmt.Errorf("failed to update code_source_path %q to %q: %w", cs.value, remote, err) + return root, fmt.Errorf("failed to update code_source_path to %q: %w", remote, err) } } return root, nil @@ -124,6 +158,17 @@ func (m *packageCodeSource) Apply(ctx context.Context, b *bundle.Bundle) diag.Di return diags } +// packagedSource is one code source packaged into an archive, ready to overlay on the +// sync root and to rewrite into the task's code_source_path. +type packagedSource struct { + // configPath is the …ai_runtime_task.code_source_path to set to the uploaded path. + configPath dyn.Path + location dyn.Location + // relArchive is the archive's sync-relative path; archive is its bytes. + relArchive string + archive []byte +} + // snapshotSubdir is the sync-relative dir the archives are placed under (dedicated // so a snapshot is never nested in the dir it snapshots). See bundle.AiCodeSnapshotDir. const snapshotSubdir = bundle.AiCodeSnapshotDir @@ -245,6 +290,137 @@ func collectLocalCodeSources(b *bundle.Bundle) ([]codeSource, diag.Diagnostics) return sources, diags } +// taskPattern matches every job task, used to locate a task by task_key when lowering +// a code_source block (keyed by task_key in Root.AiRuntimeExtras). +var taskPattern = dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), +) + +// blockSource is a DABs-native code_source block to package: the config path of the +// task's code_source_path to write, a location for diagnostics, and the block options. +type blockSource struct { + configPath dyn.Path + location dyn.Location + opts config.CodeSourceOptions +} + +// collectCodeSourceBlocks returns one blockSource per task that carries a code_source +// block (extracted into Root.AiRuntimeExtras before normalization). It matches each +// stashed (job, task_key) to the task's current position in the config tree — so it +// tracks the task after MergeJobTasks reorders tasks by task_key — and yields the +// …ai_runtime_task.code_source_path config path to set to the uploaded workspace path. +func collectCodeSourceBlocks(b *bundle.Bundle) ([]blockSource, diag.Diagnostics) { + var sources []blockSource + var diags diag.Diagnostics + if len(b.Config.AiRuntimeExtras) == 0 { + return nil, diags + } + + // Read-only pattern walk over the current tree — the callback returns each task + // unchanged and only records matches, so it does not go through Config.Mutate. + _, err := dyn.MapByPattern(b.Config.Value(), taskPattern, func(p dyn.Path, task dyn.Value) (dyn.Value, error) { + jobName := p[2].Key() + taskKey, ok := task.Get("task_key").AsString() + if !ok { + return task, nil + } + extras, ok := b.Config.AiRuntimeExtras[jobName][taskKey] + if !ok || extras.CodeSource == nil { + return task, nil + } + sources = append(sources, blockSource{ + configPath: p.Append(dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")), + location: task.Get("ai_runtime_task").Location(), + opts: *extras.CodeSource, + }) + return task, nil + }) + if err != nil { + diags = diags.Extend(diag.FromErr(err)) + } + + return sources, diags +} + +// packageBlock packages a code_source block into a content-addressed tarball and +// returns its sync-relative path plus the archive bytes. Like packageOne it performs +// no disk or workspace write; the caller overlays the bytes and file sync uploads them. +// +// A git ref archives the pinned commit via `git archive` (uncommitted changes are not +// included); otherwise the working tree is packaged via the same file-list path as a +// bare code_source_path, narrowed by include_paths. +func packageBlock(ctx context.Context, b *bundle.Bundle, cs blockSource) (string, []byte, error) { + localDir := filepath.Join(b.SyncRootPath, filepath.FromSlash(cs.opts.RootPath)) + dirName := filepath.Base(localDir) + + if cs.opts.Git != nil { + commit, err := resolveGitCommit(ctx, localDir, cs.opts.Git, cs.opts.IncludePaths) + if err != nil { + return "", nil, err + } + archive, sha, err := buildGitArchive(ctx, localDir, commit, dirName, cs.opts.IncludePaths) + if err != nil { + return "", nil, err + } + relArchive := path.Join(snapshotSubdir, fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16])) + log.Debugf(ctx, "packaged git code snapshot %s for code_source.root_path %q at %s", relArchive, cs.opts.RootPath, shortSHA(commit)) + return relArchive, archive, nil + } + + relBase, err := filepath.Rel(b.SyncRootPath, localDir) + if err != nil { + return "", nil, fmt.Errorf("code_source.root_path %q: %w", cs.opts.RootPath, 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.root_path %q: %w", cs.opts.RootPath, err) + } + if len(cs.opts.IncludePaths) > 0 { + files = filterIncludePaths(files, relBase, cs.opts.IncludePaths) + } + // An empty file list means everything was filtered out (gitignore / sync.exclude / + // include_paths) or the directory is empty. Packaging it would deploy a job with no + // code, so fail with an actionable message instead. + if len(files) == 0 { + return "", nil, fmt.Errorf("code_source.root_path %q has no files to package (all excluded by .gitignore, sync.exclude, or include_paths, or the directory is empty)", cs.opts.RootPath) + } + + 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.root_path %q: %w", cs.opts.RootPath, err) + } + relArchive := path.Join(snapshotSubdir, fmt.Sprintf("%s_%s.tar.gz", dirName, sha[:16])) + log.Debugf(ctx, "packaged code snapshot %s for code_source.root_path %q", relArchive, cs.opts.RootPath) + return relArchive, buf.Bytes(), nil +} + +// filterIncludePaths keeps only files under one of the include paths. Each include is +// relative to the code directory (relBase); f.Relative is relative to the sync root, so +// the includes are re-based onto relBase before matching. +func filterIncludePaths(files []fileset.File, relBase string, includePaths []string) []fileset.File { + prefixes := make([]string, 0, len(includePaths)) + for _, inc := range includePaths { + inc = strings.Trim(strings.TrimSpace(inc), "/") + if relBase != "." { + inc = relBase + "/" + inc + } + prefixes = append(prefixes, inc) + } + return slices.DeleteFunc(files, func(f fileset.File) bool { + for _, pre := range prefixes { + // Match an exact file entry or anything under a directory entry. + if f.Relative == pre || strings.HasPrefix(f.Relative, pre+"/") { + return false + } + } + return true + }) +} + // 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 diff --git a/bundle/config/mutator/aicode/package_code_source_test.go b/bundle/config/mutator/aicode/package_code_source_test.go index fca33a0a7bf..1789d4a664c 100644 --- a/bundle/config/mutator/aicode/package_code_source_test.go +++ b/bundle/config/mutator/aicode/package_code_source_test.go @@ -10,6 +10,7 @@ import ( "github.com/databricks/cli/bundle/config/resources" "github.com/databricks/cli/bundle/internal/bundletest" "github.com/databricks/cli/libs/dyn" + "github.com/databricks/cli/libs/fileset" "github.com/databricks/databricks-sdk-go/service/iam" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/stretchr/testify/assert" @@ -87,3 +88,44 @@ func TestCollectLocalCodeSourcesSkipsLocalFile(t *testing.T) { require.Empty(t, diags) assert.Empty(t, sources, "a local tarball file must flow through artifact upload, not aicode packaging") } + +func TestFilterIncludePaths(t *testing.T) { + mk := func(rels ...string) []fileset.File { + out := make([]fileset.File, len(rels)) + for i, r := range rels { + out[i] = fileset.File{Relative: r} + } + return out + } + rels := func(files []fileset.File) []string { + out := make([]string, len(files)) + for i, f := range files { + out[i] = f.Relative + } + return out + } + + t.Run("relBase src, keep subtrees and exact files", func(t *testing.T) { + files := mk("src/a/x.py", "src/b/y.py", "src/c.py", "src/keep.txt") + got := filterIncludePaths(files, "src", []string{"a", "keep.txt"}) + assert.Equal(t, []string{"src/a/x.py", "src/keep.txt"}, rels(got)) + }) + + t.Run("relBase . keeps top-level includes", func(t *testing.T) { + files := mk("a/x.py", "b/y.py", "c.py") + got := filterIncludePaths(files, ".", []string{"a"}) + assert.Equal(t, []string{"a/x.py"}, rels(got)) + }) + + t.Run("trailing slash and whitespace tolerated", func(t *testing.T) { + files := mk("src/a/x.py", "src/b/y.py") + got := filterIncludePaths(files, "src", []string{" a/ "}) + assert.Equal(t, []string{"src/a/x.py"}, rels(got)) + }) + + t.Run("a prefix does not match a sibling with the same prefix string", func(t *testing.T) { + files := mk("src/a/x.py", "src/ab/y.py") + got := filterIncludePaths(files, "src", []string{"a"}) + assert.Equal(t, []string{"src/a/x.py"}, rels(got)) + }) +} diff --git a/bundle/config/mutator/aicode/snapshot_git.go b/bundle/config/mutator/aicode/snapshot_git.go new file mode 100644 index 00000000000..691d61f86e2 --- /dev/null +++ b/bundle/config/mutator/aicode/snapshot_git.go @@ -0,0 +1,212 @@ +package aicode + +import ( + "bytes" + "context" + "crypto/sha256" + "encoding/hex" + "errors" + "fmt" + "os/exec" + "path/filepath" + "strings" + + "github.com/databricks/cli/bundle/config" +) + +// Local, no-network git packaging for a code_source block that pins a git ref. +// Ported from experimental/air/cmd (snapshot_git.go / snapshot_resolve.go / +// snapshot_package.go), trimmed to what a bundle deploy needs: resolve the ref to a +// local commit, then `git archive` that commit into an in-memory gzipped tarball the +// caller overlays on the sync root (same as the working-tree snapshot). The +// remote-fetch and provenance-sidecar paths are intentionally not ported — a ref must +// resolve to a commit already present locally. + +// gitRepo runs git subcommands scoped to one repository via `git -C`. Arguments are +// passed as a slice, never a shell string, so a branch/commit value can't inject. +type gitRepo struct { + path string +} + +func newGitRepo(path string) gitRepo { + return gitRepo{path: path} +} + +// runBytes executes `git -C ` and returns raw stdout; a non-zero exit +// wraps stderr. +func (g gitRepo) runBytes(ctx context.Context, args ...string) ([]byte, error) { + full := append([]string{"-C", g.path}, args...) + cmd := exec.CommandContext(ctx, "git", full...) + + var stdout, stderr bytes.Buffer + cmd.Stdout = &stdout + cmd.Stderr = &stderr + if err := cmd.Run(); err != nil { + msg := strings.TrimSpace(stderr.String()) + if msg == "" { + return nil, fmt.Errorf("git %s: %w", strings.Join(args, " "), err) + } + return nil, fmt.Errorf("git %s: %w: %s", strings.Join(args, " "), err, msg) + } + return stdout.Bytes(), nil +} + +func (g gitRepo) run(ctx context.Context, args ...string) (string, error) { + out, err := g.runBytes(ctx, args...) + return string(out), err +} + +// isRepository reports whether the path is inside a git work tree. rev-parse (not a +// .git lookup) means a subdirectory of a repo counts — the common case when root_path +// is a subfolder of a monorepo. +func (g gitRepo) isRepository(ctx context.Context) bool { + out, err := g.run(ctx, "rev-parse", "--is-inside-work-tree") + if err != nil { + return false + } + return strings.TrimSpace(out) == "true" +} + +// hasUncommittedChanges reports whether there are staged or unstaged changes under the +// repo subtree. The `-- .` pathspec scopes the check so a subfolder snapshot considers +// only changes that could land in it. +func (g gitRepo) hasUncommittedChanges(ctx context.Context) (bool, error) { + out, err := g.run(ctx, "status", "--porcelain", "--", ".") + if err != nil { + return false, err + } + return strings.TrimSpace(out) != "", nil +} + +// resolveLocalBranchSHA resolves a branch to its local-HEAD commit. No remote is +// contacted; the branch must exist locally. +func (g gitRepo) resolveLocalBranchSHA(ctx context.Context, branch string) (string, error) { + out, err := g.run(ctx, "rev-parse", "refs/heads/"+branch) + if err != nil { + return "", fmt.Errorf("failed to resolve local branch %q; ensure the branch exists locally and root_path is correct: %w", branch, err) + } + return strings.TrimSpace(out), nil +} + +// commitExistsLocally reports whether commitSHA is in the local object store, without +// triggering a promisor/lazy fetch. +func (g gitRepo) commitExistsLocally(ctx context.Context, commitSHA string) bool { + _, err := g.run(ctx, "cat-file", "-e", commitSHA) + return err == nil +} + +// hasFilesAtCommit reports whether the commit has any files (blobs) under the given +// paths (whole tree when includePaths is empty). `-r` recurses so only blobs are +// listed, not tree entries — an empty result means there is nothing to package. +func (g gitRepo) hasFilesAtCommit(ctx context.Context, commitSHA string, includePaths []string) (bool, error) { + args := append([]string{"ls-tree", "-r", "--name-only", commitSHA}, includePaths...) + out, err := g.run(ctx, args...) + if err != nil { + return false, err + } + return strings.TrimSpace(out) != "", nil +} + +// validateIncludePathsExist checks that every include path exists at commitSHA. Without +// -d, `git ls-tree` reports both blobs and trees; empty output means the path is missing. +func (g gitRepo) validateIncludePathsExist(ctx context.Context, commitSHA string, includePaths []string) error { + var missing []string + for _, p := range includePaths { + out, err := g.run(ctx, "ls-tree", commitSHA, p) + if err != nil { + return err + } + if strings.TrimSpace(out) == "" { + missing = append(missing, p) + } + } + if len(missing) > 0 { + return fmt.Errorf("code_source.include_paths do not exist at commit %s: %s", shortSHA(commitSHA), strings.Join(missing, ", ")) + } + return nil +} + +// shortSHA abbreviates a commit SHA to 8 chars for messages, tolerating a +// user-supplied abbreviation shorter than that. +func shortSHA(sha string) string { + return sha[:min(len(sha), 8)] +} + +// resolveGitCommit resolves a code_source.git ref to a commit SHA that exists locally, +// verifying include_paths exist at that commit. repoPath is the local code directory. +// +// git.commit pins a committed SHA; a dirty working tree is irrelevant (not archived). +// git.branch archives the branch's local HEAD, so a dirty tree is an error: the +// committed HEAD would not include the uncommitted changes. +func resolveGitCommit(ctx context.Context, repoPath string, git *config.CodeSourceGit, includePaths []string) (string, error) { + repo := newGitRepo(repoPath) + // Forward-slash the path in user-facing messages so output is identical across + // operating systems (matches the repo's path-output convention). + displayPath := filepath.ToSlash(repoPath) + if !repo.isRepository(ctx) { + return "", fmt.Errorf("code_source.git is set but %s is not a git repository", displayPath) + } + + var commit string + switch { + case git.Commit != "": + if !repo.commitExistsLocally(ctx, git.Commit) { + return "", fmt.Errorf("commit %q does not exist locally; fetch it (e.g. `git fetch`) before deploying — the snapshot archives your local copy and does not fetch from a remote", git.Commit) + } + commit = git.Commit + case git.Branch != "": + dirty, err := repo.hasUncommittedChanges(ctx) + if err != nil { + return "", err + } + if dirty { + return "", fmt.Errorf("uncommitted changes under %s would not be included: code_source.git.branch deploys the committed HEAD of %q. Commit your changes, or use git.commit to pin a specific revision", displayPath, git.Branch) + } + sha, err := repo.resolveLocalBranchSHA(ctx, git.Branch) + if err != nil { + return "", err + } + commit = sha + default: + // aicode.Validate guarantees exactly one of branch/commit is set. + return "", errors.New("code_source.git requires either 'branch' or 'commit'") + } + + if len(includePaths) > 0 { + if err := repo.validateIncludePathsExist(ctx, commit, includePaths); err != nil { + return "", err + } + } + + // Fail on an empty tree rather than package a codeless archive and deploy a job + // with no code. Mirrors the working-tree path's len(files)==0 guard; the + // include_paths existence check above does not catch an otherwise-empty root_path. + hasFiles, err := repo.hasFilesAtCommit(ctx, commit, includePaths) + if err != nil { + return "", err + } + if !hasFiles { + return "", fmt.Errorf("code_source.root_path has no files to package at commit %s", shortSHA(commit)) + } + + return commit, nil +} + +// buildGitArchive produces a gzipped tarball of commitSHA via `git archive`, with every +// entry prefixed by prefix/ (so it expands to /..., matching the runtime's +// /databricks/code_source/ extraction contract). When includePaths is set only +// those paths are archived. It returns the archive bytes and their SHA-256 hex digest +// for content-addressed naming. The commit is deterministic, so an identical +// (commit, include_paths) yields an identical name and skips re-upload — but git's gzip +// output is not guaranteed byte-identical across git versions, so the name is stable +// only within a git version (a one-time re-upload after a git upgrade is acceptable). +func buildGitArchive(ctx context.Context, repoPath, commitSHA, prefix string, includePaths []string) ([]byte, string, error) { + args := []string{"archive", "--format=tar.gz", "--prefix=" + prefix + "/", commitSHA} + args = append(args, includePaths...) + out, err := newGitRepo(repoPath).runBytes(ctx, args...) + if err != nil { + return nil, "", fmt.Errorf("failed to create git archive: %w", err) + } + sum := sha256.Sum256(out) + return out, hex.EncodeToString(sum[:]), nil +} diff --git a/bundle/config/mutator/aicode/validate.go b/bundle/config/mutator/aicode/validate.go index 492203bc9f4..7ce45022882 100644 --- a/bundle/config/mutator/aicode/validate.go +++ b/bundle/config/mutator/aicode/validate.go @@ -5,6 +5,9 @@ import ( "fmt" "os" "path/filepath" + "regexp" + "slices" + "strings" "github.com/databricks/cli/bundle" "github.com/databricks/cli/bundle/config" @@ -69,6 +72,14 @@ func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics if task.AiRuntimeTask == nil { continue } + + // A DABs-native code_source block (extracted into AiRuntimeExtras before + // normalization) packages a directory too, so the snapshot-dir guards apply. + if block := codeSourceBlock(b, name, task.TaskKey); block != nil { + packagesCode = true + diags = diags.Extend(v.validateCodeSourceBlock(b, job.GitSource, task, taskPath, block)) + } + codePath := taskPath.Append(dyn.Key("ai_runtime_task"), dyn.Key("code_source_path")) diags = diags.Extend(v.validateTask(b, job.GitSource, task.AiRuntimeTask.CodeSourcePath, codePath)) } @@ -81,6 +92,140 @@ func (v *validate) Apply(ctx context.Context, b *bundle.Bundle) diag.Diagnostics return diags } +// codeSourceBlock returns the DABs-native code_source block stashed for a task, or +// nil if none. It is keyed by (job, task_key), matching rewriteAiRuntimeCodeSource. +func codeSourceBlock(b *bundle.Bundle, jobName, taskKey string) *config.CodeSourceOptions { + extras, ok := b.Config.AiRuntimeExtras[jobName][taskKey] + if !ok { + return nil + } + return extras.CodeSource +} + +// validateCodeSourceBlock checks a DABs-native ai_runtime_task.code_source block. +// The field-format rules mirror the AIR CLI's train.yaml validation +// (experimental/air/cmd/runconfig.go) so the two authoring surfaces agree. +func (v *validate) validateCodeSourceBlock(b *bundle.Bundle, gitSource *jobs.GitSource, task jobs.Task, taskPath dyn.Path, block *config.CodeSourceOptions) diag.Diagnostics { + blockPath := taskPath.Append(dyn.Key("ai_runtime_task"), dyn.Key("code_source")) + locations := b.Config.GetLocations(blockPath.String()) + reject := func(summary, detail string) diag.Diagnostics { + return diag.Diagnostics{{ + Severity: diag.Error, + Summary: summary, + Detail: detail, + Locations: locations, + Paths: []dyn.Path{blockPath}, + }} + } + + // code_source and code_source_path are two ways to set the same thing; requiring + // exactly one avoids an ambiguous precedence. + if task.AiRuntimeTask.CodeSourcePath != "" { + return reject( + "ai_runtime_task sets both code_source and code_source_path", + "Set either the code_source block or code_source_path, not both", + ) + } + + if strings.TrimSpace(block.RootPath) == "" { + return reject("ai_runtime_task.code_source.root_path is required", "") + } + + // The packaged directory must live inside the bundle sync root (it is uploaded as + // part of the bundle). Same check as the code_source_path path below. + if rel, err := filepath.Rel(b.SyncRootPath, filepath.Join(b.SyncRootPath, filepath.FromSlash(block.RootPath))); err != nil || !filepath.IsLocal(rel) { + return reject( + fmt.Sprintf("code_source.root_path %q is outside the bundle root", block.RootPath), + "root_path must point at a directory inside the bundle", + ) + } + isDir, err := isExistingDir(filepath.Join(b.SyncRootPath, filepath.FromSlash(block.RootPath))) + if err != nil { + return reject(fmt.Sprintf("failed to inspect code_source.root_path %q: %v", block.RootPath, err), "") + } + if !isDir { + return reject(fmt.Sprintf("code_source.root_path %q is not a directory", block.RootPath), "") + } + + // include_paths: relative, non-empty, no parent traversal (they stay within + // root_path). Mirrors snapshotSourceConfig.validate. + for _, p := range block.IncludePaths { + p = strings.TrimSpace(p) + if p == "" { + return reject("code_source.include_paths entry cannot be empty", "") + } + if strings.HasPrefix(p, "/") { + return reject(fmt.Sprintf("code_source.include_paths must be relative paths, got %q", p), "") + } + if slices.Contains(strings.Split(p, "/"), "..") { + return reject(fmt.Sprintf("code_source.include_paths cannot contain '..' traversal, got %q", p), "") + } + } + + // remote_volume would upload the archive to a UC Volume instead of the bundle's + // workspace file path. That is a deploy-phase workspace write (the overlay path this + // mutator uses only targets workspace.file_path), and a Volume is not cleaned by + // bundle destroy — both unresolved. Reject it explicitly rather than silently ignore. + if block.RemoteVolume != "" { + return reject( + "code_source.remote_volume is not yet supported", + "Remove remote_volume; the archive is uploaded to the bundle's workspace file path", + ) + } + + if block.Git != nil { + if diags := validateCodeSourceGit(block.Git, reject); diags.HasError() { + return diags + } + // The deploy engine retrieves task files from git when git_source is set, so a + // per-task git-pinned snapshot would be silently ignored. Reject the combination. + if gitSource != nil { + return reject( + "ai_runtime_task.code_source.git cannot be combined with the job's git_source", + "Remove git_source, or remove the code_source.git ref", + ) + } + } + + return nil +} + +// validateCodeSourceGit checks a code_source.git ref: exactly one of branch/commit, +// and a safe branch name. Mirrors gitRef.validate in the AIR CLI. +func validateCodeSourceGit(git *config.CodeSourceGit, reject func(summary, detail string) diag.Diagnostics) diag.Diagnostics { + switch { + case git.Branch == "" && git.Commit == "": + return reject("code_source.git requires either 'branch' or 'commit'", "") + case git.Branch != "" && git.Commit != "": + return reject("code_source.git 'branch' and 'commit' are mutually exclusive; specify only one", "") + } + if git.Branch != "" && !gitRefRe.MatchString(git.Branch) { + return reject( + fmt.Sprintf("invalid code_source.git.branch %q", git.Branch), + "only alphanumeric characters, hyphens, dots, slashes, and underscores are allowed", + ) + } + // commit must be a hex object id, not a moving ref. A branch/tag/HEAD value would + // package a different archive (and content-addressed name) on every deploy, + // defeating the point of pinning; requiring hex also rejects a leading-"-" value + // that would otherwise be a flag-injection shape into `git archive`. + if git.Commit != "" && !gitCommitRe.MatchString(git.Commit) { + return reject( + fmt.Sprintf("invalid code_source.git.commit %q", git.Commit), + "commit must be a hex commit SHA (7-64 hex characters), not a branch, tag, or HEAD", + ) + } + return nil +} + +// gitRefRe guards a branch name that flows into git exec args against injection. +// Mirrors the AIR CLI's gitRefRe (experimental/air/cmd/runconfig.go). +var gitRefRe = regexp.MustCompile(`^[\w./-]+$`) + +// gitCommitRe requires a hex object id: full SHA-1 (40) or SHA-256 (64), or an +// abbreviated prefix down to git's conventional 7-char short form. +var gitCommitRe = regexp.MustCompile(`^[0-9a-fA-F]{7,64}$`) + // packagesLocalDir reports whether codeSourcePath is one this mutator packages: a // local path that is an existing directory. A local *file* (a pre-built tarball from // an `artifacts` block) is uploaded by the artifact path instead, so it packages no diff --git a/bundle/config/mutator/aicode/validate_test.go b/bundle/config/mutator/aicode/validate_test.go index 8ab48060af1..d84162902cc 100644 --- a/bundle/config/mutator/aicode/validate_test.go +++ b/bundle/config/mutator/aicode/validate_test.go @@ -179,3 +179,131 @@ func TestValidateSnapshotGuardsSkippedWithoutLocalCodeSource(t *testing.T) { b.Config.Sync.Exclude = []string{".air_snapshots/*"} assert.Empty(t, Validate().Apply(t.Context(), b)) } + +// setCodeSourceBlock attaches a DABs-native code_source block to the "train" task, +// as rewriteAiRuntimeCodeSource would have done at load. +func setCodeSourceBlock(b *bundle.Bundle, block config.CodeSourceOptions) { + b.Config.AiRuntimeExtras = map[string]map[string]config.AiRuntimeTaskExtras{ + "train": {"train": {CodeSource: &block}}, + } +} + +func TestValidateCodeSourceBlock(t *testing.T) { + tests := []struct { + name string + block config.CodeSourceOptions + // codeSourcePath, when set, is also placed on the task to test the conflict. + codeSourcePath string + wantErr string // "" means no error expected + }{ + { + name: "valid working tree", + block: config.CodeSourceOptions{RootPath: "src"}, + }, + { + name: "valid include_paths", + block: config.CodeSourceOptions{RootPath: "src", IncludePaths: []string{"a", "b/c"}}, + }, + { + name: "valid git commit short", + block: config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{Commit: "abc1234"}}, + }, + { + name: "valid git commit full sha", + block: config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{Commit: "0123456789abcdef0123456789abcdef01234567"}}, + }, + { + name: "git commit too short", + block: config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{Commit: "abc123"}}, + wantErr: "invalid code_source.git.commit", + }, + { + name: "git commit is a moving ref", + block: config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{Commit: "HEAD"}}, + wantErr: "invalid code_source.git.commit", + }, + { + name: "git commit non-hex (branch name)", + block: config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{Commit: "main"}}, + wantErr: "invalid code_source.git.commit", + }, + { + name: "git commit injection shape", + block: config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{Commit: "--output=/tmp/x"}}, + wantErr: "invalid code_source.git.commit", + }, + { + name: "conflict with code_source_path", + block: config.CodeSourceOptions{RootPath: "src"}, + codeSourcePath: "src", + wantErr: "both code_source and code_source_path", + }, + { + name: "empty root_path", + block: config.CodeSourceOptions{RootPath: ""}, + wantErr: "root_path is required", + }, + { + name: "root_path outside bundle", + block: config.CodeSourceOptions{RootPath: "../escape"}, + wantErr: "outside the bundle root", + }, + { + name: "include_paths absolute", + block: config.CodeSourceOptions{RootPath: "src", IncludePaths: []string{"/abs"}}, + wantErr: "must be relative paths", + }, + { + name: "include_paths traversal", + block: config.CodeSourceOptions{RootPath: "src", IncludePaths: []string{"a/../b"}}, + wantErr: "'..' traversal", + }, + { + name: "remote_volume not supported", + block: config.CodeSourceOptions{RootPath: "src", RemoteVolume: "/Volumes/main/x"}, + wantErr: "remote_volume is not yet supported", + }, + { + name: "git branch and commit both set", + block: config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{Branch: "main", Commit: "abc"}}, + wantErr: "mutually exclusive", + }, + { + name: "git neither branch nor commit", + block: config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{}}, + wantErr: "requires either 'branch' or 'commit'", + }, + { + name: "git invalid branch", + block: config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{Branch: "bad;rm -rf"}}, + wantErr: "invalid code_source.git.branch", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + b := bundleForValidate(t, tt.codeSourcePath, nil) + mkCodeDir(t, b, "src") + setCodeSourceBlock(b, tt.block) + + diags := Validate().Apply(t.Context(), b) + if tt.wantErr == "" { + assert.Empty(t, diags) + return + } + require.True(t, diags.HasError(), "expected an error") + assert.Contains(t, diags[0].Summary+" "+diags[0].Detail, tt.wantErr) + }) + } +} + +// A code_source.git ref combined with the job's git_source is rejected: the deploy +// engine would fetch task files from git_source and ignore the packaged snapshot. +func TestValidateCodeSourceBlockGitSourceConflict(t *testing.T) { + b := bundleForValidate(t, "", &jobs.GitSource{GitUrl: "https://example.invalid/repo"}) + mkCodeDir(t, b, "src") + setCodeSourceBlock(b, config.CodeSourceOptions{RootPath: "src", Git: &config.CodeSourceGit{Commit: "abc1234"}}) + diags := Validate().Apply(t.Context(), b) + require.True(t, diags.HasError()) + assert.Contains(t, diags[0].Summary, "cannot be combined with the job's git_source") +} diff --git a/bundle/config/root.go b/bundle/config/root.go index e13bf78bf15..8bf6bc35ec3 100644 --- a/bundle/config/root.go +++ b/bundle/config/root.go @@ -5,8 +5,11 @@ import ( "context" "errors" "fmt" + "maps" "os" + "path/filepath" "reflect" + "slices" "strings" "github.com/databricks/cli/bundle/config/resources" @@ -89,6 +92,14 @@ type Root struct { //nolint:recvcheck // value receivers for read-only accessors // information for every path in the configuration tree. Locations *dynloc.Locations `json:"__locations,omitempty" bundle:"internal"` + // AiRuntimeExtras holds ai_runtime_task authoring sugar (currently the + // DABs-native code_source block) that has no field on the SDK jobs.AiRuntimeTask. + // rewriteAiRuntimeCodeSource extracts it before normalization — which would + // otherwise drop it as an unknown field — keyed by job name then task_key, so it + // survives task reordering (MergeJobTasks sorts by task_key) and target/include + // merges. Internal: not user-set directly, not in the JSON schema. + AiRuntimeExtras map[string]map[string]AiRuntimeTaskExtras `json:"__ai_runtime_task_extras,omitempty" bundle:"internal"` + Scripts map[string]Script `json:"scripts,omitempty"` // Python configures loading of Python code defined with 'databricks-bundles' package. @@ -130,6 +141,14 @@ func LoadFromBytes(path string, raw []byte) (*Root, diag.Diagnostics) { return nil, diag.Errorf("failed to rewrite %s: %v", path, err) } + // Extract ai_runtime_task.code_source blocks before normalization would drop them + // as unknown fields on the SDK jobs.AiRuntimeTask. + v, err = rewriteAiRuntimeCodeSource(v) + if err != nil { + // Forward-slash the file path so the error is byte-identical across OSes. + return nil, diag.Errorf("failed to rewrite %s: %v", filepath.ToSlash(path), err) + } + // Normalize dynamic configuration tree according to configuration type. v, diags := convert.Normalize(r, v) @@ -551,6 +570,135 @@ func rewriteShorthands(v dyn.Value) (dyn.Value, error) { })) } +// aiRuntimeCodeSourceKey is the DABs-native block nested under ai_runtime_task that +// rewriteAiRuntimeCodeSource extracts; it is not a field on the SDK jobs.AiRuntimeTask. +const aiRuntimeCodeSourceKey = "code_source" + +// aiRuntimeExtrasKey is the internal top-level key the extracted blocks are stashed +// under; it maps to Root.AiRuntimeExtras. +const aiRuntimeExtrasKey = "__ai_runtime_task_extras" + +// slashLocation renders a dyn.Location with a forward-slash file path so error +// messages are identical across operating systems (acceptance output is stable on +// Windows). dyn.Location.String() emits the OS-native path verbatim. +func slashLocation(l dyn.Location) string { + l.File = filepath.ToSlash(l.File) + return l.String() +} + +// rewriteAiRuntimeCodeSource moves each `ai_runtime_task.code_source` block out of the +// SDK-typed task tree into the internal __ai_runtime_task_extras stash, keyed by job +// name then task_key, and removes the code_source key from under ai_runtime_task. +// +// It must run before convert.Normalize: code_source is not a field on the SDK +// jobs.AiRuntimeTask, so normalization would drop it (with an "unknown field" warning) +// from both the typed config and the stored dynamic tree. Keying by task_key (not +// position) keeps the stash aligned with the task after MergeJobTasks reorders tasks +// by sorted task_key. The stash is a real (internal) config field, so it rides the +// normal per-file merge and ToTyped without extra wiring. +// +// A code_source block under a targets.* override is rejected: the (job, task_key) key +// cannot represent a per-target override without collision, and this rewrite runs +// before target selection. +func rewriteAiRuntimeCodeSource(v dyn.Value) (dyn.Value, error) { + if v.Kind() != dyn.KindMap { + return v, nil + } + + // Reject a code_source authored inside a target override, before extracting the + // top-level ones (so the error fires regardless of whether a top-level block exists). + targetPattern := dyn.NewPattern( + dyn.Key("targets"), dyn.AnyKey(), + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + dyn.Key("ai_runtime_task"), + ) + var targetErr error + _, err := dyn.MapByPattern(v, targetPattern, func(p dyn.Path, task dyn.Value) (dyn.Value, error) { + if cs := task.Get(aiRuntimeCodeSourceKey); cs.Kind() != dyn.KindInvalid { + targetErr = fmt.Errorf("ai_runtime_task.code_source at %s is not supported inside a target override; move it to the task's base definition under resources.jobs", slashLocation(cs.Location())) + } + return task, nil + }) + if err != nil { + return dyn.InvalidValue, err + } + if targetErr != nil { + return dyn.InvalidValue, targetErr + } + + // extracted[jobName][taskKey] = the code_source block value. + extracted := map[string]map[string]dyn.Value{} + + taskPattern := dyn.NewPattern( + dyn.Key("resources"), dyn.Key("jobs"), dyn.AnyKey(), + dyn.Key("tasks"), dyn.AnyIndex(), + ) + out, err := dyn.MapByPattern(v, taskPattern, func(p dyn.Path, task dyn.Value) (dyn.Value, error) { + airt := task.Get("ai_runtime_task") + if airt.Kind() != dyn.KindMap { + return task, nil + } + cs := airt.Get(aiRuntimeCodeSourceKey) + if cs.Kind() == dyn.KindInvalid { + return task, nil + } + + // A task with a code_source block must have a string task_key so the stash can + // be keyed by it (and stay aligned after MergeJobTasks reorders by task_key). + taskKey, ok := task.Get("task_key").AsString() + if !ok || taskKey == "" { + return dyn.InvalidValue, fmt.Errorf("ai_runtime_task.code_source at %s requires the task to set a task_key", slashLocation(cs.Location())) + } + + // p is resources.jobs..tasks[]; the job name is component 2. + jobName := p[2].Key() + if extracted[jobName] == nil { + extracted[jobName] = map[string]dyn.Value{} + } + // Catches a duplicate (job, task_key) within this file only — the rewrite runs + // per file, before includes/targets are merged. A collision across files is + // resolved by the normal AiRuntimeExtras map merge (last block wins), the same + // as any other duplicated resource field. + if _, dup := extracted[jobName][taskKey]; dup { + return dyn.InvalidValue, fmt.Errorf("duplicate task_key %q in job %q with an ai_runtime_task.code_source block", taskKey, jobName) + } + extracted[jobName][taskKey] = cs + + // Drop the code_source key from ai_runtime_task so normalization sees only real + // SDK fields (no "unknown field" warning). Rebuild the mapping without it. + newAirt := dyn.NewMapping() + for _, pair := range airt.MustMap().Pairs() { + if pair.Key.MustString() == aiRuntimeCodeSourceKey { + continue + } + newAirt.SetLoc(pair.Key.MustString(), pair.Key.Locations(), pair.Value) + } + return dyn.Set(task, "ai_runtime_task", dyn.NewValue(newAirt, airt.Locations())) + }) + if err != nil { + return dyn.InvalidValue, err + } + if len(extracted) == 0 { + return out, nil + } + + // Build the stash tree __ai_runtime_task_extras...code_source and set + // it on the root. Deterministic key order keeps the dynamic tree stable. + jobsMapping := dyn.NewMapping() + for _, jobName := range slices.Sorted(maps.Keys(extracted)) { + tasksMapping := dyn.NewMapping() + for _, taskKey := range slices.Sorted(maps.Keys(extracted[jobName])) { + cs := extracted[jobName][taskKey] + extras := dyn.NewMapping() + extras.SetLoc(aiRuntimeCodeSourceKey, cs.Locations(), cs) + tasksMapping.SetLoc(taskKey, cs.Locations(), dyn.NewValue(extras, cs.Locations())) + } + jobsMapping.SetLoc(jobName, nil, dyn.NewValue(tasksMapping, nil)) + } + return dyn.Set(out, aiRuntimeExtrasKey, dyn.NewValue(jobsMapping, nil)) +} + // validateVariableOverrides checks that all variables specified // in the target override are also defined in the root. func validateVariableOverrides(root, target dyn.Value) (err error) { diff --git a/bundle/config/root_test.go b/bundle/config/root_test.go index 42fae49d98c..36f6590ac0a 100644 --- a/bundle/config/root_test.go +++ b/bundle/config/root_test.go @@ -31,6 +31,92 @@ func TestRootLoad(t *testing.T) { assert.Equal(t, "basic", root.Bundle.Name) } +func TestRewriteAiRuntimeCodeSourceExtractsBlock(t *testing.T) { + yml := []byte(` +resources: + jobs: + train: + tasks: + - task_key: t1 + ai_runtime_task: + experiment: my-exp + code_source: + root_path: ./src + include_paths: [src, configs] + git: + commit: abc123 +`) + root, diags := LoadFromBytes("databricks.yml", yml) + require.NoError(t, diags.Error()) + + // The block is stashed keyed by (job, task_key). + extras, ok := root.AiRuntimeExtras["train"]["t1"] + require.True(t, ok, "expected code_source stashed under train/t1") + require.NotNil(t, extras.CodeSource) + assert.Equal(t, "./src", extras.CodeSource.RootPath) + assert.Equal(t, []string{"src", "configs"}, extras.CodeSource.IncludePaths) + require.NotNil(t, extras.CodeSource.Git) + assert.Equal(t, "abc123", extras.CodeSource.Git.Commit) + + // The code_source key is removed from the task, and no unknown-field warning fired. + task := root.Resources.Jobs["train"].Tasks[0] + require.NotNil(t, task.AiRuntimeTask) + assert.Equal(t, "my-exp", task.AiRuntimeTask.Experiment) + assert.Empty(t, diags, "expected no diagnostics (no unknown-field warning)") +} + +func TestRewriteAiRuntimeCodeSourceLeavesCodeSourcePathAlone(t *testing.T) { + yml := []byte(` +resources: + jobs: + train: + tasks: + - task_key: t1 + ai_runtime_task: + experiment: my-exp + code_source_path: ./src +`) + root, diags := LoadFromBytes("databricks.yml", yml) + require.NoError(t, diags.Error()) + assert.Nil(t, root.AiRuntimeExtras) + assert.Equal(t, "./src", root.Resources.Jobs["train"].Tasks[0].AiRuntimeTask.CodeSourcePath) +} + +func TestRewriteAiRuntimeCodeSourceRejectsMissingTaskKey(t *testing.T) { + yml := []byte(` +resources: + jobs: + train: + tasks: + - ai_runtime_task: + experiment: my-exp + code_source: + root_path: ./src +`) + _, diags := LoadFromBytes("databricks.yml", yml) + require.Error(t, diags.Error()) + assert.Contains(t, diags.Error().Error(), "requires the task to set a task_key") +} + +func TestRewriteAiRuntimeCodeSourceRejectsTargetOverride(t *testing.T) { + yml := []byte(` +targets: + dev: + resources: + jobs: + train: + tasks: + - task_key: t1 + ai_runtime_task: + experiment: my-exp + code_source: + root_path: ./src +`) + _, diags := LoadFromBytes("databricks.yml", yml) + require.Error(t, diags.Error()) + assert.Contains(t, diags.Error().Error(), "not supported inside a target override") +} + func TestInitializeVariables(t *testing.T) { fooDefault := "abc" root := &Root{ diff --git a/bundle/internal/schema/main.go b/bundle/internal/schema/main.go index 3f468852a83..82615403e5e 100644 --- a/bundle/internal/schema/main.go +++ b/bundle/internal/schema/main.go @@ -116,6 +116,48 @@ func removeJobsFields(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema { return s } +// addAiRuntimeCodeSource injects the DABs-native `code_source` block into the +// jobs.AiRuntimeTask schema. code_source is not a field on the SDK type — it is +// extracted from the config tree before normalization (see +// rewriteAiRuntimeCodeSource) and lowered into code_source_path at deploy — so it +// would otherwise be rejected by the type's additionalProperties: false. The schema is +// hand-authored (mirrors config.CodeSourceOptions) so editors document and validate it. +func addAiRuntimeCodeSource(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema { + if typ != reflect.TypeFor[jobs.AiRuntimeTask]() { + return s + } + if s.Properties == nil { + s.Properties = map[string]*jsonschema.Schema{} + } + strSchema := func(desc string) *jsonschema.Schema { + return &jsonschema.Schema{Type: jsonschema.StringType, Description: desc} + } + s.Properties["code_source"] = &jsonschema.Schema{ + Type: jsonschema.ObjectType, + Description: "DABs-native code source for this task. At deploy the CLI packages the directory (honoring .gitignore and sync include/exclude, plus git ref and include_paths when set), uploads it, and sets code_source_path. Mutually exclusive with code_source_path.", + AdditionalProperties: false, + Required: []string{"root_path"}, + Properties: map[string]*jsonschema.Schema{ + "root_path": strSchema("Local directory to package, relative to the bundle sync root."), + "include_paths": { + Type: jsonschema.ArrayType, + Description: "Package only these subtrees of root_path (relative paths, no \"..\") instead of the whole directory.", + Items: strSchema(""), + }, + "git": { + Type: jsonschema.ObjectType, + Description: "Pin the snapshot to a committed git revision instead of the working tree. Specify exactly one of branch or commit.", + AdditionalProperties: false, + Properties: map[string]*jsonschema.Schema{ + "branch": strSchema("Package the local HEAD of this branch. Requires a clean working tree."), + "commit": strSchema("Package this commit (must exist locally)."), + }, + }, + }, + } + return s +} + func removePipelineFields(typ reflect.Type, s jsonschema.Schema) jsonschema.Schema { switch typ { case reflect.TypeFor[resources.Pipeline](): @@ -282,6 +324,7 @@ func generateSchema(workdir, outputFile, cliJSONFile string, docsMode bool) { removePipelineFields, removeDeploymentFields, makeVolumeTypeOptional, + addAiRuntimeCodeSource, a.addAnnotations, removeOutputOnlyFields, } diff --git a/bundle/schema/jsonschema.json b/bundle/schema/jsonschema.json index 2c4827b141d..23f351ad0c2 100644 --- a/bundle/schema/jsonschema.json +++ b/bundle/schema/jsonschema.json @@ -6977,6 +6977,42 @@ "type": "object", "description": "AiRuntimeTask: multi-node GPU compute task definition for Databricks AI\nRuntime workloads.\n\nJobs-framework-level concepts (retries, per-task timeout, idempotency\ntoken, usage/budget policy, permissions) live on the surrounding\nTaskSettings / run-submit request and are intentionally NOT duplicated\nhere. Users compose `ai_runtime_task` with the standard Jobs/DABs task\nwrapper to get those.", "properties": { + "code_source": { + "type": "object", + "description": "DABs-native code source for this task. At deploy the CLI packages the directory (honoring .gitignore and sync include/exclude, plus git ref and include_paths when set), uploads it, and sets code_source_path. Mutually exclusive with code_source_path.", + "properties": { + "git": { + "type": "object", + "description": "Pin the snapshot to a committed git revision instead of the working tree. Specify exactly one of branch or commit.", + "properties": { + "branch": { + "type": "string", + "description": "Package the local HEAD of this branch. Requires a clean working tree." + }, + "commit": { + "type": "string", + "description": "Package this commit (must exist locally)." + } + }, + "additionalProperties": false + }, + "include_paths": { + "type": "array", + "description": "Package only these subtrees of root_path (relative paths, no \"..\") instead of the whole directory.", + "items": { + "type": "string" + } + }, + "root_path": { + "type": "string", + "description": "Local directory to package, relative to the bundle sync root." + } + }, + "additionalProperties": false, + "required": [ + "root_path" + ] + }, "code_source_path": { "description": "[Private Preview] Workspace or UC volume path of the code-source archive, unpacked on\neach node and exposed through `$CODE_SOURCE`. Set by first-party\ntooling; not for direct callers.", "$ref": "#/$defs/string", diff --git a/pyrefly.toml b/pyrefly.toml new file mode 100644 index 00000000000..2d50edae1cb --- /dev/null +++ b/pyrefly.toml @@ -0,0 +1,11 @@ +project-includes = [ + "**/*.py" +] + +project-excludes = [] + +search-path = [] + +disable-search-path-heuristics = true +ignore-missing-imports = ["*"] +ignore-errors-in-generated-code = true diff --git a/python/codegen/codegen/jsonschema.py b/python/codegen/codegen/jsonschema.py index e0e57eb60d4..6cac8a4a695 100644 --- a/python/codegen/codegen/jsonschema.py +++ b/python/codegen/codegen/jsonschema.py @@ -98,13 +98,18 @@ def _parse_bool(value) -> Optional[bool]: return value for k, v in schema.get("properties", {}).items(): + # Properties without a $ref are inline schemas the Go side injects for + # DABs-only authoring sugar that has no SDK type (e.g. ai_runtime_task's + # code_source block). pydabs has no model for them, so skip rather than + # assert — the field is available in YAML DABs but not the Python API. + if v.get("$ref") is None: + continue + assert v.get("type") is None assert v.get("anyOf") is None assert v.get("properties") is None assert v.get("items") is None - assert v.get("$ref") - prop = Property( ref=v["$ref"], description=v.get("description"),