diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 54189fba95..4310b11774 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 76ab3aeb3a..9781b28101 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.