AIR CLI Integration: image wiring (registered image used with client side) - #6165
Open
riddhibhagwat-db wants to merge 22 commits into
Open
AIR CLI Integration: image wiring (registered image used with client side)#6165riddhibhagwat-db wants to merge 22 commits into
riddhibhagwat-db wants to merge 22 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)
```
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
This reverts commit b632473.
## 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? --> ## Tests <!-- How have you tested the changes? --> <!-- If your PR needs to be included in the release notes for next release, add a changelog fragment: create .nextchanges/<section>/<name>.md with a one-line description (e.g. .nextchanges/cli/quickstart.md). See .nextchanges/README.md. -->
Reverts #6121 ("Air cli drop requirements yaml"). #6121 moved `air run` dependencies onto `environments[].spec.dependencies` and dropped the co-located `requirements.yaml` upload. Reverting so the equivalent change can land via #6077, which additionally: - resolves the **version declared inside a file-form `requirements.yaml`**. In #6121 `requirementsDoc.Version` is decoded but never used, so `dependencies: ./reqs.yaml` with `version: 5` inside silently falls back to the default runtime image (`cfg.runtimeVersion()` returns `ok=false` for the file form). - rejects `-r`/`--requirement` includes in a requirements file, which reference a second file that is never uploaded with the run and so cannot resolve on the node. - adds acceptance coverage (`acceptance/experimental/air/run-submit-deps`) asserting the deps on the wire, the file-form version, and that no requirements file is uploaded. ## Tests Verified on this branch after the revert: `go build ./experimental/air/...`, the `experimental/air/cmd` unit suite, and the air acceptance suite (`TestAccept/experimental/air`) all pass. This pull request and its description were written by Isaac.
…6077) ## Changes & Why `air run` now carries the user's declared dependencies (which may be an inline list, or read from a requirements.yaml file) on the submission's environments[].spec.dependencies, and no longer uploads a requirements.yaml artifact at all. This is the follow up PR to https://github.com/databricks-eng/universe/pull/2178617?timeline_per_page=5 (implementing this method in the python CLI) and https://github.com/databricks-eng/universe/pull/2297011?timeline_per_page=5 (follow up backend changes to unblock the new path; installs the inline deps via --deps-config and treats a missing co-located requirements.yaml as "no requirements"). This removes the vestigial empty requirements.yaml that a no-dependency run used to upload just to satisfy the launcher's derived path. When no dependencies are declared, spec.dependencies is omitted and the payload is unchanged. A -r/--requirement include in a requirements file is rejected, since the referenced file is never uploaded with the run. ## Tests Unit tests: - Upload side: TestBuildArtifacts_CommandAndConfig, TestBuildArtifacts_ParametersButNoRequirements, TestBuildArtifacts_RequirementsFileNotUploaded, TestBuildArtifacts_EnvVarsAndSecrets, TestBuildArtifacts_OversizeConfigRejected - Submit side: TestBuildSubmitPayloadInlineDependencies, TestEnvironmentDependencies, TestReadRequirementsDependencies, TestEnvironmentDependencies_MissingRequirementsFile - End-to-end (unit): TestSubmitWorkload / TestSubmitWorkloadWithCodeSource �� Acceptance tests: `acceptance/experimental/air/run-submit-deps/` -> submits with inline deps, golden asserts: - `spec.dependencies: [numpy, torch==2.3.0]` on the runs/submit wire �� - Only `command.sh` + `training_config.yaml` uploaded no requirements.yaml � Can verify tests using: ``` go test ./experimental/air/cmd/ â�� pass go test ./acceptance -run TestAccept/experimental/air â�� pass gofmt clean ``` Manual verification: Instantiates a run succesfully with/without req.yaml dependencies declared: <img width="1673" height="1033" alt="Screenshot 2026-07-27 at 1 45 20 PM" src="https://github.com/user-attachments/assets/61ff9e24-225f-4cd2-87af-9b63030980b3" /> <img width="1166" height="354" alt="Screenshot 2026-07-27 at 1 46 08 PM" src="https://github.com/user-attachments/assets/97a0ffc9-a1e6-4e65-809f-a583bce6e083" />
## Changes Ports the `dcs register-image` capability (image registration) from the Python `ai-compute/cli` into the Go CLI as `air register-image`, under `experimental/air/cmd`. Mirrors a Docker image into the workspace registry. - `air register-image IMAGE_URL` registers an image and waits for it to become AVAILABLE, reporting the manifest digest (text or `-o json` envelope). - Registration always re-checks the source registry for the latest digest. - Credentials for private images are discovered from the local Docker config (`docker login` → `~/.docker/config.json`: credHelpers → credsStore → inline auth) and auto-stored in a per-user Databricks secret (creator-only ACL). If stored credentials are rejected, it retries once anonymously in case the image is public. ## Why Brings image registration to the Go `air` CLI so users on the Go binary can register private and public images. The credential-flag removal narrows the surface to a single, secure path so that creds are read from an existing `docker login` and stored per-user (never workspace-readable), so a registry PAT is never passed on the command line or leaked to other workspace members. ## Tests - Unit tests: URL normalization, status parsing, credential resolution order (incl. a credential-helper subprocess stub), secret scope/key storage + quota, error classification, and the anonymous-retry fallback. - Acceptance test (`acceptance/experimental/air/register-image/`) covers the registration flow, credential discovery (asserting the secret reference reaches the POST while the raw PAT never appears in output), and flag validation. - Manual Verification:
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. |
environment.docker_image parsed only `url`, dropping the tag_policy and credential fields the Python config supports. Add TagPolicy, CredentialsScope, and CredentialsKey with the same validation: tag_policy must be auto or latest, and the credential scope/key must be provided together. Also drop the stale TODO on dockerImageURL (image registration has landed) and add a dockerImage() accessor for the block. Nothing consumes these yet; the run preflight and submit plumbing follow. Co-authored-by: Isaac
Port the pre-submit image checks from the Python cli/docker_utils.py into
rundockerimage.go:
- waitForRegisteredImage requires an existing registration and blocks while it
is still PENDING/IMPORTING. A missing or FAILED registration is an error
pointing at `air register-image`, so a run fails here with a clear cause
rather than deep in the launch/pod stage.
- resolveLatestDockerImage re-registers when tag_policy is "latest" so the run
picks up the tag's newest digest, using the config's credentials when set and
otherwise the local Docker config, with the same stale-credential anonymous
retry as `air register-image`.
- prepareDockerImage sequences the two.
Nothing calls prepareDockerImage yet; the submit wiring follows.
Co-authored-by: Isaac
Complete the end-to-end path: `air run` now verifies the custom image before
doing any upload work, and passes it to the Jobs submit call.
- submitWorkload calls prepareDockerImage right after the idempotency token is
resolved, so an unregistered, failed, or still-importing image fails (or
blocks) before artifacts are uploaded.
- buildSubmitPayload sets ai_runtime_task.docker_image_url from
environment.docker_image.url, matching the Python jobs client.
NOTE: this does not compile against databricks-sdk-go v0.165.0 —
jobs.AiRuntimeTask does not model docker_image_url yet. The field is written as
DockerImageUrl in anticipation of the pending SDK PR; bump the SDK in go.mod once
it merges and this builds as-is. Everything else in the branch (config parsing
and the preflight helpers) was verified green with the field removed.
Co-authored-by: Isaac
Trim the docker-image comments to what the code does not already say. Co-authored-by: Isaac
Review fixes:
- The "registration in progress" and "re-resolving" messages used log.Infof,
which is silent at the CLI's default WARN level, so a run could block for up
to imageReadyTimeout with no output at all. Print them with cmdio.LogString
like the rest of `air run`.
- dockerImageConfig.validate checked a trimmed URL but stored the raw one, so
a padded `url:` passed validation and then rode the submitted task untrimmed.
Trim URL and the credential fields in place, matching the Python validator;
this also stops a blank credentials_scope from suppressing discovery.
- Move the preflight below ensureExperimentDirectory/userWorkspaceDir so a bad
experiment directory fails immediately instead of after a tag_policy=latest
refresh, while still preceding any upload.
Adds coverage for the two previously-unexercised branches: credential
auto-discovery on the latest path (asserting the reference rides the POST) and
the storage-denied path (asserting the error names that cause, not `docker
login`).
Co-authored-by: Isaac
- Only retry anonymously when the credentials were auto-discovered. The gate
was `scope != ""`, so credentials the user configured explicitly were also
retried away and then blamed on a missing `docker login`. Carry the
distinction in an imageCredentials struct; an explicitly-named secret that is
rejected now reports that secret instead.
- Reject docker_image.credentials_scope/credentials_key under the default tag
policy. They are only consulted when re-resolving the tag, so accepting them
otherwise silently ignored them.
- Note on waitForRegisteredImage that Python's :validateImageAccess preflight
is deliberately not ported and should be implemented in the backend, so every
client benefits and the CLI does not pay a round trip per submit.
Co-authored-by: Isaac
riddhibhagwat-db
force-pushed
the
air-integration-m6-2
branch
from
August 4, 2026 23:57
2ab6d00 to
1f550e7
Compare
Collaborator
Integration test reportCommit: 1f550e7
|
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
3 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
Wires a registered Docker image end-to-end through
air run, so an image registered withair register-imageis actually verified and used by a run.environment.docker_imagenow parsestag_policy,credentials_scope, andcredentials_key(previously onlyurl), with the same validation as the Python config:tag_policymust beautoorlatest, and the credential scope/key must be provided together.rundockerimage.goadds the pre-submit checks ported from the Pythoncli/docker_utils.py:PENDING/IMPORTING; a missing orFAILEDregistration errors withair register-imageguidance.tag_policy: latest, re-register so the run picks up the tag's newest digest, using the configured credentials or else the local Docker config.submitWorkloadruns those checks after the cheap workspace calls but before any artifact upload, andbuildSubmitPayloadsetsai_runtime_task.docker_image_url.Why
air runpreviously parsedenvironment.docker_image.urlandrunconfig_launch.gocarried a TODO saying full support needed image registration, which landed in the parent PR.Tests
FAILED/ waits-while-importing,autonot re-registering vslatestre-registering, credential auto-discovery (asserting the secret reference rides the POST), storage-denied reporting that cause rather thandocker login, and an explicitly-named secret being rejected without an anonymous retry.tag_policyvalues, credential pairing, credentials rejected under the default policy, and a blank scope not counting as set.buildSubmitPayloadtests (with and without a docker image) are included but can't compile until the SDK field lands. Everything else was verified green (go test,golangci-lint,gofmt) with that one field removed.