AIR CLI Integration: Port Register Image Functionality - #6090
Open
riddhibhagwat-db wants to merge 10 commits into
Open
AIR CLI Integration: Port Register Image Functionality#6090riddhibhagwat-db wants to merge 10 commits into
riddhibhagwat-db wants to merge 10 commits into
Conversation
Port the AI Compute Manager image-registration client from the Python `air`
CLI (cli/image_registration/image_client.py) into Go, as the first phase of
the `air register-image` port.
This adds image_client.go, a raw client.Do client (the SDK does not model the
ai-compute-manager service) with:
- normalizeDockerImageURL: canonicalizes image URLs for consistent hashing,
matching the Python registry-detection, digest, and default-tag rules.
- imageRegistration/imageStatus: the response model, with an unknown "state"
degrading to PENDING.
- createImage / getImage / checkImageAccess / waitForImageReady, using the
literal ":get" / ":checkImageAccess" endpoint verbs the backend expects and
mapping apierr.ErrNotFound to a nil registration.
The command still returns notImplemented; wiring lands in phase 2.
Co-authored-by: Isaac
Implement `air register-image` end-to-end using the phase-1 image client,
porting the resolution logic from cli_entrypoint.py::handle_register and
image_policy.py.
- Flag validation: reject an empty IMAGE_URL, --scope/--key given apart, and
-i combined with --scope/--key, each with an actionable INVALID_ARGS error.
- resolveImage flattens the Python ImagePolicyHandler: reuse a cached
AVAILABLE image under `auto`, re-register under `latest` or when the stored
status isn't AVAILABLE, and report whether the digest changed.
- Text and JSON-envelope output match the Python CLI (docker_image_url,
manifest_sha256, status, image_updated, cached).
Interactive credential setup and local ~/.docker/config.json discovery are
deferred to later phases; -i is rejected as not-yet-supported rather than
silently ignored.
Drops register-image from the stub and unimplemented tests and adds an
acceptance test covering the cached/auto, latest-recheck, and flag-validation
paths.
Co-authored-by: Isaac
Deprecate `auto` image-resolution mode. Registration now always re-registers
and checks the source registry for the latest digest; the prior image's status
is no longer consulted (only the freshly-created registration's status gates
polling).
- --tag-policy is hidden and deprecated. "latest" and the empty default are
accepted no-ops for backward compatibility; "auto" is rejected with guidance
(it used to reuse a cached image, so silently remapping it would hide a
behavior change). No internal caller passes the flag.
- resolveImage drops the cached-reuse fast path and status branching; it
fetches the prior digest only to report image_updated.
- Text output collapses to "Image registered" / "Image already up to date".
Adds register_image_test.go covering validateTagPolicy and the fresh /
unchanged / changed-digest branches; updates the acceptance test to model a
fresh registration and the auto-rejection.
Co-authored-by: Isaac
Credentials for private images will come solely from the user's local Docker configuration (`docker login` + ~/.docker/config.json); there are no user-facing credential flags. Remove --scope, --key, and --interactive-authenticate along with their validation and threading. Registration currently proceeds with no credentials (public images only). Local Docker credential discovery and internal auto-storage of the secret the registration API requires land in the next phase, restoring private-image support with no credential flags. The image client's createImage still accepts a credentials scope/key (passed empty for now) so that phase can supply the auto-stored secret reference. Co-authored-by: Isaac
Restore private-image support with no user-facing credential flags: creds are
discovered from the local Docker config and stored in a per-user secret whose
reference is passed to registration.
- docker_config_creds.go: read ~/.docker/config.json (honoring DOCKER_CONFIG)
following Docker's own resolution order — per-registry credHelper, global
credsStore, then inline base64 auth. Helper subprocess runs via
exec.CommandContext with a timeout; any failure falls through to the next
source and is never fatal.
- image_credentials.go: store discovered creds in a per-user secret scope
(docker-credentials-<user>, creator-only ACL so a PAT can't leak to other
members) under a -local-suffixed key. Secret-scope quota errors surface;
other storage failures fall through to the public-image path.
- register-image wiring: probe checkImageAccess before storing (skip creds for
public images), then register; on an auth failure with stored creds, retry
anonymously in case the image is public. Auth failures surface actionable
`docker login` guidance. Auth classification uses the SDK's typed sentinels
rather than substring-matching 401/403.
Local creds are read once and threaded to storage so a credential helper is
never invoked twice. Adds unit tests (incl. a credential-helper stub, skipped
on Windows) and an acceptance case asserting the credential reference reaches
the registration POST while the secret value never appears in output.
Co-authored-by: Isaac
discoverCredentials passed the raw IMAGE_URL to readDockerCredentials, but that lookup keys off the registry host and requires the normalized URL. A bare Docker Hub image (e.g. `ubuntu`) resolved to the wrong host, so `docker login` credentials were silently missed. Normalize once before lookup, matching the Python handler. Also drop a stale doc-comment remnant on storeDockerCredentials, trim the images-API path comment to the verbs actually used, and give the secret-scope quota error an actionable message. Co-authored-by: Isaac
renderError for a failed registration hardcoded kind=TRANSIENT/retryable=true, so permanent failures (auth, not-found, cancellation, a terminal FAILED upload) were reported as retryable in JSON mode — a consumer keying on `retryable` would retry an unrecoverable auth failure forever, contradicting the `docker login` guidance in the same message. Add classifyRegistrationError: auth/not-found/canceled/upload-failed → PERMANENT/non-retryable; poll timeout and unclassified API errors → TRANSIENT/retryable. A terminal FAILED upload now wraps errImageUploadFailed so it classifies distinctly from a transient poll error. Found by review. Adds table-driven coverage for the classification. Co-authored-by: Isaac
…retry
Second review pass: classifyRegistrationError's catch-all marked every
unclassified failure (400 bad-request, 409 conflict, and any unmapped error) as
TRANSIENT/retryable, so a consumer would retry a request that can't succeed.
Invert it: positively identify transient errors (rate limit, server error,
deadline, and the poll timeout) and default everything else to
PERMANENT/non-retryable, matching json_output.ErrorKind.INTERNAL.
- The poll timeout is now the errImageWaitTimeout sentinel rather than a
formatted string, so classification keys off errors.Is, not prose.
- Extract registerWithCredentialFallback so the stale-credential anonymous
retry is a named, testable function; add tests for the retry (creds
rejected -> anonymous success) and the no-creds path.
- Extend TestClassifyRegistrationError with bad-request/conflict/rate-limit/
server-error/unknown rows.
Co-authored-by: Isaac
Final review nits: reject --timeout-minutes <= 0 as INVALID_ARGS (a 0/negative value otherwise makes waitForImageReady time out immediately and misreport the misconfiguration as a retryable transient error). Python has the same gap, so this is a small improvement, not a parity change. Also add the two uncovered transient classification arms (ErrTemporarilyUnavailable, ErrDeadlineExceeded) to TestClassifyRegistrationError and an acceptance case for the timeout rejection. Co-authored-by: Isaac
Contributor
Approval status: pending
|
Collaborator
Integration test reportCommit: 35d7b8e
|
| } | ||
|
|
||
| scope = fmt.Sprintf("%s-%s", dockerCredsScopePrefix, me.UserName) | ||
| key = fmt.Sprintf("dockerio-%s%s", username, localManagedKeySuffix) |
Contributor
There was a problem hiding this comment.
We can't assume the host is always dockerio. It could also be nvcr.io. You can get the actual host from dockerImageURL, something like:
host, _, _ := strings.Cut(normalizeDockerImageURL(dockerImageURL), "/")
Alternatively you can pass in the already normalized version
vinchenzo-db
added a commit
that referenced
this pull request
Jul 30, 2026
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
vinchenzo-db
added a commit
that referenced
this pull request
Jul 30, 2026
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 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
Ports the
dcs register-imagecapability (image registration) from the Pythonai-compute/cliinto the Go CLI asair register-image, underexperimental/air/cmd. Mirrors a Docker image into the workspace registry.air register-image IMAGE_URLregisters an image and waits for it to become AVAILABLE, reporting the manifest digest (text or-o jsonenvelope).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
airCLI 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 existingdocker loginand stored per-user (never workspace-readable), so a registry PAT is never passed on the command line or leaked to other workspace members.Tests
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.