AIR CLI Integration: Adding support for air run configuration#5657
Merged
Conversation
riddhibhagwat-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
riddhibhagwat-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
Collaborator
Integration test reportCommit: f764fa9
21 interesting tests: 13 SKIP, 7 KNOWN, 1 RECOVERED
Top 28 slowest tests (at least 2 minutes):
|
riddhibhagwat-db
force-pushed
the
air-integration-m2-2
branch
from
73088e7 to
373988d
Compare
riddhibhagwat-db
force-pushed
the
air-integration-m2-2
branch
from
373988d to
226d41a
Compare
riddhibhagwat-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
riddhibhagwat-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
|
|
||
| // validate runs structural validation over the whole config, returning the first | ||
| // failure. Fields are checked in declaration order to keep error output stable. | ||
| func (c *runConfig) validate() error { |
Author
There was a problem hiding this comment.
| DockerImage *dockerImageConfig `yaml:"docker_image"` | ||
| } | ||
|
|
||
| func (e *environmentConfig) validate() error { |
Author
There was a problem hiding this comment.
There was a problem hiding this comment.
this will change slight as Maggie figures out name for image url but this is good for now
There was a problem hiding this comment.
Just FYI: In the future DockerImage will not be mutually exclusive with Dependencies. And the we might give it a new name.
| URL string `yaml:"url"` | ||
| } | ||
|
|
||
| func (d *dockerImageConfig) validate() error { |
Author
There was a problem hiding this comment.
| Snapshot *snapshotSourceConfig `yaml:"snapshot"` | ||
| } | ||
|
|
||
| func (c *codeSourceConfig) validate() error { |
Author
There was a problem hiding this comment.
| IncludePaths []string `yaml:"include_paths"` | ||
| } | ||
|
|
||
| func (s *snapshotSourceConfig) validate() error { |
Author
There was a problem hiding this comment.
| Remote gitRemote `yaml:"remote"` | ||
| } | ||
|
|
||
| func (g *gitRef) validate() error { |
Author
There was a problem hiding this comment.
| Level string `yaml:"level"` | ||
| } | ||
|
|
||
| func (p *permission) validate() error { |
Author
There was a problem hiding this comment.
There was a problem hiding this comment.
I wonder if there's a struct here for permissions already? This mirrors permissions field for dabs
| assert.Equal(t, 1, cfg.Compute.NumAccelerators) | ||
| } | ||
|
|
||
| func TestLoadRunConfig_FullFeatured(t *testing.T) { |
Author
There was a problem hiding this comment.
vinchenzo-db
left a comment
vinchenzo-db
left a comment
There was a problem hiding this comment.
maybe a question for @maggiewang-db or @ben-hansen-db, is there a way to make sure this run config doesn't diverge from the run config in the python cli?
| type gitRef struct { | ||
| Branch *string `yaml:"branch"` | ||
| Commit *string `yaml:"commit"` | ||
| Remote gitRemote `yaml:"remote"` |
There was a problem hiding this comment.
I'm not very good at Go but is the gitRemote flag used anywhere outside of tests?
Author
There was a problem hiding this comment.
gitRemote is used in production code, but only inside runconfig.go itself and no other file references it as of now. It is just that the feature it validates (code_source) is itself deferred, so gitRemote (and the rest of gitRef/snapshotSourceConfig) currently is there only as part of config-parsing/validation, not the run mechanics (yet)
Comment on lines
+40
to
+41
| EnvVariables map[string]string `yaml:"env_variables"` | ||
| Secrets map[string]string `yaml:"secrets"` |
There was a problem hiding this comment.
Make sure you check that these aren't pointing to the same env variable.... and if so I imagine EnvVariable should take precedence, but ask @maggiewang-db
Reasoning is that if user accidentally sets env var and secret, and doesn't know which one, if secret takes precedence and they try to read it they make accidentally leak their secret.
Author
There was a problem hiding this comment.
This makes sense, I can implement this change if @maggiewang-db can sign off on the precedence order
There was a problem hiding this comment.
If there's a collision we should just reject the request at validation with a clear error message.
| MaxRetries *int `yaml:"max_retries"` | ||
| TimeoutMinutes *int `yaml:"timeout_minutes"` | ||
| IdempotencyToken *string `yaml:"idempotency_token"` | ||
| Parameters map[string]any `yaml:"parameters"` |
There was a problem hiding this comment.
Are you sure map[string]any is best here? Maybe some validation to make sure people aren't injecting bad inputs.
Author
There was a problem hiding this comment.
this is a free form map without content validation right now (and mirrors python cli which says parameters: Optional[Dict[str, Any]] = Field(default=None) and does not have a validator (sdk/config.py:443) ).
I dont think we need to add an explicit validation here since there is no real privilege boundary (the parameters themselves are the user's own hyperparameters and they are not interpolated into another user's context or into a shell command. The yaml file that they land in is read by the users own training script. Also, the yaml.Marshal escapes/quotes keys and values, so a value like "; rm -rf /" would just become a quoted yaml string and nothing that is actualy executable.
Comment on lines
+153
to
+155
| // validateExperimentName enforces the Databricks Jobs API task_key constraints: | ||
| // the experiment_name becomes a task key, which caps at 100 characters and allows | ||
| // only alphanumerics, hyphens, and underscores. |
There was a problem hiding this comment.
do we need to do something similar with mlflow_run_name?
Author
There was a problem hiding this comment.
I have this right now for mlflow_run_name (see line 112)
if c.MLflowRunName != nil { v := strings.TrimSpace(*c.MLflowRunName) if v == "" { return errors.New("mlflow_run_name cannot be empty") } if !taskKeyRe.MatchString(v) { return fmt.Errorf("invalid mlflow_run_name %q: only alphanumeric characters, hyphens, and underscores are allowed", v) } }
If we want to also validate for 100 chars or less, I can add that in.
There was a problem hiding this comment.
That's a @ben-hansen-db question :P
riddhibhagwat-db
force-pushed
the
air-integration-m2-2
branch
from
226d41a to
c992624
Compare
riddhibhagwat-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
riddhibhagwat-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
ben-hansen-db
approved these changes
Jun 30, 2026
| MLflowRunName *string `yaml:"mlflow_run_name"` | ||
| MLflowExperimentDirectory *string `yaml:"mlflow_experiment_directory"` | ||
| Permissions []permission `yaml:"permissions"` | ||
| UsagePolicyName *string `yaml:"usage_policy_name"` |
There was a problem hiding this comment.
small update, we take usagepolicyname or usage policy id, they are mutually exclusive. underthehood we resolve to the id
| return err | ||
| } | ||
|
|
||
| if err := validateSecretRefs(c.Secrets); err != nil { |
There was a problem hiding this comment.
I see, you are validating everything here already
| DockerImage *dockerImageConfig `yaml:"docker_image"` | ||
| } | ||
|
|
||
| func (e *environmentConfig) validate() error { |
There was a problem hiding this comment.
this will change slight as Maggie figures out name for image url but this is good for now
| URL string `yaml:"url"` | ||
| } | ||
|
|
||
| func (d *dockerImageConfig) validate() error { |
| Level string `yaml:"level"` | ||
| } | ||
|
|
||
| func (p *permission) validate() error { |
There was a problem hiding this comment.
I wonder if there's a struct here for permissions already? This mirrors permissions field for dabs
| git: | ||
| branch: main | ||
| remote: origin | ||
| include_paths: |
maggiewang-db
approved these changes
Jun 30, 2026
Port the run YAML schema and its structural validation from the Python CLI's sdk/config.py: the top-level runConfig plus the environment, docker_image, code_source/snapshot/git, and permission blocks. loadRunConfig decodes a YAML file with KnownFields (mirroring pydantic extra="forbid") and runs the validation pass. "Structural" covers types, required fields, and format/cross-field rules that need no workspace access. Online checks (compute pool resolution, GPU availability), git/filesystem checks, _bases_ composition, and CLI --override handling are deferred to later milestones. Two deliberate divergences from the Python schema, both following from the training-service-only port: the compute pool fields were already dropped, and the top-level priority field is dropped here since it is a node-pool queue-ordering knob with no meaning for serverless workloads. Co-authored-by: Isaac
riddhibhagwat-db
force-pushed
the
air-integration-m2-2
branch
from
c992624 to
f764fa9
Compare
riddhibhagwat-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
riddhibhagwat-db
temporarily deployed
to
test-trigger-is
GitHub Actions
Inactive
riddhibhagwat-db
added a commit
that referenced
this pull request
Jun 30, 2026
## Changes Implements the `air run` happy path on top of the config schema (#5657), submitting a one-time training run through the Jobs API. Five commits, one per phase: 1. run config launch accessors: flatten the validated config into launch values (timeout seconds, retry default, requirements file-vs-inline, runtime version). 2. wire run command (load, validate, dry-run): air run -f <config> loads + structurally validates the YAML; `--dry-run` validates offline (no workspace/auth) and returns; `--override/--watch` are rejected for now with clear errors (ported in future PR). 3. pre-submit resolution: resolve current user / workspace home / a unique cli_launch dir, and ensure a custom `experiment_directory` exists. 4. upload launch artifacts: write training_config.yaml (1 MB cap), command.sh, requirements.yaml (file or synthesized from inline deps), `env_vars.json` / `secret_env_vars.json`, and hyperparameters.yaml into the launch dir via a workspace filer. 5. assemble + submit: build the native `ai_runtime_task` payload and `POST /api/2.2/jobs/runs/submit` directly, then print the run id + dashboard URL (or a JSON envelope). Submission uses the **native `ai_runtime_task`** task (BYOT task type) and it talks only to the Jobs API (which internally routes to training service endpoint) and has no genai-mapi forwarding (the MAPI path is deprecated). It isn't modeled by the typed SDK in go, so the payload is a custom struct posted to the raw endpoint. The proto is lean: env vars and secrets ship as co-located `env_vars.json` / `secret_env_vars.json` files rather than inline, and `requirements.yaml` / `hyperparameters.yaml` are derived server-side from the command directory. **Deferred, with explicit "not yet supported" errors (no silent drops):** `code_source` snapshot packaging, `--watch` log streaming, and `usage_policy_name`. `environment.docker_image` is accepted by the schema as scaffolding but not conveyed in the payload (the native path has no docker field). `node_pool_id` / `pool_name` / `priority` remain dropped (new AIR CLI does not support pool placement). ## Why `air run` is the core of the migration for AIR CLI. Splitting it into per-phase commits keeps each reviewable in isolation, and stacking on the schema PR keeps that PR focused. Regarding some specific decisions: - We maintain the native ai_runtime_task (and not the genai_compute_task interfacing with mapi) as a hand built struct posted to the raw endpoint. This is so that we can interface with jobs directly (and jobs.SubmitTask only knows gen_ai_compute_task and this typed struct also omits the env-vars/secrets/requirements fields that are needed for the run) and make sure we also stay off the deprecated genai-mapi forwarding path. - `--dry-run` is decoupled from auth. It validates the config locally and returns before any workspace call, so config validation works fully offline (matching the Python CLI). Only actual submission requires an authenticated workspace client. ## Tests - Unit tests for every phase: launch accessors, pre-submit resolution (incl. ensureExperimentDirectory create/exists/not-a-directory), artifact assembly + upload, payload assembly, and submitWorkload end-to-end against a fake workspace. - New acceptance/experimental/air/run test covering --dry-run (text + JSON), the --override/--watch guards, an invalid config, and missing --file. - Updated the unimplemented acceptance test (removed run, now implemented). `go test ./experimental/air/...`, `go test ./acceptance -run TestAccept/experimental/air`, and `./task lint-q` all pass. **Manual verification tests (all pass):** - Dry run (offline, no auth) > - command only > - full run config > - json output - actual run submission > - throws error when profile is not set > - submission loop: submitted, can see the run in `air list` and `air get` and mlflow environment was created > - same run id gets ouputted when run submitted with the SAME idempotency key > - new run gets created when run submitted with SAME config but DIFFERENT idempotency key - `--watch` and `--override` return an informative error message (since they are not supported yet, but are valid flags) - usage_policy_name set in config throws error: usage_policy_name is not yet supported - code_source set in config throws error: code_source is not yet supported - missing --file throws informative error: required flag(s) "file" not set - invalid config (e.g. experiment_name: bad.name, or num_accelerators not a multiple of the per-node count) throws field-specific validation error **How to test locally for manual verification:** Checkout & build: ```bash git fetch origin git checkout air-integration-m2-3 # this PR (stacked on air-integration-m2-2) ./task build ``` Sample configs: ```bash cat > /tmp/min.yaml <<'YAML' experiment_name: air-cuj command: python train.py compute: {accelerator_type: GPU_1xH100, num_accelerators: 1} YAML ``` ```bash cat > /tmp/full.yaml <<'YAML' experiment_name: full-run command: | pip install -r requirements.txt python train.py compute: {accelerator_type: GPU_8xH100, num_accelerators: 16} environment: {dependencies: [torch==2.3.0], version: 5} env_variables: {WANDB_PROJECT: demo} secrets: {HF_TOKEN: my_scope/hf_token} parameters: {lr: 0.001, epochs: 3} mlflow_run_name: full-run-v2 max_retries: 2 timeout_minutes: 120 YAML ``` Automated tests ```bash go test ./experimental/air/... # unit (incl. submitWorkload vs a fake workspace) go test ./acceptance -run TestAccept/experimental/air # acceptance (run + unimplemented) ./task lint-q # lint changed files ``` Dry run: ```bash ./cli experimental air run -f /tmp/min.yaml --dry-run # note that this command will, in the final version, be databricks experimental air run ./cli experimental air run -f /tmp/full.yaml --dry-run ./cli experimental air run -f /tmp/min.yaml --dry-run -o json ``` Actual run submission: ```bash PROFILE=<your-dev-profile> # no auth configured → fails fast (exit 1) env -u DATABRICKS_HOST -u DATABRICKS_TOKEN ./cli experimental air run -f /tmp/min.yaml #> Error: ... (cannot configure default credentials / auth) # submit → prints run_id + dashboard URL ./cli experimental air run -f /tmp/min.yaml -p $PROFILE -o json #> { "data": { "status":"SUBMITTED", "run_id":"<id>", "dashboard_url":"<host>/jobs/runs/<id>" } } # verify in the workspace: open dashboard_url (run exists), and the MLflow experiment was created. ./cli experimental air get <run_id> -p $PROFILE # run state ./cli experimental air list -p $PROFILE # run appears in the list # idempotency — SAME key returns the SAME run_id (no new run) ./cli experimental air run -f /tmp/min.yaml -p $PROFILE --idempotency-key demo-key-1 -o json # run_id = X ./cli experimental air run -f /tmp/min.yaml -p $PROFILE --idempotency-key demo-key-1 -o json # run_id = X (same) # idempotency — DIFFERENT key creates a NEW run ./cli experimental air run -f /tmp/min.yaml -p $PROFILE --idempotency-key demo-key-2 -o json # run_id = Y (new) ``` Unsupported flags (asserting that error is thrown): ```bash ./cli experimental air run -f /tmp/min.yaml --dry-run --watch #> Error: --watch is not yet supported ./cli experimental air run -f /tmp/min.yaml --dry-run --override compute.num_accelerators=8 #> Error: --override is not yet supported # usage_policy_name (needs a workspace to reach the submit guard) printf 'experiment_name: t\ncommand: x\ncompute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\nusage_policy_name: my-policy\n' > /tmp/policy.yaml ./cli experimental air run -f /tmp/policy.yaml -p $PROFILE #> Error: usage_policy_name is not yet supported # code_source printf 'experiment_name: t\ncommand: x\ncompute: {accelerator_type: GPU_1xH100, num_accelerators: 1}\ncode_source: {type: snapshot, snapshot: {root_path: .}}\n' > /tmp/code.yaml air run -f /tmp/code.yaml -p $PROFILE #> Error: code_source is not yet supported ``` Validation errors for field-specific message (exit 1, offline): ```bash # missing --file air run --dry-run #> Error: required flag(s) "file" not set # invalid experiment_name + num_accelerators not a multiple of the per-node count printf 'experiment_name: bad.name\ncommand: x\ncompute: {accelerator_type: GPU_8xH100, num_accelerators: 3}\n' > /tmp/bad.yaml air run -f /tmp/bad.yaml --dry-run #> Error: invalid experiment_name "bad.name": only alphanumeric characters, hyphens (-), and underscores (_) are allowed # (and, once the name is fixed: compute.num_accelerators for GPU_8xH100 must be a multiple of 8, got 3) ```
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
5 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
Ports the air run YAML config schema and its structural validation from the Python CLI (cli/sdk/config.py) to Go, under experimental/air/cmd/.
Two deliberate divergences from the Python schema, both following from the training-service-only port:
Why
"Structural" validation (types, required fields, format/cross-field rules) needs no workspace access, so it's a self-contained, fully unit-testable unit that's worth landing on its own ahead of the launch logic. Splitting it out keeps the upcoming handle_run PR focused on orchestration rather than mixing in ~900 lines of schema.
The extra="forbid" / KnownFields behavior is load-bearing: it's what turns a typo'd or stale config key into an actionable error instead of a silently-ignored field, so it's preserved faithfully. This is stacked on air-integration-m2-1 (the compute model).
Tests
New unit tests in runconfig_test.go (62 subtests, table-driven), covering:
go test ./experimental/air/... passes; ./task lint-q reports 0 issues.