air: add convert-to-dabs (run YAML -> Databricks Asset Bundle) - #6102
air: add convert-to-dabs (run YAML -> Databricks Asset Bundle)#6102vinchenzo-db wants to merge 17 commits into
Conversation
- Remove a hardcoded email from render_test.go (use user@example.com). - list_tui.go: use the existing cmdio.IsPromptSupported instead of the air-added IsPagerSupported; drop IsPagerSupported from libs/cmdio/io.go since it is no longer used. - format.go: standardize on termenv.Hyperlink and drop the hand-rolled osc8Link helper (and its test). - Centralize EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] into a shared acceptance/experimental/air/test.toml; remove the per-dir duplication (deleting the test.toml files left with nothing else). Co-authored-by: Isaac
## Changes Ports `--override KEY=VALUE` from the Python CLI. Overrides apply to the parsed YAML map before re-decode + validate, so path existence, type coercion, and semantic validate() rules all run. A reflection-based path check names the exact --override key and lists available fields on error. ## Why --override allows users to tweak a training configuration at launch time without editing the file and this allows for ease of use with sweep hyperparameters or scalable compute without having to maintain a forked configuration for each run. ## Tests unit tests: - TestParseOverrides: parsing KEY=VALUE - TestValidateOverridePaths: dotted-path validation against the schema - TestLoadRunConfigWithOverrides: end-to-end through the loader: scalar coercion, multiple overrides, free-form env-var-as-string, auto-created intermediate maps, unknown-path rejection, semantic re-validation after override, type mismatch, malformed override. - TestSubmitWorkload: the harness pattern being extended to prove overrides reach the actual POST body sent to /api/2.2/jobs/runs/submit, verified against the in-process testserver that records the request. - TestSubmitWorkloadHonorsOverride: proves a --override actually changes what gets sent to the Jobs API on a real submit, not just during dry-run validation. acceptance tests: - successful override that logs the change and then validates - an unknown field override that errors with an actionable message (see screenshots below) - override that passes type checking but fails schema validation (so we know that validate() still runs and works properly) Manual verification: <img width="1894" height="888" alt="Screenshot 2026-07-14 at 10 41 22 AM" src="https://github.com/user-attachments/assets/69187c62-ad60-4144-9744-79e52787447c" />
…fig (#5968) Post-merge cleanup for the experimental AIR CLI (follow-up to #5847, which squash-merged `air-cli` into `main`). This commit was made after that merge, so it is not yet in `main`. - Remove a hardcoded email from `render_test.go` (use `user@example.com`). - `list_tui.go`: use the existing `cmdio.IsPromptSupported` instead of the air-added `IsPagerSupported`; drop `IsPagerSupported` from `libs/cmdio/io.go` since it is no longer used. - `format.go`: standardize on `termenv.Hyperlink` and drop the hand-rolled `osc8Link` helper (and its test). - Centralize `EnvMatrix.DATABRICKS_BUNDLE_ENGINE = []` into a shared `acceptance/experimental/air/test.toml`; remove the per-dir duplication. Stacked below the `air logs` port branch (`air-logs-m4`), which depends on the centralized `test.toml` introduced here. This pull request and its description were written by Isaac. --------- Signed-off-by: Lennart Kats <lennart.kats@databricks.com> Co-authored-by: radakam <55745584+radakam@users.noreply.github.com> Co-authored-by: Lennart Kats (databricks) <lennart.kats@databricks.com> Co-authored-by: Jan N Rose <janniklas.rose@gmail.com> Co-authored-by: Grigory Panov <grigory.panov@databricks.com> Co-authored-by: Andrew Nester <andrew.nester.dev@gmail.com> Co-authored-by: Pieter Noordhuis <pieter.noordhuis@databricks.com>
## Changes Implements the air logs JOB_RUN_ID command (previously a notImplemented stub) for the experimental AIR CLI. It fetches a run's training logs with a Bricklens-first, MLflow-fallback strategy: - Bricklens (primary): streams logs from the AiTraining log endpoint following an active run to completion, or tailing a completed one. - MLflow (fallback): when Bricklens is unavailable & gated off by a backend flag (FEATURE_DISABLED), not deployed (ENDPOINT_NOT_FOUND/404), or persistently failing. The command falls back to reading the run's MLflow log artifacts (chunk discovery + credential-vended download). The flag is evaluated server-side; the CLI only reads the response error code. Flags: - `--minutes` restricts the fetch to the last N minutes (Bricklens time window). - `--lines <N>` is the tail the last N lines of a completed run. Mutually exclusive with --minutes. - `--node`, `--retry` is used to select a node / retry attempt. A past retry of a still-active run renders its (immutable) logs once instead of following the run. - `--download-to` / `--review` are rejected with a clear "not implemented" error (next PR follow up). ## Why `air logs` was the last unimplemented read command in the AIR CLI port. Bricklens is the primary log source, but it's behind a backend feature flag and isn't universally deployed, so the command must degrade gracefully to MLflow artifacts rather than fail. Bricklens is time-indexed (hence --minutes), while MLflow stores fixed log chunks (hence line-based --lines). the two flags map to what each backend can actually do. ## Tests - Unit: classifyLogError fallback classification, --minutes/--lines window math, run-status projection, bounded dedup set, page draining + tail ordering, MLflow chunk listing/sorting and log-path discovery, flag validation, and an end-to-end Bricklens MLflow fallback through a mock server. - Acceptance: acceptance/experimental/air/logs/ (text + JSON streaming, --minutes, --lines, --retry, mutual-exclusion error, invalid ID, --download-to rejection) and logs-mlflow-fallback/ (Bricklens FEATURE_DISABLED, MLflow fallback, no-logs, text + JSON). <img width="1628" height="860" alt="Screenshot 2026-07-24 at 4 14 21 PM" src="https://github.com/user-attachments/assets/1f70943a-8a5d-47f2-8797-ec6f3e91a8fe" />
…6080) ## Changes & Why After submitting a workload, --watch follows the run's logs to completion and exits with the run's outcome, reusing the same Bricklens-with-MLflow-fallback pipeline as `air logs`. - Text mode: prints "Submitted run", the dashboard link, and "Monitoring run and streaming logs...", then streams the logs. - JSON mode: emits a SUBMITTED event with the run id, a STATUS event on each lifecycle transition, the streamed LOG/ALERT events, and a closing terminal-status envelope (SUCCESS/FAILED/CANCELED) — matching the Python CLI's --watch JSONL contract. - Without --watch, the plain submit path now prints a tip about --watch. - STATUS events are watch-scoped (opt-in via logRequest.onStatusChange), so the merged `air logs` output is unchanged. --dry-run still takes precedence over --watch (nothing is submitted or streamed). ## Tests Unit tests (experimental/air/cmd/) - logbricklens_test.go: Bricklens client query/path serialization + time_unix_nano parsing - logstream_test.go: fallback classification, status projection, --minutes/tail math, page dedup/ordering, retry-then-fallback, JSONL/ALERT emit, Ctrl-C exit - logmlflow_test.go: MLflow chunk discovery/listing, attempt-prefix layout, no-logs exit-code parity - logs_test.go: air logs command: flag validation, completed-run tail, Bricklens→MLflow fallback, past-retry static view - run_watch_test.go: air run --watch: text stream, JSON SUBMITTED→STATUS→LOG→terminal envelope, failed-run exit code, dry-run precedence Acceptance tests (acceptance/experimental/air/) - logs/ : text/JSON streaming, --minutes, --lines, --lines 0, --retry, mutual-exclusion errors, invalid id, negative node, --download-to - logs-mlflow-fallback/ : Bricklens FEATURE_DISABLED → MLflow fallback → no-logs (text & JSON) - run/ : dry-run, --override, config validation, --watch ignored under --dry-run - run-submit/ : real submit payload + --watch tip line - help/ : air --help, air logs --help command-tree pins Manual verification: Properly monitors and outputs logs from runs on manual test: <img width="1166" height="112" alt="Screenshot 2026-07-27 at 3 00 15 PM" src="https://github.com/user-attachments/assets/f4522965-f200-410c-8ddc-307092c6b731" />
Brings the air-cli feature branch current with main. air-cli had drifted 110 commits behind, spanning Grigory Panov's localenv/environments redesign (cmd/localenv renamed to cmd/environments, new JobTaskEnvironment API, --cluster-name/--job-task flags, uv provisioning) and other infra work. Conflict resolution rule: - Infra (libs/localenv, cmd/environments, libs/filer, bundle/config/validate) and their acceptance goldens: take main. air-cli only carried an older snapshot of this shared code (via the #5968 squash); it had no AIR-specific edits there, so main is authoritative. - AIR (experimental/air/**, acceptance/experimental/air/**): keep air-cli. This is Riddhi's AIR CLI work and the reason the branch exists. The per-dir air test.toml deletions are honored (engine matrix centralized in the shared acceptance/experimental/air/test.toml). - libs/localenv/target_test.go aligned to main to match main's target.go API. Verified: go build ./... clean; 123 packages pass across libs/, cmd/, bundle/config/ (0 failures); localenv acceptance green. Co-authored-by: Isaac
The main catch-up merge brought in a new acceptance validator that rejects EnvMatrix.DATABRICKS_BUNDLE_ENGINE = [] (it would run on both direct and terraform CI runners). The shared acceptance/experimental/air/test.toml still had [], which broke every air acceptance test on air-cli post-merge. Pin to ["direct"] as the validator directs: no air command deploys a bundle, so running once on a single engine is the intent (an empty list is now disallowed). Co-authored-by: Isaac
…ommand The main catch-up merge renamed the local-env command to `environments setup-local` (--job-task <id>.<key> model). Three air-cli-only tests (job-ambiguous-compute, job-multicluster-mismatch, job-serverless-version-mismatch) still invoked the removed `local-env python sync --job ... --check` surface, so they broke post-merge with `unknown command "local-env"`. Delete them: they targeted the pre-redesign CLI, and main's newer localenv suite already covers these cases (cluster-name-ambiguous, job-task-missing-key, job-task-jobcluster, serverless-check, etc.) via the new command surface. Co-authored-by: Isaac
Follow-up to the earlier shared test.toml fix ([] -> ["direct"]): the derived out.test.toml for the logs and logs-mlflow-fallback dirs still recorded the old [] value, so CI's "changed files" check (git diff --exit-code after regenerating out.test.toml) failed. Regenerate them to match. Co-authored-by: Isaac
Package the code_source working tree into a tarball and upload it
through DABs' artifact-upload plumbing (libraries.ReplaceWithRemotePath
+ libraries.Upload over a minimal in-memory bundle), rewriting
ai_runtime_task.code_source_path to the uploaded remote path. The
packaging + upload orchestration is CLI-owned (experimental/air/cmd,
OWNERS = us); it only reuses DABs' uploader so we don't reimplement
workspace/volume upload.
snapshot_dabs.go: build the plain-tar tarball (createPlainTarball),
carry it as a file-valued code_source_path on a minimal bundle, and
drive the DABs upload. runsubmit.go swaps the old raw-filer snapshot
upload for this. Removes the retired raw-filer upload path (snapshot.go
uploader, snapshot_test.go).
Tar snapshotting only; git pinning follows in the next PR (its git
helpers are removed here and reintroduced there).
Co-authored-by: Isaac
## Changes
<!-- Brief summary of your changes that is easy to understand -->
## Why
<!-- Why are these changes needed? Provide the context that the reviewer
might be missing.
For example, were there any decisions behind the change that are not
reflected in the code itself? -->
## Testing
### Unit + acceptance
`experimental/air/cmd/...` and `acceptance/experimental/air/run-submit`
— working-tree,
git-pinned, and remote-Volume submits each assert the tarball lands
under `.internal/`
and the rewritten `code_source_path` rides the submitted
`ai_runtime_task`; plus the tar
builders (`.gitignore`, `.git` exclusion, `include_paths`) and the
no-`code_source`
nil-guard. All green.
### Live E2E — staging `dbc-04ac0685-8857` (GPU_1xA10)
5/5 runs SUCCESS, one per packaging mode. All runs are `CAN_VIEW` for
the workspace
`users` group, so every link below is openable by anyone in the
workspace.
**Setup** — a tiny project with a gitignored file (`debug.log`) to prove
exclusion:
```bash
mkdir -p /tmp/air-demo/proj/src/pkg && cd /tmp/air-demo/proj
printf 'import os\nprint("train ran; cwd:", os.getcwd(), "CODE_SOURCE_PATH:", os.environ.get("CODE_SOURCE_PATH"))\n' > src/train.py
echo 'def helper(): return 1' > src/pkg/util.py
echo '*.log' > src/.gitignore
echo 'noise' > src/debug.log # gitignored — must never be uploaded
cat > wt.yaml <<'YAML'
experiment_name: vchen_demo_wt
command: cd $CODE_SOURCE_PATH && python train.py
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
code_source: {type: snapshot, snapshot: {root_path: src}}
YAML
```
**1. Working-tree tarball** — plain tar of the working tree, honoring
`.gitignore`.
Run:
[994765091508414](https://dbc-04ac0685-8857.staging.cloud.databricks.com/jobs/runs/994765091508414)
```bash
dbcli experimental air run -f wt.yaml -p dbc-04ac0685-8857
# Uploading src.tar.gz... → Submitted run 994765091508414
dbcli workspace export /Workspace/Users/v.chen@databricks.com/.air/repo_snapshots/.internal/src.tar.gz --file /tmp/wt.tar.gz -p dbc-04ac0685-8857
tar tzf /tmp/wt.tar.gz
# → src/train.py, src/pkg/util.py, src/.gitignore (NO debug.log ✓ gitignore honored)
```
**2. Git-pinned commit** — `git archive` of a pinned SHA. An uncommitted
file is created
*after* the commit to prove the archive captures the commit, not the
dirty working tree.
Run:
[304463075281818](https://dbc-04ac0685-8857.staging.cloud.databricks.com/jobs/runs/304463075281818)
```bash
git init -q && git add -A && git commit -qm init
SHA=$(git rev-parse HEAD)
echo "print('uncommitted')" > src/uncommitted.py # created AFTER the commit
cat > git.yaml <<YAML
experiment_name: vchen_demo_git
command: cd \$CODE_SOURCE_PATH && python train.py
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
code_source: {type: snapshot, snapshot: {root_path: src, git: {commit: $SHA}}}
YAML
dbcli experimental air run -f git.yaml -p dbc-04ac0685-8857
# Uploading src.tar.gz... → Submitted run 304463075281818
dbcli workspace export /Workspace/Users/v.chen@databricks.com/.air/repo_snapshots/.internal/src.tar.gz --file /tmp/git.tar.gz -p dbc-04ac0685-8857
tar tzf /tmp/git.tar.gz
# → src/train.py, src/pkg/util.py, src/.gitignore
# (NO uncommitted.py, NO debug.log ✓ archived the commit, not the working tree)
```
**3. UC Volume destination** — `remote_volume` routes the upload to a UC
Volume via the
Files API (`/api/2.0/fs/files/...`), natively, with no special-casing in
the CLI.
Run:
[438683652713410](https://dbc-04ac0685-8857.staging.cloud.databricks.com/jobs/runs/438683652713410)
```bash
dbcli volumes create main default vchen_demo MANAGED -p dbc-04ac0685-8857
cat > vol.yaml <<'YAML'
experiment_name: vchen_demo_vol
command: cd $CODE_SOURCE_PATH && python train.py
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
code_source: {type: snapshot, snapshot: {root_path: src, remote_volume: /Volumes/main/default/vchen_demo}}
YAML
dbcli experimental air run -f vol.yaml -p dbc-04ac0685-8857 --debug 2>&1 | grep -iE "submitted|code_source_path|/api/2.0/fs/files"
# → "code_source_path": "/Volumes/main/default/vchen_demo/.internal/src.tar.gz"
# → Submitted run 438683652713410
dbcli fs ls dbfs:/Volumes/main/default/vchen_demo/.internal -p dbc-04ac0685-8857
# → src.tar.gz (uploaded to the Volume ✓)
```
**4. `include_paths` subset** — only the listed paths are packaged.
Run:
[261250771126835](https://dbc-04ac0685-8857.staging.cloud.databricks.com/jobs/runs/261250771126835)
```bash
cat > inc.yaml <<'YAML'
experiment_name: vchen_demo_inc
command: cd $CODE_SOURCE_PATH && python train.py
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
code_source: {type: snapshot, snapshot: {root_path: src, include_paths: [pkg]}}
YAML
dbcli experimental air run -f inc.yaml -p dbc-04ac0685-8857
# Uploading src.tar.gz... → Submitted run 261250771126835
dbcli workspace export /Workspace/Users/v.chen@databricks.com/.air/repo_snapshots/.internal/src.tar.gz --file /tmp/inc.tar.gz -p dbc-04ac0685-8857
tar tzf /tmp/inc.tar.gz
# → src/pkg/util.py ONLY (train.py / .gitignore excluded ✓)
```
**5. No `code_source`** — nothing is uploaded; `code_source_path` is
left empty (nil-guard).
Run:
[138286541552104](https://dbc-04ac0685-8857.staging.cloud.databricks.com/jobs/runs/138286541552104)
```bash
cat > none.yaml <<'YAML'
experiment_name: vchen_demo_none
command: echo hello
compute: {accelerator_type: GPU_1xA10, num_accelerators: 1}
environment: {version: "4", dependencies: []}
YAML
dbcli experimental air run -f none.yaml -p dbc-04ac0685-8857
# → Submitted run 138286541552104 (no "Uploading" line; code_source_path empty)
```
Approval status: pending
|
Integration test reportCommit: 05c7e4b
8 interesting tests: 4 SKIP, 3 KNOWN, 1 RECOVERED
Top 3 slowest tests (at least 2 minutes):
|
3db3098 to
6c2d198
Compare
a3e3771 to
bd6e3a1
Compare
The main->air-cli catch-up merge bumped the SDK to v0.165 and pulled in updated generated pydabs models, but the checked-in python/databricks/bundles/** was formatted by an older pinned ruff. The current ruff pin reformats them (e.g. "Self" -> 'Self', import wrapping), so `task generate-check` / the validate-generated CI job drifts on every PR into air-cli (#6102, #6090). Regenerated with `task pydabs-codegen` so the checked-in models match the generators + pinned ruff. Generated-only change (python/databricks/bundles/**). Co-authored-by: Isaac
bd6e3a1 to
6e7fa87
Compare
The main->air-cli catch-up merge bumped the SDK to v0.165 and pulled in updated generated pydabs models, but the checked-in python/databricks/bundles/** was formatted by an older pinned ruff. The current ruff pin reformats them (e.g. "Self" -> 'Self', import wrapping), so `task generate-check` / the validate-generated CI job drifts on every PR into air-cli (#6102, #6090). Regenerated with `task pydabs-codegen` so the checked-in models match the generators + pinned ruff. Generated-only change (python/databricks/bundles/**). Co-authored-by: Isaac
6e7fa87 to
2ae2e0d
Compare
This reverts commit b632473.
Adds `databricks experimental air convert-to-dabs <yaml_path>`: a purely local
transform that turns an AIR CLI run YAML into a Databricks Asset Bundle, so a
workload authored for `air run` can be deployed and managed with DABs.
Design:
- Round-trips: `bundle validate` accepts the output and `bundle deploy`
reproduces the same ai_runtime_task workload `air run` submits. Verified
end-to-end on staging (convert -> deploy -> run -> SUCCESS); the acceptance
test runs `bundle validate` on the output.
- The DABs ai_runtime_task is the SDK jobs.AiRuntimeTask (strict schema:
experiment + deployments[].{command_path,compute} + code_source_path).
Framework fields (retries, timeout) are emitted on the task wrapper, where the
schema puts them, not inside ai_runtime_task.
- code_source_path is a TARBALL, not a directory: the backend expects an archive
("unpacked on each node") and bundle deploy's uploader only handles single
files. convert tars the resolved root_path into code_source.tar.gz (reusing the
snapshot path's createPlainTarball, honoring include_paths). command.sh and the
env/secret/param sidecars are loose files at the bundle root. All are emitted
"./"-prefixed so IsLibraryLocal classifies them as local uploads (a bare,
extensionless path is otherwise read as a PyPI package name and skipped).
- No manual upload step: bundle deploy uploads the tarball + command.sh; the user
runs one deploy.
- env_variables / secrets / parameters have no inline ai_runtime_task fields (the
native jobs env_variables field is the eventual north star but isn't usable from
the Go SDK yet), so they ride as sidecar files, matching `air run`'s layout.
- docker_image and usage_policy_name are rejected with a clear message.
Output defaults to an ephemeral, per-experiment temp dir that is reset each run
(idempotent re-runs, no project clutter); --output-dir pins a permanent location.
Tests: unit tests over the pure conversion (field mapping, wrapper-vs-task
placement, omit-when-unset, sidecar staging, permissions, rejections) and an
acceptance test that converts the torchrun-a10 example and validates the bundle.
Co-authored-by: Isaac
…host) createPlainTarball shelled out to `tar -czf <outputTarball> ...`. On Windows an absolute output path like `C:\out\code_source.tar.gz` makes tar interpret the `C:` as a remote host (host:path syntax) and fail with "Cannot connect to C:". This surfaced on the Windows CI lane for the new convert-to-dabs acceptance test (the snapshot path never hit it — its acceptance test is cloud-gated and doesn't run on the Windows unit lane). Fix: pass only the archive basename to -f and run tar with its working directory set to the (absolute) output directory, so no drive-letter colon appears in the -f operand. The source -C parent is made absolute so it still resolves under the new working dir. Works for GNU tar and bsdtar alike; no per-flavor flag. Co-authored-by: Isaac
2ae2e0d to
503aa8d
Compare
Round out the printed deploy sequence with step 4 (databricks bundle run <job-key>) so the guidance is a complete deploy-and-run flow, using the actual job resource key. Destroy/cleanup guidance is retained. Co-authored-by: Isaac
503aa8d to
9b39e83
Compare
ben-hansen-db
left a comment
There was a problem hiding this comment.
A few behavioral things.
- new folder of databricks.yml should by default go in the same folder as the original yaml
- The user can specify the placement of the new folder using absolute paths
- code source shouldn't be pre-created imo; it should use the work you already for code source snapshot so snapshotting can happen during deployment.
- env vars and secrets are missing
- How does this handle parameters field?
- Can the generated databricks.yml say in a command what fields may need to be filled in by a user or which were ignored? (for instance if they have git snapshotting we just do regular snapshotting).
… notes Addresses review feedback: - Default the bundle to a <experiment>-bundle folder next to the input YAML instead of a /tmp scratch dir; --output-dir (absolute or relative) overrides. The bundle is the user's artifact to keep and manage, so it no longer lands in temp. Re-runs refuse to clobber an existing dir rather than silently wiping it. - Print a "Notes:" section listing what the conversion transformed or staged out-of-band, so a user migrating from `air run` can see what changed: the code_source snapshot, and env_variables/secrets/parameters staged as sidecar files (no native ai_runtime_task field for them yet). env vars / secrets keep riding as env_vars.json / secret_env_vars.json sidecars: SDK v0.165 has no native env-vars field on ai_runtime_task/Task/Environment, so a native mapping isn't possible yet (it's the documented BYOT north-star, not in the SDK). code_source still snapshots at convert time; a deploy-time snapshot mutator is tracked as a separate change. Co-authored-by: Isaac
Adds
databricks experimental air convert-to-dabs <yaml_path>: a purely local transform that turns an AIR CLI run YAML into a Databricks Asset Bundle, so a workload authored forair runcan be deployed and managed with DABs.Design:
bundle validateaccepts it andbundle deployreproduces the same ai_runtime_task workloadair runwould submit. The acceptance test runsbundle validateon the output to pin this.bundle deployuploads them via the existing collectLocalLibraries path (aiRuntimeCodeSourcePattern). The user runs one deploy.air runexactly.Tests: unit tests over the pure conversion (field mapping, wrapper-vs-task placement, omitted-when-unset, sidecar staging, permissions, rejections) and an acceptance test that converts the torchrun-a10 example and validates the bundle.
Co-authored-by: Isaac
Changes
Why
Tests
E2E test:
Install CLI
Create simple train yaml
run the
convert-to-dabscommandInvestigate what gets put in the tmp folder
Run databricks bundle validate
Run databricks bundle deploy + run
Run bundle delete
Deep scrub from Claude: