AIR CLI: attach docker_image_url to the run (follow-up to #6165)

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

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.