Run AIR CLI via DABs#5937
Open
vinchenzo-db wants to merge 6 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
Waiting for approvalBased on git history, these people are best suited to review:
Eligible reviewers: Suggestions based on git history. See OWNERS for ownership rules. |
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: e8780b8
10 interesting tests: 4 flaky, 4 SKIP, 2 RECOVERED
Top 9 slowest tests (at least 2 minutes):
|
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