From 1e780a03351a8e2b5f23f310876846b34ebc10ff Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Tue, 11 Aug 2026 18:40:23 +0000
Subject: [PATCH] [air][m6-3] Pass docker_image_url on the AI runtime task

Follow-up to the docker-image wiring: set ai_runtime_task.docker_image_url from
environment.docker_image.url so a registered custom image is actually used by
the run (the parent PR verifies the image but does not yet attach it).

REQUIRES an SDK bump: docker_image_url was added to jobs.AiRuntimeTask in
databricks-sdk-go after v0.170.0, and the CLI is pinned to v0.166.0. This one
line does not compile until the CLI bumps the SDK (a separate PR that also
regenerates the command stubs). Merge only after that bump lands; the code was
verified green against the SDK commit that has the field.

Re-adds the two buildSubmitPayload tests covering the field set and omitted.

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

diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go
index 54189fba956..4310b117749 100644
--- a/experimental/air/cmd/runsubmit.go
+++ b/experimental/air/cmd/runsubmit.go
@@ -80,11 +80,8 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage, usagePolicyID stri
 			},
 		}},
 		CodeSourcePath: snap.CodeSourcePath,
-		// 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.
+		// Verified as registered by prepareDockerImage.
+		DockerImageUrl: cfg.dockerImageURL(),
 	}
 	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 76ab3aeb3ac..9781b281010 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.