From 0045b0bbd2fe7f2c14798ea5c7e21b37204e08d2 Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Mon, 3 Aug 2026 22:06:47 +0000
Subject: [PATCH 1/7] [air][m6-2] Parse docker_image tag_policy and credentials
 (phase 1)

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
---
 experimental/air/cmd/runconfig.go        | 24 +++++++++++++++++++
 experimental/air/cmd/runconfig_launch.go | 15 ++++++++----
 experimental/air/cmd/runconfig_test.go   | 30 ++++++++++++++++++++++++
 3 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go
index 4cfbf3736d1..19ed25fb2e9 100644
--- a/experimental/air/cmd/runconfig.go
+++ b/experimental/air/cmd/runconfig.go
@@ -305,12 +305,36 @@ func (s *stringOrInt) UnmarshalYAML(node *yaml.Node) error {
 // dockerImageConfig is environment.docker_image.
 type dockerImageConfig struct {
 	URL string `yaml:"url"`
+	// TagPolicy is "latest" to re-check the registry for the tag's newest digest
+	// on every run, or "auto" (the default) to reuse the existing registration.
+	// Registrations are per-user, so "latest" keeps a shared YAML current.
+	TagPolicy string `yaml:"tag_policy"`
+	// Registry credentials for a private image that "latest" re-resolves. When
+	// unset they are discovered from the local Docker config.
+	CredentialsScope string `yaml:"credentials_scope"`
+	CredentialsKey   string `yaml:"credentials_key"`
 }
 
+// dockerTagPolicy values for environment.docker_image.tag_policy.
+const (
+	dockerTagPolicyAuto   = "auto"
+	dockerTagPolicyLatest = "latest"
+)
+
 func (d *dockerImageConfig) validate() error {
 	if strings.TrimSpace(d.URL) == "" {
 		return errors.New("docker_image.url cannot be empty")
 	}
+
+	switch strings.ToLower(strings.TrimSpace(d.TagPolicy)) {
+	case "", dockerTagPolicyAuto, dockerTagPolicyLatest:
+	default:
+		return fmt.Errorf("invalid docker_image.tag_policy %q: must be %q or %q", d.TagPolicy, dockerTagPolicyAuto, dockerTagPolicyLatest)
+	}
+
+	if (d.CredentialsScope != "") != (d.CredentialsKey != "") {
+		return errors.New("docker_image.credentials_scope and docker_image.credentials_key must be provided together")
+	}
 	return nil
 }
 
diff --git a/experimental/air/cmd/runconfig_launch.go b/experimental/air/cmd/runconfig_launch.go
index 1408b600736..60b9352bbba 100644
--- a/experimental/air/cmd/runconfig_launch.go
+++ b/experimental/air/cmd/runconfig_launch.go
@@ -26,16 +26,21 @@ func (c *runConfig) maxRetries() int {
 }
 
 // dockerImageURL returns the custom docker image URL, or "" when none is set.
-//
-// TODO: not wired into submission yet — the native ai_runtime_task carries no
-// docker field, and full support needs image registration (pending the DCS work).
 func (c *runConfig) dockerImageURL() string {
-	if c.Environment != nil && c.Environment.DockerImage != nil {
-		return c.Environment.DockerImage.URL
+	if img := c.dockerImage(); img != nil {
+		return img.URL
 	}
 	return ""
 }
 
+// dockerImage returns the environment.docker_image block, or nil when none is set.
+func (c *runConfig) dockerImage() *dockerImageConfig {
+	if c.Environment == nil {
+		return nil
+	}
+	return c.Environment.DockerImage
+}
+
 // requirementsFile returns the path to a requirements file when
 // environment.dependencies is a string, and whether it was set.
 func (c *runConfig) requirementsFile() (string, bool) {
diff --git a/experimental/air/cmd/runconfig_test.go b/experimental/air/cmd/runconfig_test.go
index 45cae7cd0ae..af263d9581a 100644
--- a/experimental/air/cmd/runconfig_test.go
+++ b/experimental/air/cmd/runconfig_test.go
@@ -287,6 +287,36 @@ func TestEnvironmentConfigValidate(t *testing.T) {
 			environmentConfig{DockerImage: &dockerImageConfig{URL: "  "}},
 			"docker_image.url cannot be empty",
 		},
+		{
+			"tag policy latest ok",
+			environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", TagPolicy: "latest"}},
+			"",
+		},
+		{
+			"tag policy auto ok",
+			environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", TagPolicy: "AUTO"}},
+			"",
+		},
+		{
+			"invalid tag policy",
+			environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", TagPolicy: "newest"}},
+			`invalid docker_image.tag_policy "newest"`,
+		},
+		{
+			"credentials scope without key",
+			environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", CredentialsScope: "s"}},
+			"must be provided together",
+		},
+		{
+			"credentials key without scope",
+			environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", CredentialsKey: "k"}},
+			"must be provided together",
+		},
+		{
+			"credentials pair ok",
+			environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", CredentialsScope: "s", CredentialsKey: "k"}},
+			"",
+		},
 		{
 			"version with file deps",
 			environmentConfig{

From 92e93284e0677d720303eb86be24f7117634be36 Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Mon, 3 Aug 2026 22:08:25 +0000
Subject: [PATCH 2/7] [air][m6-2] Add docker image preflight helpers (phase 2)

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
---
 experimental/air/cmd/runconfig.go           |   6 ++
 experimental/air/cmd/rundockerimage.go      | 107 ++++++++++++++++++++
 experimental/air/cmd/rundockerimage_test.go |  87 ++++++++++++++++
 3 files changed, 200 insertions(+)
 create mode 100644 experimental/air/cmd/rundockerimage.go
 create mode 100644 experimental/air/cmd/rundockerimage_test.go

diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go
index 19ed25fb2e9..d2c5e614e14 100644
--- a/experimental/air/cmd/runconfig.go
+++ b/experimental/air/cmd/runconfig.go
@@ -338,6 +338,12 @@ func (d *dockerImageConfig) validate() error {
 	return nil
 }
 
+// wantsLatest reports whether the image should be re-resolved against the source
+// registry before the run.
+func (d *dockerImageConfig) wantsLatest() bool {
+	return strings.EqualFold(strings.TrimSpace(d.TagPolicy), dockerTagPolicyLatest)
+}
+
 // codeSourceConfig is the `code_source` block. Only the "snapshot" type exists.
 type codeSourceConfig struct {
 	Type     string                `yaml:"type"`
diff --git a/experimental/air/cmd/rundockerimage.go b/experimental/air/cmd/rundockerimage.go
new file mode 100644
index 00000000000..f5c1b22657c
--- /dev/null
+++ b/experimental/air/cmd/rundockerimage.go
@@ -0,0 +1,107 @@
+package aircmd
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"time"
+
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/databricks-sdk-go"
+)
+
+// imageReadyTimeout bounds how long a submit waits for a registration that is
+// still importing. Registration itself can take several minutes for a large
+// image, so this is generous.
+const imageReadyTimeout = time.Hour
+
+// digestDisplay abbreviates a manifest digest for logs, or reports it as unknown
+// when the registration carries none.
+func digestDisplay(sha string) string {
+	if sha == "" {
+		return "digest unknown"
+	}
+	return shortManifestSHA(sha)
+}
+
+// notRegisteredError explains that the image must be registered before a run can
+// use it, pointing at the command that does it.
+func notRegisteredError(dockerImageURL string) error {
+	return fmt.Errorf("docker image not registered: %s\nregister it first: databricks experimental air register-image %s", dockerImageURL, dockerImageURL)
+}
+
+// prepareDockerImage runs the pre-submit checks for a run's custom image: it
+// re-resolves the tag when tag_policy is "latest", then requires a registration
+// that is (or becomes) AVAILABLE. It returns an error rather than letting the run
+// fail later in the launch/pod stage, where the cause is much harder to see.
+func prepareDockerImage(ctx context.Context, w *databricks.WorkspaceClient, img *dockerImageConfig) error {
+	c, err := newImageClient(w)
+	if err != nil {
+		return err
+	}
+
+	if img.wantsLatest() {
+		if err := resolveLatestDockerImage(ctx, w, c, img); err != nil {
+			return err
+		}
+	}
+
+	return waitForRegisteredImage(ctx, c, img.URL)
+}
+
+// resolveLatestDockerImage re-registers the image so the run picks up the tag's
+// newest digest. Credentials come from the config when set, else from the local
+// Docker config; stale discovered credentials fall back to an anonymous retry the
+// same way `air register-image` does.
+func resolveLatestDockerImage(ctx context.Context, w *databricks.WorkspaceClient, c *imageClient, img *dockerImageConfig) error {
+	scope, key := img.CredentialsScope, img.CredentialsKey
+	var credErr error
+	if scope == "" {
+		scope, key, credErr = discoverCredentials(ctx, w, c, img.URL)
+		if credErr != nil {
+			log.Debugf(ctx, "could not store local Docker credentials: %v", credErr)
+		}
+	}
+
+	log.Infof(ctx, "re-resolving %s against the source registry (tag_policy=latest)", img.URL)
+	if _, _, err := registerWithCredentialFallback(ctx, c, img.URL, scope, key, imageReadyTimeout); err != nil {
+		return registrationError(img.URL, err, credErr)
+	}
+	return nil
+}
+
+// waitForRegisteredImage requires an existing registration and blocks while it is
+// still importing. A missing or failed registration is an error the user must fix
+// by (re-)registering.
+func waitForRegisteredImage(ctx context.Context, c *imageClient, dockerImageURL string) error {
+	reg, err := c.getImage(ctx, dockerImageURL)
+	if err != nil {
+		return err
+	}
+	if reg == nil {
+		return notRegisteredError(dockerImageURL)
+	}
+
+	switch reg.Status {
+	case imageStatusAvailable:
+		log.Infof(ctx, "using image %s (%s)", dockerImageURL, digestDisplay(reg.ManifestSHA256))
+		return nil
+	case imageStatusFailed:
+		msg := reg.StatusMessage
+		if msg == "" {
+			msg = "unknown error"
+		}
+		return fmt.Errorf("docker image registration failed: %s\nfix the issue and re-register: databricks experimental air register-image %s", msg, dockerImageURL)
+	case imageStatusPending, imageStatusImporting:
+		log.Infof(ctx, "docker image registration in progress (%s); waiting for it to become available", reg.Status)
+		final, err := c.waitForImageReady(ctx, dockerImageURL, imageReadyTimeout, imagePollInterval)
+		if err != nil {
+			return err
+		}
+		log.Infof(ctx, "image ready (%s)", digestDisplay(final.ManifestSHA256))
+		return nil
+	}
+
+	// Unreachable: normalizeStatus maps every unknown state to PENDING.
+	return errors.New("unexpected image status " + string(reg.Status))
+}
diff --git a/experimental/air/cmd/rundockerimage_test.go b/experimental/air/cmd/rundockerimage_test.go
new file mode 100644
index 00000000000..ebd3d5361ba
--- /dev/null
+++ b/experimental/air/cmd/rundockerimage_test.go
@@ -0,0 +1,87 @@
+package aircmd
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDigestDisplay(t *testing.T) {
+	assert.Equal(t, "digest unknown", digestDisplay(""))
+	assert.Equal(t, "abc", digestDisplay("abc"))
+	assert.Equal(t, "0123456789abcdef...", digestDisplay("0123456789abcdefghij"))
+}
+
+func TestWaitForRegisteredImageAvailable(t *testing.T) {
+	url := imageServer(t, `{}`, `{"state":"AVAILABLE","manifest_sha256":"abc"}`)
+	require.NoError(t, waitForRegisteredImage(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0"))
+}
+
+func TestWaitForRegisteredImageNotRegistered(t *testing.T) {
+	// :get 404s, so the run must stop with registration guidance.
+	url := imageServer(t, `{}`, "")
+	err := waitForRegisteredImage(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "docker image not registered")
+	assert.Contains(t, err.Error(), "air register-image nvcr.io/org/img:1.0")
+}
+
+func TestWaitForRegisteredImageFailed(t *testing.T) {
+	url := imageServer(t, `{}`, `{"state":"FAILED","status_message":"manifest not found"}`)
+	err := waitForRegisteredImage(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0")
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "registration failed: manifest not found")
+	assert.Contains(t, err.Error(), "re-register")
+}
+
+func TestWaitForRegisteredImageWaitsWhileImporting(t *testing.T) {
+	// First poll is still importing; the next reports AVAILABLE.
+	url := imageServer(t, `{}`,
+		`{"state":"IMPORTING"}`,
+		`{"state":"AVAILABLE","manifest_sha256":"abc"}`)
+	require.NoError(t, waitForRegisteredImage(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0"))
+}
+
+// latestImageServer serves a registry where the image is registered and
+// AVAILABLE, counting POSTs so a test can assert whether a re-registration
+// happened.
+func latestImageServer(t *testing.T, posts *int) string {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case imagesAPIPath + ":get":
+			_, _ = w.Write([]byte(`{"state":"AVAILABLE","manifest_sha256":"abc"}`))
+		case imagesAPIPath:
+			*posts++
+			_, _ = w.Write([]byte(`{"image":{"state":"AVAILABLE","manifest_sha256":"abc"}}`))
+		default:
+			_, _ = w.Write([]byte(`{}`))
+		}
+	}))
+	t.Cleanup(srv.Close)
+	return srv.URL
+}
+
+func TestPrepareDockerImageAutoDoesNotReregister(t *testing.T) {
+	var posts int
+	w := newTestWorkspaceClient(t, latestImageServer(t, &posts))
+	img := &dockerImageConfig{URL: "nvcr.io/org/img:1.0"}
+	require.NoError(t, prepareDockerImage(t.Context(), w, img))
+	assert.Zero(t, posts, "default tag policy must not re-register")
+}
+
+func TestPrepareDockerImageLatestReregisters(t *testing.T) {
+	var posts int
+	w := newTestWorkspaceClient(t, latestImageServer(t, &posts))
+	img := &dockerImageConfig{
+		URL:              "nvcr.io/org/img:1.0",
+		TagPolicy:        "latest",
+		CredentialsScope: "scope",
+		CredentialsKey:   "key",
+	}
+	require.NoError(t, prepareDockerImage(t.Context(), w, img))
+	assert.Equal(t, 1, posts, "tag_policy=latest must re-register against the source registry")
+}

From d5b1064fc7422c9c269f66e8e48e9ec470462069 Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Mon, 3 Aug 2026 22:09:59 +0000
Subject: [PATCH 3/7] [air][m6-2] Wire the docker image into run submit (phase
 3)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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
---
 experimental/air/cmd/runsubmit.go      | 11 +++++++++
 experimental/air/cmd/runsubmit_test.go | 31 ++++++++++++++++++++++++++
 2 files changed, 42 insertions(+)

diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go
index 3aa2436b827..4e5e5ba2455 100644
--- a/experimental/air/cmd/runsubmit.go
+++ b/experimental/air/cmd/runsubmit.go
@@ -80,6 +80,8 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri
 			},
 		}},
 		CodeSourcePath: snap.CodeSourcePath,
+		// The image must already be registered; prepareDockerImage verified that.
+		DockerImageUrl: cfg.dockerImageURL(),
 	}
 	if cfg.MLflowRunName != nil {
 		task.MlflowRun = *cfg.MLflowRunName
@@ -186,6 +188,15 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
 	if err != nil {
 		return 0, "", err
 	}
+
+	// After the cheap workspace checks (a tag_policy=latest refresh can block for
+	// minutes) but before any upload, so a bad image wastes no artifact work.
+	if img := cfg.dockerImage(); img != nil {
+		if err := prepareDockerImage(ctx, w, img); err != nil {
+			return 0, "", err
+		}
+	}
+
 	runName := ""
 	if cfg.MLflowRunName != nil {
 		runName = *cfg.MLflowRunName
diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go
index 76ab3aeb3ac..3c4fd6f2798 100644
--- a/experimental/air/cmd/runsubmit_test.go
+++ b/experimental/air/cmd/runsubmit_test.go
@@ -71,6 +71,37 @@ func TestBuildSubmitPayload(t *testing.T) {
 	assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu8xH100, AcceleratorCount: 16}, at.Deployments[0].Compute)
 }
 
+func TestBuildSubmitPayloadDockerImage(t *testing.T) {
+	cfg := &runConfig{
+		ExperimentName: "exp",
+		Command:        new("x"),
+		Compute:        &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1},
+		Environment: &environmentConfig{
+			DockerImage: &dockerImageConfig{URL: "nvcr.io/org/img:1.0"},
+		},
+	}
+
+	p := buildSubmitPayload(cfg, "/d/command.sh", "5", snapshotResult{}, nil)
+	require.Len(t, p.Tasks, 1)
+	require.NotNil(t, p.Tasks[0].AiRuntimeTask)
+	assert.Equal(t, "nvcr.io/org/img:1.0", p.Tasks[0].AiRuntimeTask.DockerImageUrl)
+}
+
+func TestBuildSubmitPayloadNoDockerImage(t *testing.T) {
+	// Without an environment.docker_image block the field stays empty (omitempty),
+	// so the runtime-managed environment is used.
+	cfg := &runConfig{
+		ExperimentName: "exp",
+		Command:        new("x"),
+		Compute:        &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1},
+	}
+
+	p := buildSubmitPayload(cfg, "/d/command.sh", "5", snapshotResult{}, nil)
+	require.Len(t, p.Tasks, 1)
+	require.NotNil(t, p.Tasks[0].AiRuntimeTask)
+	assert.Empty(t, p.Tasks[0].AiRuntimeTask.DockerImageUrl)
+}
+
 func TestBuildSubmitPayloadDefaultRetries(t *testing.T) {
 	// max_retries unset defaults to 3 (matching the Python native path), so both
 	// retry fields are sent.

From 2f4f599d97e520b134a7f9a36dd5f8383d416ee4 Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Mon, 3 Aug 2026 22:59:35 +0000
Subject: [PATCH 4/7] [air][m6-2] Tighten comments

Trim the docker-image comments to what the code does not already say.

Co-authored-by: Isaac
---
 experimental/air/cmd/runconfig.go      | 13 +++++--------
 experimental/air/cmd/rundockerimage.go | 23 +++++++----------------
 experimental/air/cmd/runsubmit.go      |  2 +-
 3 files changed, 13 insertions(+), 25 deletions(-)

diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go
index d2c5e614e14..f743407ff1c 100644
--- a/experimental/air/cmd/runconfig.go
+++ b/experimental/air/cmd/runconfig.go
@@ -305,17 +305,15 @@ func (s *stringOrInt) UnmarshalYAML(node *yaml.Node) error {
 // dockerImageConfig is environment.docker_image.
 type dockerImageConfig struct {
 	URL string `yaml:"url"`
-	// TagPolicy is "latest" to re-check the registry for the tag's newest digest
-	// on every run, or "auto" (the default) to reuse the existing registration.
-	// Registrations are per-user, so "latest" keeps a shared YAML current.
+	// "latest" re-checks the registry for the tag's newest digest each run;
+	// "auto" (default) reuses the existing registration.
 	TagPolicy string `yaml:"tag_policy"`
-	// Registry credentials for a private image that "latest" re-resolves. When
-	// unset they are discovered from the local Docker config.
+	// Credentials for a private image that "latest" re-resolves; discovered from
+	// the local Docker config when unset.
 	CredentialsScope string `yaml:"credentials_scope"`
 	CredentialsKey   string `yaml:"credentials_key"`
 }
 
-// dockerTagPolicy values for environment.docker_image.tag_policy.
 const (
 	dockerTagPolicyAuto   = "auto"
 	dockerTagPolicyLatest = "latest"
@@ -338,8 +336,7 @@ func (d *dockerImageConfig) validate() error {
 	return nil
 }
 
-// wantsLatest reports whether the image should be re-resolved against the source
-// registry before the run.
+// wantsLatest reports whether the image should be re-resolved before the run.
 func (d *dockerImageConfig) wantsLatest() bool {
 	return strings.EqualFold(strings.TrimSpace(d.TagPolicy), dockerTagPolicyLatest)
 }
diff --git a/experimental/air/cmd/rundockerimage.go b/experimental/air/cmd/rundockerimage.go
index f5c1b22657c..896ab7d2272 100644
--- a/experimental/air/cmd/rundockerimage.go
+++ b/experimental/air/cmd/rundockerimage.go
@@ -10,13 +10,10 @@ import (
 	"github.com/databricks/databricks-sdk-go"
 )
 
-// imageReadyTimeout bounds how long a submit waits for a registration that is
-// still importing. Registration itself can take several minutes for a large
-// image, so this is generous.
+// imageReadyTimeout is generous: replicating a large image takes minutes.
 const imageReadyTimeout = time.Hour
 
-// digestDisplay abbreviates a manifest digest for logs, or reports it as unknown
-// when the registration carries none.
+// digestDisplay abbreviates a manifest digest for logs.
 func digestDisplay(sha string) string {
 	if sha == "" {
 		return "digest unknown"
@@ -24,16 +21,12 @@ func digestDisplay(sha string) string {
 	return shortManifestSHA(sha)
 }
 
-// notRegisteredError explains that the image must be registered before a run can
-// use it, pointing at the command that does it.
 func notRegisteredError(dockerImageURL string) error {
 	return fmt.Errorf("docker image not registered: %s\nregister it first: databricks experimental air register-image %s", dockerImageURL, dockerImageURL)
 }
 
-// prepareDockerImage runs the pre-submit checks for a run's custom image: it
-// re-resolves the tag when tag_policy is "latest", then requires a registration
-// that is (or becomes) AVAILABLE. It returns an error rather than letting the run
-// fail later in the launch/pod stage, where the cause is much harder to see.
+// prepareDockerImage verifies a run's custom image before submit, so a bad image
+// fails here instead of deep in the launch stage where the cause is obscured.
 func prepareDockerImage(ctx context.Context, w *databricks.WorkspaceClient, img *dockerImageConfig) error {
 	c, err := newImageClient(w)
 	if err != nil {
@@ -50,9 +43,8 @@ func prepareDockerImage(ctx context.Context, w *databricks.WorkspaceClient, img
 }
 
 // resolveLatestDockerImage re-registers the image so the run picks up the tag's
-// newest digest. Credentials come from the config when set, else from the local
-// Docker config; stale discovered credentials fall back to an anonymous retry the
-// same way `air register-image` does.
+// newest digest, using the config's credentials when set and otherwise the local
+// Docker config.
 func resolveLatestDockerImage(ctx context.Context, w *databricks.WorkspaceClient, c *imageClient, img *dockerImageConfig) error {
 	scope, key := img.CredentialsScope, img.CredentialsKey
 	var credErr error
@@ -71,8 +63,7 @@ func resolveLatestDockerImage(ctx context.Context, w *databricks.WorkspaceClient
 }
 
 // waitForRegisteredImage requires an existing registration and blocks while it is
-// still importing. A missing or failed registration is an error the user must fix
-// by (re-)registering.
+// still importing. Missing or FAILED is an error the user fixes by re-registering.
 func waitForRegisteredImage(ctx context.Context, c *imageClient, dockerImageURL string) error {
 	reg, err := c.getImage(ctx, dockerImageURL)
 	if err != nil {
diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go
index 4e5e5ba2455..4310b117749 100644
--- a/experimental/air/cmd/runsubmit.go
+++ b/experimental/air/cmd/runsubmit.go
@@ -80,7 +80,7 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri
 			},
 		}},
 		CodeSourcePath: snap.CodeSourcePath,
-		// The image must already be registered; prepareDockerImage verified that.
+		// Verified as registered by prepareDockerImage.
 		DockerImageUrl: cfg.dockerImageURL(),
 	}
 	if cfg.MLflowRunName != nil {

From 3057cfbd6461c915f9ac285a8782e329b6f9771f Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Mon, 3 Aug 2026 23:22:24 +0000
Subject: [PATCH 5/7] [air][m6-2] Make the image wait visible; trim the config
 URL

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
---
 experimental/air/cmd/runconfig.go           |  8 +-
 experimental/air/cmd/rundockerimage.go      | 10 ++-
 experimental/air/cmd/rundockerimage_test.go | 81 +++++++++++++++++++--
 3 files changed, 89 insertions(+), 10 deletions(-)

diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go
index f743407ff1c..9a6a5fb9c80 100644
--- a/experimental/air/cmd/runconfig.go
+++ b/experimental/air/cmd/runconfig.go
@@ -320,7 +320,13 @@ const (
 )
 
 func (d *dockerImageConfig) validate() error {
-	if strings.TrimSpace(d.URL) == "" {
+	// Store the trimmed values: URL rides the submitted task, and the credential
+	// pairing check below must not treat blank-but-present as set.
+	d.URL = strings.TrimSpace(d.URL)
+	d.CredentialsScope = strings.TrimSpace(d.CredentialsScope)
+	d.CredentialsKey = strings.TrimSpace(d.CredentialsKey)
+
+	if d.URL == "" {
 		return errors.New("docker_image.url cannot be empty")
 	}
 
diff --git a/experimental/air/cmd/rundockerimage.go b/experimental/air/cmd/rundockerimage.go
index 896ab7d2272..4f624776f50 100644
--- a/experimental/air/cmd/rundockerimage.go
+++ b/experimental/air/cmd/rundockerimage.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"time"
 
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/databricks/cli/libs/log"
 	"github.com/databricks/databricks-sdk-go"
 )
@@ -55,7 +56,8 @@ func resolveLatestDockerImage(ctx context.Context, w *databricks.WorkspaceClient
 		}
 	}
 
-	log.Infof(ctx, "re-resolving %s against the source registry (tag_policy=latest)", img.URL)
+	// Visible at the default log level (WARN): this can block for minutes.
+	cmdio.LogString(ctx, fmt.Sprintf("Re-resolving %s against the source registry (tag_policy=latest)...", img.URL))
 	if _, _, err := registerWithCredentialFallback(ctx, c, img.URL, scope, key, imageReadyTimeout); err != nil {
 		return registrationError(img.URL, err, credErr)
 	}
@@ -84,12 +86,14 @@ func waitForRegisteredImage(ctx context.Context, c *imageClient, dockerImageURL
 		}
 		return fmt.Errorf("docker image registration failed: %s\nfix the issue and re-register: databricks experimental air register-image %s", msg, dockerImageURL)
 	case imageStatusPending, imageStatusImporting:
-		log.Infof(ctx, "docker image registration in progress (%s); waiting for it to become available", reg.Status)
+		// This wait runs up to imageReadyTimeout, so it must be visible at the
+		// default log level (WARN); log.Infof would be silent.
+		cmdio.LogString(ctx, fmt.Sprintf("Docker image registration in progress (%s); waiting for it to become available...", reg.Status))
 		final, err := c.waitForImageReady(ctx, dockerImageURL, imageReadyTimeout, imagePollInterval)
 		if err != nil {
 			return err
 		}
-		log.Infof(ctx, "image ready (%s)", digestDisplay(final.ManifestSHA256))
+		cmdio.LogString(ctx, "Image ready ("+digestDisplay(final.ManifestSHA256)+")")
 		return nil
 	}
 
diff --git a/experimental/air/cmd/rundockerimage_test.go b/experimental/air/cmd/rundockerimage_test.go
index ebd3d5361ba..f0d4ef81701 100644
--- a/experimental/air/cmd/rundockerimage_test.go
+++ b/experimental/air/cmd/rundockerimage_test.go
@@ -1,10 +1,12 @@
 package aircmd
 
 import (
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"testing"
 
+	"github.com/databricks/cli/libs/cmdio"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -17,13 +19,13 @@ func TestDigestDisplay(t *testing.T) {
 
 func TestWaitForRegisteredImageAvailable(t *testing.T) {
 	url := imageServer(t, `{}`, `{"state":"AVAILABLE","manifest_sha256":"abc"}`)
-	require.NoError(t, waitForRegisteredImage(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0"))
+	require.NoError(t, waitForRegisteredImage(cmdio.MockDiscard(t.Context()), newTestImageClient(t, url), "nvcr.io/org/img:1.0"))
 }
 
 func TestWaitForRegisteredImageNotRegistered(t *testing.T) {
 	// :get 404s, so the run must stop with registration guidance.
 	url := imageServer(t, `{}`, "")
-	err := waitForRegisteredImage(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0")
+	err := waitForRegisteredImage(cmdio.MockDiscard(t.Context()), newTestImageClient(t, url), "nvcr.io/org/img:1.0")
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "docker image not registered")
 	assert.Contains(t, err.Error(), "air register-image nvcr.io/org/img:1.0")
@@ -31,7 +33,7 @@ func TestWaitForRegisteredImageNotRegistered(t *testing.T) {
 
 func TestWaitForRegisteredImageFailed(t *testing.T) {
 	url := imageServer(t, `{}`, `{"state":"FAILED","status_message":"manifest not found"}`)
-	err := waitForRegisteredImage(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0")
+	err := waitForRegisteredImage(cmdio.MockDiscard(t.Context()), newTestImageClient(t, url), "nvcr.io/org/img:1.0")
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "registration failed: manifest not found")
 	assert.Contains(t, err.Error(), "re-register")
@@ -42,7 +44,7 @@ func TestWaitForRegisteredImageWaitsWhileImporting(t *testing.T) {
 	url := imageServer(t, `{}`,
 		`{"state":"IMPORTING"}`,
 		`{"state":"AVAILABLE","manifest_sha256":"abc"}`)
-	require.NoError(t, waitForRegisteredImage(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0"))
+	require.NoError(t, waitForRegisteredImage(cmdio.MockDiscard(t.Context()), newTestImageClient(t, url), "nvcr.io/org/img:1.0"))
 }
 
 // latestImageServer serves a registry where the image is registered and
@@ -69,10 +71,77 @@ func TestPrepareDockerImageAutoDoesNotReregister(t *testing.T) {
 	var posts int
 	w := newTestWorkspaceClient(t, latestImageServer(t, &posts))
 	img := &dockerImageConfig{URL: "nvcr.io/org/img:1.0"}
-	require.NoError(t, prepareDockerImage(t.Context(), w, img))
+	require.NoError(t, prepareDockerImage(cmdio.MockDiscard(t.Context()), w, img))
 	assert.Zero(t, posts, "default tag policy must not re-register")
 }
 
+// TestPrepareDockerImageLatestDiscoversCredentials covers the config-without-
+// credentials path: creds come from the local Docker config and ride the POST.
+func TestPrepareDockerImageLatestDiscoversCredentials(t *testing.T) {
+	var postBodies []string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case imagesAPIPath + ":get":
+			_, _ = w.Write([]byte(`{"state":"AVAILABLE","manifest_sha256":"abc"}`))
+		case imagesAPIPath + ":checkImageAccess":
+			_, _ = w.Write([]byte(`{"publicly_accessible": false}`))
+		case imagesAPIPath:
+			body, _ := io.ReadAll(r.Body)
+			postBodies = append(postBodies, string(body))
+			_, _ = w.Write([]byte(`{"image":{"state":"AVAILABLE","manifest_sha256":"abc"}}`))
+		case "/api/2.0/secrets/scopes/list":
+			_, _ = w.Write([]byte(`{"scopes":[]}`))
+		default:
+			_, _ = w.Write([]byte(`{}`))
+		}
+	}))
+	t.Cleanup(srv.Close)
+
+	ctx := cmdio.MockDiscard(writeDockerConfig(t, `{"auths":{"nvcr.io":{"auth":"`+b64(t, "bob:secret")+`"}}}`))
+	w := newTestWorkspaceClient(t, srv.URL)
+	img := &dockerImageConfig{URL: "nvcr.io/org/img:1.0", TagPolicy: "latest"}
+
+	require.NoError(t, prepareDockerImage(ctx, w, img))
+	require.Len(t, postBodies, 1)
+	assert.Contains(t, postBodies[0], `"credentials_key":"nvcr.io-bob-local"`)
+	assert.Contains(t, postBodies[0], `"credentials_scope":"docker-credentials-`)
+}
+
+// TestPrepareDockerImageLatestStorageDeniedReportsCause covers the journey where
+// creds exist locally but can't be stored: the error must name that cause rather
+// than tell the user to `docker login`.
+func TestPrepareDockerImageLatestStorageDeniedReportsCause(t *testing.T) {
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case imagesAPIPath + ":get":
+			w.WriteHeader(http.StatusNotFound)
+			_, _ = w.Write([]byte(`{"error_code":"NOT_FOUND","message":"not registered"}`))
+		case imagesAPIPath + ":checkImageAccess":
+			_, _ = w.Write([]byte(`{"publicly_accessible": false}`))
+		case "/api/2.0/secrets/scopes/create":
+			w.WriteHeader(http.StatusForbidden)
+			_, _ = w.Write([]byte(`{"error_code":"PERMISSION_DENIED","message":"denied"}`))
+		case imagesAPIPath:
+			w.WriteHeader(http.StatusForbidden)
+			_, _ = w.Write([]byte(`{"error_code":"PERMISSION_DENIED","message":"cannot pull: unauthorized"}`))
+		case "/api/2.0/secrets/scopes/list":
+			_, _ = w.Write([]byte(`{"scopes":[]}`))
+		default:
+			_, _ = w.Write([]byte(`{}`))
+		}
+	}))
+	t.Cleanup(srv.Close)
+
+	ctx := cmdio.MockDiscard(writeDockerConfig(t, `{"auths":{"nvcr.io":{"auth":"`+b64(t, "bob:secret")+`"}}}`))
+	w := newTestWorkspaceClient(t, srv.URL)
+	img := &dockerImageConfig{URL: "nvcr.io/org/img:1.0", TagPolicy: "latest"}
+
+	err := prepareDockerImage(ctx, w, img)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "could not be stored")
+	assert.NotContains(t, err.Error(), "run `docker login`")
+}
+
 func TestPrepareDockerImageLatestReregisters(t *testing.T) {
 	var posts int
 	w := newTestWorkspaceClient(t, latestImageServer(t, &posts))
@@ -82,6 +151,6 @@ func TestPrepareDockerImageLatestReregisters(t *testing.T) {
 		CredentialsScope: "scope",
 		CredentialsKey:   "key",
 	}
-	require.NoError(t, prepareDockerImage(t.Context(), w, img))
+	require.NoError(t, prepareDockerImage(cmdio.MockDiscard(t.Context()), w, img))
 	assert.Equal(t, 1, posts, "tag_policy=latest must re-register against the source registry")
 }

From 74954e861229dd93588b281aab729a68fdf8855b Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Tue, 4 Aug 2026 00:07:07 +0000
Subject: [PATCH 6/7] [air][m6-2] Address review nits

  - 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
---
 experimental/air/cmd/register_image.go      | 65 ++++++++++++---------
 experimental/air/cmd/register_image_test.go | 14 ++++-
 experimental/air/cmd/runconfig.go           |  6 ++
 experimental/air/cmd/runconfig_test.go      | 12 +++-
 experimental/air/cmd/rundockerimage.go      | 19 ++++--
 experimental/air/cmd/rundockerimage_test.go | 39 +++++++++++++
 6 files changed, 119 insertions(+), 36 deletions(-)

diff --git a/experimental/air/cmd/register_image.go b/experimental/air/cmd/register_image.go
index d4195d0d8e2..c3537102a00 100644
--- a/experimental/air/cmd/register_image.go
+++ b/experimental/air/cmd/register_image.go
@@ -111,18 +111,15 @@ configuration (run ` + "`docker login`" + ` first); there are no credential flag
 
 		timeout := time.Duration(timeoutMinutes) * time.Minute
 
-		// Discover credentials from the local Docker config and store them in a
-		// per-user secret for the registration call. If storage fails, registration
-		// still proceeds without credentials — a public image succeeds — and
-		// credErr is reported as the cause if the registry rejects anonymous access.
-		// credErr is not fatal on its own: it is reported by registrationError only
-		// if the registry then rejects anonymous access.
-		scope, key, credErr := discoverCredentials(ctx, w, c, dockerImageURL)
+		// credErr is not fatal on its own: registration proceeds without
+		// credentials (a public image still succeeds) and registrationError reports
+		// it only if the registry rejects anonymous access.
+		creds, credErr := discoverCredentials(ctx, w, c, dockerImageURL)
 		if credErr != nil {
 			log.Debugf(ctx, "could not store local Docker credentials: %v", credErr)
 		}
 
-		updated, sha, err := registerWithCredentialFallback(ctx, c, dockerImageURL, scope, key, timeout)
+		updated, sha, err := registerWithCredentialFallback(ctx, c, dockerImageURL, creds, timeout)
 		if err != nil {
 			kind, retryable := classifyRegistrationError(err)
 			return renderError(ctx, cmd, "REGISTRATION_FAILED", kind, retryable,
@@ -142,16 +139,25 @@ configuration (run ` + "`docker login`" + ` first); there are no credential flag
 	return cmd
 }
 
+// imageCredentials is a secret reference passed to registration. discovered
+// marks it as auto-discovered from the local Docker config rather than
+// configured by the user, which decides whether a rejection may be retried
+// anonymously.
+type imageCredentials struct {
+	scope      string
+	key        string
+	discovered bool
+}
+
 // discoverCredentials resolves registry credentials from the local Docker config
-// and stores them in a per-user secret, returning the (scope, key) reference for
-// registration. It first probes whether the image is public: if so, no
-// credentials are stored (avoiding a throwaway secret). Returns empty scope/key
-// when the image is public or no local credentials exist, both with a nil error.
-// A non-nil error means credentials were found but could not be stored (e.g. the
-// user lacks permission to create a secret scope); it is advisory, so the caller
-// can still attempt an anonymous registration and report this as the cause if
-// that fails.
-func discoverCredentials(ctx context.Context, w *databricks.WorkspaceClient, c *imageClient, dockerImageURL string) (scope, key string, err error) {
+// and stores them in a per-user secret, returning the reference for registration.
+// It first probes whether the image is public: if so, no credentials are stored
+// (avoiding a throwaway secret). Returns empty credentials when the image is
+// public or no local credentials exist, both with a nil error. A non-nil error
+// means credentials were found but could not be stored (e.g. the user lacks
+// permission to create a secret scope); it is advisory, so the caller can still
+// attempt an anonymous registration and report this as the cause if that fails.
+func discoverCredentials(ctx context.Context, w *databricks.WorkspaceClient, c *imageClient, dockerImageURL string) (creds imageCredentials, err error) {
 	// readDockerCredentials keys off the registry host, so it needs the normalized
 	// URL (e.g. bare "ubuntu" resolves to the Docker Hub host).
 	normalized := normalizeDockerImageURL(dockerImageURL)
@@ -161,20 +167,20 @@ func discoverCredentials(ctx context.Context, w *databricks.WorkspaceClient, c *
 	// twice.
 	username, password, ok := readDockerCredentials(ctx, normalized)
 	if !ok {
-		return "", "", nil
+		return imageCredentials{}, nil
 	}
 
 	if public := c.checkImageAccess(ctx, dockerImageURL); public != nil && *public {
 		log.Infof(ctx, "image is publicly accessible; skipping local Docker credentials")
-		return "", "", nil
+		return imageCredentials{}, nil
 	}
 
-	scope, key, err = storeDockerCredentials(ctx, w, normalized, username, password)
+	scope, key, err := storeDockerCredentials(ctx, w, normalized, username, password)
 	if err != nil {
-		return "", "", err
+		return imageCredentials{}, err
 	}
 	log.Infof(ctx, "using Docker credentials from local config (stored as %s/%s)", scope, key)
-	return scope, key, nil
+	return imageCredentials{scope: scope, key: key, discovered: true}, nil
 }
 
 // isAuthError reports whether err is an authentication or permission failure,
@@ -218,14 +224,15 @@ func registrationError(dockerImageURL string, err, credErr error) error {
 	return fmt.Errorf("image %q was not found or requires credentials: run `docker login` for its registry, then retry: %w", dockerImageURL, err)
 }
 
-// registerWithCredentialFallback registers the image and, if the stored
+// registerWithCredentialFallback registers the image and, if auto-discovered
 // credentials are rejected as an auth failure, retries once anonymously so a
-// public image isn't blocked by stale local creds (e.g. a revoked PAT from an
-// old `docker login`). The retry only fires when credentials were supplied.
-func registerWithCredentialFallback(ctx context.Context, c *imageClient, dockerImageURL, scope, key string, timeout time.Duration) (updated bool, sha string, err error) {
-	updated, sha, err = resolveImage(ctx, c, dockerImageURL, scope, key, timeout)
-	if err != nil && scope != "" && isAuthError(err) {
-		log.Warnf(ctx, "stored Docker credentials were rejected (%v); retrying without credentials in case the image is public", err)
+// public image isn't blocked by stale local creds (e.g. a revoked PAT from an old
+// `docker login`). Credentials the user configured explicitly are never retried
+// away: they asked for those specifically, so the rejection is the real answer.
+func registerWithCredentialFallback(ctx context.Context, c *imageClient, dockerImageURL string, creds imageCredentials, timeout time.Duration) (updated bool, sha string, err error) {
+	updated, sha, err = resolveImage(ctx, c, dockerImageURL, creds.scope, creds.key, timeout)
+	if err != nil && creds.discovered && isAuthError(err) {
+		log.Warnf(ctx, "Docker credentials discovered from your local config were rejected (%v); retrying without credentials in case the image is public", err)
 		return resolveImage(ctx, c, dockerImageURL, "", "", timeout)
 	}
 	return updated, sha, err
diff --git a/experimental/air/cmd/register_image_test.go b/experimental/air/cmd/register_image_test.go
index 8ad5383940b..5b7adb8779a 100644
--- a/experimental/air/cmd/register_image_test.go
+++ b/experimental/air/cmd/register_image_test.go
@@ -167,7 +167,7 @@ func credRejectingImageServer(t *testing.T, credentialedPOSTs *int) string {
 func TestRegisterWithCredentialFallbackRetriesAnonymously(t *testing.T) {
 	var credentialedPOSTs int
 	url := credRejectingImageServer(t, &credentialedPOSTs)
-	updated, sha, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", "scope", "key", time.Second)
+	updated, sha, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", imageCredentials{scope: "scope", key: "key", discovered: true}, time.Second)
 	require.NoError(t, err)
 	assert.True(t, updated)
 	assert.Equal(t, "pubsha", sha)
@@ -179,7 +179,17 @@ func TestRegisterWithCredentialFallbackNoRetryWithoutCreds(t *testing.T) {
 	// failure surfaces directly.
 	var credentialedPOSTs int
 	url := credRejectingImageServer(t, &credentialedPOSTs)
-	_, _, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", "", "", time.Second)
+	_, _, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", imageCredentials{}, time.Second)
 	require.NoError(t, err) // anonymous POST succeeds on this server
 	assert.Equal(t, 0, credentialedPOSTs)
 }
+
+func TestRegisterWithCredentialFallbackNoRetryForExplicitCreds(t *testing.T) {
+	// The user named these credentials, so a rejection is the real answer: do not
+	// silently retry without them.
+	var credentialedPOSTs int
+	url := credRejectingImageServer(t, &credentialedPOSTs)
+	_, _, err := registerWithCredentialFallback(t.Context(), newTestImageClient(t, url), "nvcr.io/org/img:1.0", imageCredentials{scope: "scope", key: "key"}, time.Second)
+	require.Error(t, err)
+	assert.Equal(t, 1, credentialedPOSTs, "should try once with creds and stop")
+}
diff --git a/experimental/air/cmd/runconfig.go b/experimental/air/cmd/runconfig.go
index 9a6a5fb9c80..96dcf96d5a1 100644
--- a/experimental/air/cmd/runconfig.go
+++ b/experimental/air/cmd/runconfig.go
@@ -339,6 +339,12 @@ func (d *dockerImageConfig) validate() error {
 	if (d.CredentialsScope != "") != (d.CredentialsKey != "") {
 		return errors.New("docker_image.credentials_scope and docker_image.credentials_key must be provided together")
 	}
+
+	// Credentials are only consulted when re-resolving the tag, so accepting them
+	// under the default policy would silently ignore them.
+	if d.CredentialsScope != "" && !d.wantsLatest() {
+		return fmt.Errorf("docker_image.credentials_scope/credentials_key only apply with tag_policy %q; the image is otherwise used as already registered", dockerTagPolicyLatest)
+	}
 	return nil
 }
 
diff --git a/experimental/air/cmd/runconfig_test.go b/experimental/air/cmd/runconfig_test.go
index af263d9581a..cd7804cd4af 100644
--- a/experimental/air/cmd/runconfig_test.go
+++ b/experimental/air/cmd/runconfig_test.go
@@ -313,8 +313,18 @@ func TestEnvironmentConfigValidate(t *testing.T) {
 			"must be provided together",
 		},
 		{
-			"credentials pair ok",
+			"credentials pair ok with latest",
+			environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", TagPolicy: "latest", CredentialsScope: "s", CredentialsKey: "k"}},
+			"",
+		},
+		{
+			"credentials without latest rejected",
 			environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", CredentialsScope: "s", CredentialsKey: "k"}},
+			`only apply with tag_policy "latest"`,
+		},
+		{
+			"blank credentials scope is not treated as set",
+			environmentConfig{DockerImage: &dockerImageConfig{URL: "org/repo:tag", CredentialsScope: "  "}},
 			"",
 		},
 		{
diff --git a/experimental/air/cmd/rundockerimage.go b/experimental/air/cmd/rundockerimage.go
index 4f624776f50..9645d367770 100644
--- a/experimental/air/cmd/rundockerimage.go
+++ b/experimental/air/cmd/rundockerimage.go
@@ -47,10 +47,12 @@ func prepareDockerImage(ctx context.Context, w *databricks.WorkspaceClient, img
 // newest digest, using the config's credentials when set and otherwise the local
 // Docker config.
 func resolveLatestDockerImage(ctx context.Context, w *databricks.WorkspaceClient, c *imageClient, img *dockerImageConfig) error {
-	scope, key := img.CredentialsScope, img.CredentialsKey
+	// Credentials the user configured are used as-is (discovered=false), so a
+	// rejection is not retried anonymously and reports the secret they named.
+	creds := imageCredentials{scope: img.CredentialsScope, key: img.CredentialsKey}
 	var credErr error
-	if scope == "" {
-		scope, key, credErr = discoverCredentials(ctx, w, c, img.URL)
+	if creds.scope == "" {
+		creds, credErr = discoverCredentials(ctx, w, c, img.URL)
 		if credErr != nil {
 			log.Debugf(ctx, "could not store local Docker credentials: %v", credErr)
 		}
@@ -58,7 +60,10 @@ func resolveLatestDockerImage(ctx context.Context, w *databricks.WorkspaceClient
 
 	// Visible at the default log level (WARN): this can block for minutes.
 	cmdio.LogString(ctx, fmt.Sprintf("Re-resolving %s against the source registry (tag_policy=latest)...", img.URL))
-	if _, _, err := registerWithCredentialFallback(ctx, c, img.URL, scope, key, imageReadyTimeout); err != nil {
+	if _, _, err := registerWithCredentialFallback(ctx, c, img.URL, creds, imageReadyTimeout); err != nil {
+		if !creds.discovered && creds.scope != "" && isAuthError(err) {
+			return fmt.Errorf("the credentials in secret %s/%s were rejected for image %q: %w", creds.scope, creds.key, img.URL, err)
+		}
 		return registrationError(img.URL, err, credErr)
 	}
 	return nil
@@ -66,6 +71,12 @@ func resolveLatestDockerImage(ctx context.Context, w *databricks.WorkspaceClient
 
 // waitForRegisteredImage requires an existing registration and blocks while it is
 // still importing. Missing or FAILED is an error the user fixes by re-registering.
+//
+// An AVAILABLE registration whose stored credentials have since lost registry
+// access is accepted here and only fails at pod launch. The Python CLI catches
+// that with a :validateImageAccess probe (cli/sdk/_submit.py), which is not ported
+// deliberately: this check belongs in the backend, so every client gets it and
+// the CLI doesn't pay a round trip per submit. Port it there rather than here.
 func waitForRegisteredImage(ctx context.Context, c *imageClient, dockerImageURL string) error {
 	reg, err := c.getImage(ctx, dockerImageURL)
 	if err != nil {
diff --git a/experimental/air/cmd/rundockerimage_test.go b/experimental/air/cmd/rundockerimage_test.go
index f0d4ef81701..362e0030a3c 100644
--- a/experimental/air/cmd/rundockerimage_test.go
+++ b/experimental/air/cmd/rundockerimage_test.go
@@ -4,6 +4,7 @@ import (
 	"io"
 	"net/http"
 	"net/http/httptest"
+	"strings"
 	"testing"
 
 	"github.com/databricks/cli/libs/cmdio"
@@ -142,6 +143,44 @@ func TestPrepareDockerImageLatestStorageDeniedReportsCause(t *testing.T) {
 	assert.NotContains(t, err.Error(), "run `docker login`")
 }
 
+// TestPrepareDockerImageLatestExplicitCredsRejected asserts a rejected secret the
+// user named reports that secret, and is not retried anonymously or blamed on a
+// missing `docker login`.
+func TestPrepareDockerImageLatestExplicitCredsRejected(t *testing.T) {
+	var credentialedPOSTs int
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		switch r.URL.Path {
+		case imagesAPIPath + ":get":
+			w.WriteHeader(http.StatusNotFound)
+			_, _ = w.Write([]byte(`{"error_code":"NOT_FOUND","message":"not registered"}`))
+		case imagesAPIPath:
+			body, _ := io.ReadAll(r.Body)
+			if strings.Contains(string(body), "credentials_scope") {
+				credentialedPOSTs++
+			}
+			w.WriteHeader(http.StatusForbidden)
+			_, _ = w.Write([]byte(`{"error_code":"PERMISSION_DENIED","message":"denied"}`))
+		default:
+			_, _ = w.Write([]byte(`{}`))
+		}
+	}))
+	t.Cleanup(srv.Close)
+
+	w := newTestWorkspaceClient(t, srv.URL)
+	img := &dockerImageConfig{
+		URL:              "nvcr.io/org/img:1.0",
+		TagPolicy:        "latest",
+		CredentialsScope: "myscope",
+		CredentialsKey:   "mykey",
+	}
+
+	err := prepareDockerImage(cmdio.MockDiscard(t.Context()), w, img)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "credentials in secret myscope/mykey were rejected")
+	assert.NotContains(t, err.Error(), "docker login")
+	assert.Equal(t, 1, credentialedPOSTs, "explicit credentials must not be retried anonymously")
+}
+
 func TestPrepareDockerImageLatestReregisters(t *testing.T) {
 	var posts int
 	w := newTestWorkspaceClient(t, latestImageServer(t, &posts))

From 45978376b916ff32d1610d29529b9c329416cabe Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Tue, 11 Aug 2026 18:38:16 +0000
Subject: [PATCH 7/7] [air][m6-2] Defer docker_image_url payload field to a
 follow-up
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The docker_image_url field on jobs.AiRuntimeTask was added to databricks-sdk-go
after v0.170.0; the CLI is pinned to v0.166.0 and moving to the newer SDK needs a
bump + full command-stub regen, which is a separate PR. Setting the field here
would not compile until then.

Leave a NOTE where the field will be set and drop the two payload tests that
assert on it. Everything else in the docker-image path — config parsing and the
prepareDockerImage preflight (verify registered, wait while importing, re-resolve
under tag_policy=latest) — is unaffected and still runs.

Co-authored-by: Isaac
---
 experimental/air/cmd/runsubmit.go      |  7 ++++--
 experimental/air/cmd/runsubmit_test.go | 31 --------------------------
 2 files changed, 5 insertions(+), 33 deletions(-)

diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go
index 4310b117749..54189fba956 100644
--- a/experimental/air/cmd/runsubmit.go
+++ b/experimental/air/cmd/runsubmit.go
@@ -80,8 +80,11 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri
 			},
 		}},
 		CodeSourcePath: snap.CodeSourcePath,
-		// Verified as registered by prepareDockerImage.
-		DockerImageUrl: cfg.dockerImageURL(),
+		// NOTE: docker_image_url is intentionally not set here yet. The field was
+		// added to jobs.AiRuntimeTask in databricks-sdk-go after v0.170.0, which the
+		// CLI has not bumped to. prepareDockerImage already verifies the image is
+		// registered; passing it on the task lands in the follow-up PR once the SDK
+		// bump + codegen is in.
 	}
 	if cfg.MLflowRunName != nil {
 		task.MlflowRun = *cfg.MLflowRunName
diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go
index 3c4fd6f2798..76ab3aeb3ac 100644
--- a/experimental/air/cmd/runsubmit_test.go
+++ b/experimental/air/cmd/runsubmit_test.go
@@ -71,37 +71,6 @@ func TestBuildSubmitPayload(t *testing.T) {
 	assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu8xH100, AcceleratorCount: 16}, at.Deployments[0].Compute)
 }
 
-func TestBuildSubmitPayloadDockerImage(t *testing.T) {
-	cfg := &runConfig{
-		ExperimentName: "exp",
-		Command:        new("x"),
-		Compute:        &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1},
-		Environment: &environmentConfig{
-			DockerImage: &dockerImageConfig{URL: "nvcr.io/org/img:1.0"},
-		},
-	}
-
-	p := buildSubmitPayload(cfg, "/d/command.sh", "5", snapshotResult{}, nil)
-	require.Len(t, p.Tasks, 1)
-	require.NotNil(t, p.Tasks[0].AiRuntimeTask)
-	assert.Equal(t, "nvcr.io/org/img:1.0", p.Tasks[0].AiRuntimeTask.DockerImageUrl)
-}
-
-func TestBuildSubmitPayloadNoDockerImage(t *testing.T) {
-	// Without an environment.docker_image block the field stays empty (omitempty),
-	// so the runtime-managed environment is used.
-	cfg := &runConfig{
-		ExperimentName: "exp",
-		Command:        new("x"),
-		Compute:        &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1},
-	}
-
-	p := buildSubmitPayload(cfg, "/d/command.sh", "5", snapshotResult{}, nil)
-	require.Len(t, p.Tasks, 1)
-	require.NotNil(t, p.Tasks[0].AiRuntimeTask)
-	assert.Empty(t, p.Tasks[0].AiRuntimeTask.DockerImageUrl)
-}
-
 func TestBuildSubmitPayloadDefaultRetries(t *testing.T) {
 	// max_retries unset defaults to 3 (matching the Python native path), so both
 	// retry fields are sent.