diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index b32ebd622c6..49f51a8aa66 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/experimental/air/cmd/runsubmit.go @@ -66,6 +66,13 @@ type jobEnvironment struct { } // submitTask is the single task air submits: a native ai_runtime_task. +// +// max_retries is always sent (including 0) so the user's YAML value is honored: +// setting it to 0 explicitly disables retries rather than falling back to the +// server default. retry_on_timeout is sent only when retries are allowed, and is +// omitempty so the wire form matches the Python CLI (which never emits a bare +// "false"). Jobs performs the retries — each attempt is a fresh AI Runtime +// workload. type submitTask struct { TaskKey string `json:"task_key"` RunIf string `json:"run_if"` @@ -125,8 +132,9 @@ func buildSubmitPayload(cfg *runConfig, commandPath, dlImage string) jobsSubmitR EnvironmentKey: aiRuntimeEnvironmentKey, MaxRetries: cfg.maxRetries(), } - // max_retries 0 (no retries) is sent explicitly; retry_on_timeout only - // applies when retries are allowed. + // retry_on_timeout only makes sense when retries are allowed; otherwise omit + // it (matches Python's native path, which sets retry_on_timeout only under + // the same > 0 gate). st.RetryOnTimeout = st.MaxRetries > 0 return jobsSubmitRun{ diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index bfd92dcd58e..d57b14f80a6 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -61,22 +61,37 @@ func TestBuildSubmitPayload(t *testing.T) { assert.Equal(t, aiRuntimeCompute{AcceleratorType: "GPU_8xH100", AcceleratorCount: 16}, at.Deployments[0].Compute) } -func TestBuildSubmitPayload_NoRetries(t *testing.T) { +func TestBuildSubmitPayloadDefaultRetries(t *testing.T) { + // max_retries unset defaults to 3 (matching the Python native path), so both + // retry fields are sent. cfg := &runConfig{ ExperimentName: "exp", Command: new("x"), Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, - MaxRetries: new(0), } + task := buildSubmitPayload(cfg, "/d/command.sh", "4").Tasks[0] + assert.Equal(t, defaultMaxRetries, task.MaxRetries) + assert.True(t, task.RetryOnTimeout) +} +func TestBuildSubmitPayloadNoRetries(t *testing.T) { + // max_retries: 0 must be sent explicitly so Jobs honors "no retries" instead + // of applying the server default. retry_on_timeout is omitted when retries + // aren't allowed. + cfg := &runConfig{ + ExperimentName: "exp", + Command: new("x"), + Compute: &computeConfig{AcceleratorType: "GPU_1xH100", NumAccelerators: 1}, + MaxRetries: new(0), + } task := buildSubmitPayload(cfg, "/d/command.sh", "4").Tasks[0] assert.Equal(t, 0, task.MaxRetries) assert.False(t, task.RetryOnTimeout) - // max_retries: 0 must be sent, not omitted, so the server honors "no retries". b, err := json.Marshal(task) require.NoError(t, err) assert.Contains(t, string(b), `"max_retries":0`) + assert.NotContains(t, string(b), "retry_on_timeout") } func TestSubmitToken(t *testing.T) {