AIR CLI Integration: image wiring (registered image used with client side) - #6165
Open
riddhibhagwat-db wants to merge 6 commits into
Open
AIR CLI Integration: image wiring (registered image used with client side)#6165riddhibhagwat-db wants to merge 6 commits into
riddhibhagwat-db wants to merge 6 commits into
Conversation
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
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
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.