Run AIR CLI via DABs#5937
Open
vinchenzo-db wants to merge 11 commits into
Open
Conversation
…-> run) with env vars
Wires the experimental `air` CLI to submit and list workloads via Databricks
Asset Bundles (persistent Jobs resources) alongside the ephemeral runs/submit
path. Gated behind --via-dabs / AIR_VIA_DABS=1; the ephemeral path stays the
default and is untouched (additive).
- run: `air run --via-dabs` converts train.yaml -> databricks.yml (reusing the
exportbundle.go converter), then deploys the bundle and runs it. env vars/secrets
ride the common Jobs env-var API (API-293) as a job-level environment_variables
profile the ai_runtime_task references by key; secrets emit as
{{secrets/scope/key}} refs. Verified against staging that this shape persists on
runs/get, so the converter's env-var/secret rejection is lifted.
- list: `air list --via-dabs` drops the SUBMIT_RUN server-side restriction so
DABs-submitted AIR runs (RunType JOB_RUN) are included; isAirRun still selects by
task shape. The AiTrainingService index path already returns the caller's runs
regardless of submit verb, so only the jobs-scan fallback needed widening.
- cancel / get: work unchanged for DABs runs (both key on a run_id that a
bundle run still produces); documented inline.
Prototype scope: deploy/run shell out to `databricks bundle` (production would call
cmd/bundle/utils.ProcessBundle in-process); run_id parsing and code_source staging
into the bundle folder are TODO(prototype). Persistent-job GC is accepted for now.
Co-authored-by: Isaac
Rewrites `air run` to submit exclusively via a Databricks Asset Bundle (convert train.yaml -> databricks.yml -> deploy -> run), removing the ephemeral runs/submit path and the experimental --via-dabs flag. This honors the AIR-CLI/DABs agreement: DECO requires AIR to integrate with DABs and submit through the public Jobs API, not an internal proxy. Deploy and run are driven IN-PROCESS via the bundle libraries (cmd/bundle/utils.ProcessBundle + Jobs.RunNow), the same entry points the `databricks bundle` commands use and the pattern the pipelines CLI established (cmd/pipelines). There is no child `databricks` process, so the deploy always uses this build's ai_runtime_task-aware bundle schema. - rundabs.go: submitWorkload converts the config to a bundle in a temp root, deploys it with ProcessBundle (rooted via DATABRICKS_BUNDLE_ROOT on the context, direct engine), then RunNow on the resolved job to return the run_id (fire-and-return, no wait). A synthetic carrier command supplies the flags ProcessBundle reads and forwards air's resolved auth profile. - run.go: DABs is the only path; --via-dabs removed. --dry-run always renders the generated bundle for transparency. - list.go: the jobs scan no longer restricts to SUBMIT_RUN, since AIR runs are now runs of a persistent DABs job (JOB_RUN); --via-dabs gate removed. isAirRun still selects by task shape. - runsubmit.go: ephemeral submit path retired, leaving only the runtime- image defaults the converter shares. Verified end-to-end on staging (e2-dogfood): run -> get -> list -> cancel all work against a real in-process deploy+run. go build + go test ./experimental/air/... pass. Co-authored-by: Isaac
Contributor
Approval status: pending
|
Wires code_source snapshots through the DABs content-addressed snapshot
path (experimental.immutable_folder) so `air run` uploads the user's code
tree the same zip-and-fingerprint way AIR does today — giving an
apples-to-apples benchmark against the mechanism the uploads-vs-DABs
scoping doc compares (DECO-27087), not the per-file bundle sync.
- exportbundle.go: emit experimental.immutable_folder=true; emit
command_path as a bundle-root-relative path so translate_paths rewrites
it to the deployed snapshot location (emitting ${workspace.file_path}
directly is validated as a local file and fails). The convertibility
gate now allows working-tree code_source snapshots but still rejects
git-pinned (working-tree upload can't pin a commit) and remote_volume
(immutable_folder is Workspace-Files-only) snapshots.
- rundabs.go: writeBundleProject copies the code_source working tree into
the bundle root (honoring include_paths) so deploy uploads it in the
immutable snapshot.
- Trim AI-narration comments across the air files (provenance dumps,
verification history, migration storytelling) down to what/why.
Verified on staging: `air run` with a code_source snapshot deploys via
"Uploading immutable bundle snapshot..." and submits a run. (The AI
Runtime execution then fails on a backend PrincipalContext/entity-owner
error, independent of code delivery and of the CLI path.)
Co-authored-by: Isaac
Covers the parts of rundabs.go that don't need a live workspace: - writeBundleProject: databricks.yml + command.sh land, immutable_folder is set, code_source tree is staged, a user command.sh doesn't shadow the generated one, and the convertibility gate fails before any temp dir is created. - newBundleCarrierCommand: forwards air's profile, declares the flags ProcessBundle reads, and seeds bundle-root + direct engine on the context. - bundleTargetsBlock and isRunnableJob. deployBundle / runDeployedJob still need a live deploy+RunNow and stay covered by staging validation. Co-authored-by: Isaac
Match the Python CLI (#2237609): env vars and secrets ride the job-level
environment_variables profile the converter emits, not the legacy
env_vars.json / secret_env_vars.json workspace side-channel. Drop that
sidecar staging from buildArtifacts and delete the now-dead entry
encoders and filename constants.
The other launch artifacts buildArtifacts stages (command.sh,
training_config.yaml, requirements.yaml, hyperparameters.yaml) are
unchanged — those fixed the "requirements.yaml not found" harness error
and are unrelated to env vars.
Known gap (verified on staging): the vendored databricks-sdk-go
jobs.JobSettings/Task has no environment_variables field, so the DABs
bundle schema strips it client-side ("unknown field") before the wire.
Env vars therefore aren't delivered on the DABs path until the SDK
carries the common Jobs env-var API. The Jobs server already accepts the
field (gated by allowEnvironmentVariables); this is purely a
DABs/SDK-schema lag.
Co-authored-by: Isaac
Mirror the Python CLI's rollout dual-write (#2237609): env vars / secrets ride BOTH the common Jobs env-var API (the converter's environment_variables profile) and the env_vars.json / secret_env_vars.json sidecars the scheduler falls back to. The prior commit removed the sidecars; restore them. The sidecar fallback matters on two fronts: (1) the allowEnvironmentVariables SAFE flag is off in prod, so Jobs drops the profile field, and (2) on the DABs path the bundle SDK schema strips the profile field client-side regardless of the flag (vendored jobs.JobSettings has no environment_variables). Until the flag ramps AND the SDK carries the field, the sidecar is the only working env-var delivery on the DABs path. Verified on staging: a code-free run with env_variables + secrets deploys env_vars.json / secret_env_vars.json into the immutable snapshot with the correct shapes, next to command.sh. Co-authored-by: Isaac
Collaborator
Integration test reportCommit: 2a7f1ab
8 interesting tests: 4 RECOVERED, 4 SKIP
Top 5 slowest tests (at least 2 minutes):
|
- exportbundle.go: replace the env-var copy loop with maps.Copy (modernize mapsloop). - exportbundle_test.go: assert.Contains instead of assert.True(strings.Contains(...)) (testifylint). Co-authored-by: Isaac
Root cause (found via driver logs on staging): a DABs-deployed ai_runtime_task failed at execution with `python train.py: No such file or directory`. The AI Runtime harness invokes command.sh from an unrelated runtime work dir, not the snapshot dir where bundle deploy syncs the code, so relative references can't find code_source files. Code-free commands (e.g. `python -c ...`) were unaffected, which is why earlier smoke tests passed. Fix: generated command.sh now prepends `cd "$(dirname "$0")"`, running the user's command from command.sh's own directory (the synced code_source location) — the DABs-native equivalent of the ephemeral path's `cd $CODE_SOURCE_PATH`. Also fixes forbidigo lint in rundabs_test.go (errors.Is/fs.ErrNotExist over os.IsNotExist). Verified end-to-end on staging (dbc-04ac0685-8857): a code_source run that failed 4x now SUCCEEDs — driver log shows the cwd contains train.py and "numpy 2.1.3" printed. Co-authored-by: Isaac
Retiring the ephemeral runs/submit path orphaned AIR's own tarball/git-archive code-upload subsystem — the DABs path delivers code via bundle deploy (immutable_folder), not a tarball. golangci `unused` flagged the entry points; removing them cascades to the whole subsystem. Delete snapshot_git.go, snapshot_package.go, snapshot_cachekey.go, snapshot_resolve.go and their tests; gut snapshot.go to just resolveRootPath (the one piece still used, by rundabs.go's stageCodeSource). Net -1585 lines. Verified: build, `go test ./experimental/air/...`, and the CI-exact lint (go tool -modfile=tools/go.mod golangci-lint run) all pass with 0 issues. Co-authored-by: Isaac
The run config already parsed a permissions block but the converter silently
dropped it, so it never took effect. Wire cfg.Permissions through to the emitted
bundle's job (level + one of user_name/group_name/service_principal_name), and
add a repeatable --permissions flag ("CAN_VIEW=group_name:users") that merges
into the config's block; flag and YAML grants are re-validated together.
This matters because a DABs-deployed job is stamped edit_mode=UI_LOCKED — its
permissions cannot be edited via the UI or permissions API after deploy ("Modify
bundle sources and redeploy to edit this job"), so declaring them in the bundle
is the only way to grant access on AIR-via-DABs jobs.
Verified on staging (dbc-04ac0685-8857): `air run --permissions
"CAN_VIEW=group_name:users"` produces a job whose ACL grants the users group
CAN_VIEW.
Co-authored-by: Isaac
air run no longer submits via runs/submit — it converts to a bundle and deploys/runs. Bring the acceptance suite in line: - Delete run-submit: it tested the retired ephemeral runs/submit path with a git code_source (POST /jobs/runs/submit + tarball upload). air run now deploys via DABs and rejects git code_source, so the scenario no longer exists. - Update run/output.txt: --dry-run now renders the generated databricks.yml (bundle preview) instead of "configuration is valid". - run/with-code-source.yaml: drop the git block so "code_source passes validation" stays a positive case — the DABs path supports working-tree snapshots but rejects git-pinned ones (covered by git-remote.yaml's rejection). Co-authored-by: Isaac
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Changes
Why
Tests