From f17e6884fd997a6ad7fb71e6771dfdfabd284391 Mon Sep 17 00:00:00 2001
From: riddhibhagwat-db <riddhi.bhagwat@databricks.com>
Date: Thu, 2 Jul 2026 18:50:03 +0000
Subject: [PATCH] experimental/air: send max_retries explicitly so 0 disables
 retries
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The ai_runtime_task path DOES honor the Jobs task max_retries field —
Jobs performs the retry and each attempt is created as a fresh AI Runtime
(AICM) workload (see ai-training AiTrainingHandlers: "Each Jobs retry
attempt is created as its own AICM workload"). So the retry knob is real.

Send max_retries whenever it is set (including 0), so the user's YAML
value is honored end-to-end:

- max_retries unset -> defaults to 3 (matches Python: sdk/config.py
  Field(default=3)).
- max_retries: N>0 -> sent, retry_on_timeout: true.
- max_retries: 0 -> sent explicitly ("max_retries":0), retry_on_timeout
  omitted. This diverges from the Python CLI (which drops the field
  under `if max_retries > 0`) so that 0 actually means "no retries"
  instead of silently falling back to the server default.

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

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) {