From 755eb96310ceff5dedae7c244b04968b49207517 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Wed, 4 Nov 2026 22:56:01 +0000 Subject: [PATCH 1/7] Run AIR preflight validation in parallel with artifact staging `databricks air run` used to wait for the server-side config:validate call before staging the launch artifacts. The backend check doesn't read the uploaded files, so run it concurrently with usage-policy resolution, the experiment directory check, mkdirs, and the uploads. Submission still waits for both, and a failure on either side cancels the other, so an invalid config is never submitted (it may now leave staged artifacts). The "validation unavailable" warning is printed after staging so it doesn't interleave with the staging spinner. Co-authored-by: Isaac --- .../air/parallel-preflight-validation.md | 1 + cmd/air/runsubmit.go | 92 +++++++++++-------- cmd/air/runsubmit_test.go | 46 +++++++++- cmd/air/validateconfig.go | 19 ++-- cmd/air/validateconfig_test.go | 41 +++++---- 5 files changed, 131 insertions(+), 68 deletions(-) create mode 100644 .nextchanges/air/parallel-preflight-validation.md diff --git a/.nextchanges/air/parallel-preflight-validation.md b/.nextchanges/air/parallel-preflight-validation.md new file mode 100644 index 00000000000..e3fca898b67 --- /dev/null +++ b/.nextchanges/air/parallel-preflight-validation.md @@ -0,0 +1 @@ +* `databricks air run` now validates the config server-side while it uploads the run artifacts, instead of before, which shortens submission. diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index d1e94e068dc..b514574557e 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -327,66 +327,38 @@ func stageRunArtifacts(ctx context.Context, launchWriter fileWriter, items []upl return snap, err } -// submitWorkload runs the submit happy path: ensure the experiment directory, -// upload the launch artifacts, assemble the Jobs payload, and submit it. It -// returns the new run_id and its dashboard URL. showProgress enables the stderr -// staging spinner (text mode only). -func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (runID int64, dashboardURL string, err error) { - start := time.Now() - var snap snapshotResult - defer func() { logRunEvent(ctx, cfg, snap, runID, time.Since(start), err) }() - - // Compute and validate the actual submission path before creating artifacts. - base, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg) - if err != nil { - return 0, "", err - } - containers := submittedContainers(cfg, funcDir) - token := effectiveIdempotencyToken(idempotencyKey, cfg) - - // Validate server-side before uploading, so invalid configs leave no artifacts. - if err := preflightValidate(ctx, w, cfg, commandPath, containers, token); err != nil { - return 0, "", err - } - - // Retain the local guard after backend validation so an unavailable endpoint - // cannot let a bad key reach upload or submission. - if err := validateIdempotencyToken(token); err != nil { - return 0, "", err - } - - // Resolve the usage policy to its id next, so a bad name fails fast with a +// stageRun resolves the usage policy, ensures the experiment directory, and +// uploads the launch artifacts and optional code snapshot. +func stageRun(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, base, funcDir string, showProgress bool) (usagePolicyID string, snap snapshotResult, err error) { + // Resolve the usage policy to its id first, so a bad name fails fast with a // clear (caller-fixable) message before we upload any artifacts. Validation // guarantees name and id are mutually exclusive: a literal id is used as-is, a // name is resolved against the workspace. - usagePolicyID := "" if cfg.UsagePolicyID != nil { usagePolicyID = strings.TrimSpace(*cfg.UsagePolicyID) } if cfg.UsagePolicyName != nil { usagePolicyID, err = resolveUsagePolicyIDByName(ctx, w, *cfg.UsagePolicyName) if err != nil { - return 0, "", err + return "", snap, err } } - deps, _ := cfg.inlineDependencies() - experimentDir := "" if cfg.MLflowExperimentDirectory != nil { experimentDir = *cfg.MLflowExperimentDirectory } if err := ensureExperimentDirectory(ctx, w, experimentDir); err != nil { - return 0, "", err + return "", snap, err } fc, err := filer.NewWorkspaceFilesClient(w, funcDir) if err != nil { - return 0, "", err + return "", snap, err } items, err := buildArtifacts(cfg) if err != nil { - return 0, "", err + return "", snap, err } // Package and upload the code snapshot, if any; the remote code_source_path @@ -408,9 +380,57 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run snap, stageErr = stageRunArtifacts(ctx, fc, items, stageSnapshot) return stageErr }) + return usagePolicyID, snap, err +} + +// submitWorkload runs the submit happy path: ensure the experiment directory, +// upload the launch artifacts, assemble the Jobs payload, and submit it. It +// returns the new run_id and its dashboard URL. showProgress enables the stderr +// staging spinner (text mode only). +func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (runID int64, dashboardURL string, err error) { + start := time.Now() + var snap snapshotResult + defer func() { logRunEvent(ctx, cfg, snap, runID, time.Since(start), err) }() + + // Compute and validate the actual submission path before creating artifacts. + base, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg) if err != nil { return 0, "", err } + containers := submittedContainers(cfg, funcDir) + token := effectiveIdempotencyToken(idempotencyKey, cfg) + + // Validate server-side while the run artifacts stage. The backend check does + // not read the uploaded files, so the two can overlap. Submission waits for + // both, so an invalid config may leave staged artifacts but is never + // submitted. A failure on either side cancels the other. + group, groupCtx := errgroup.WithContext(ctx) + var validationUnavailable bool + group.Go(func() (err error) { + validationUnavailable, err = preflightValidate(groupCtx, w, cfg, commandPath, containers, token) + return err + }) + var usagePolicyID string + group.Go(func() (err error) { + usagePolicyID, snap, err = stageRun(groupCtx, w, cfg, configPath, base, funcDir, showProgress) + return err + }) + if err := group.Wait(); err != nil { + return 0, "", err + } + // Warn only after staging finishes so the warning doesn't interleave with the + // staging spinner. + if validationUnavailable && cmdio.HasIO(ctx) { + cmdio.LogString(ctx, "Warning: server-side config validation was unavailable; continuing with submission.") + } + + // Retain the local guard after backend validation so an unavailable endpoint + // cannot let a bad key reach submission. + if err := validateIdempotencyToken(token); err != nil { + return 0, "", err + } + + deps, _ := cfg.inlineDependencies() runtimeVersion, _ := cfg.runtimeVersion() payload := buildSubmitPayload(cfg, commandPath, dlRuntimeImage(ctx, runtimeVersion), usagePolicyID, snap, deps) diff --git a/cmd/air/runsubmit_test.go b/cmd/air/runsubmit_test.go index d722edac163..c3703e6a321 100644 --- a/cmd/air/runsubmit_test.go +++ b/cmd/air/runsubmit_test.go @@ -409,6 +409,7 @@ func TestSubmitWorkloadValidationOutcome(t *testing.T) { response any wantError string wantSubmit bool + wantStderr string }{ { name: "field error", @@ -421,6 +422,7 @@ func TestSubmitWorkloadValidationOutcome(t *testing.T) { name: "service unavailable", response: testserver.Response{StatusCode: 503, Body: map[string]string{"error_code": "TEMPORARILY_UNAVAILABLE", "message": "try again"}}, wantSubmit: true, + wantStderr: "Warning: server-side config validation was unavailable; continuing with submission.\n", }, } @@ -446,23 +448,63 @@ func TestSubmitWorkloadValidationOutcome(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) + ctx, stderr := cmdio.NewTestContextWithStderr(t.Context()) + _, _, err = submitWorkload(ctx, w, cfg, cfgPath, "idem-key", false) if tt.wantError != "" { require.ErrorContains(t, err, tt.wantError) } else { require.NoError(t, err) } + // Uploads overlap validation, so a rejected config may still have + // staged artifacts; only submission is gated on validation. if tt.wantSubmit { assert.Positive(t, uploads.Load()) assert.Equal(t, int32(1), submissions.Load()) } else { - assert.Zero(t, uploads.Load()) assert.Zero(t, submissions.Load()) } + assert.Equal(t, tt.wantStderr, stderr.String()) }) } } +// Validation and artifact staging must run concurrently: the validate handler +// only answers once an upload has arrived, which can't happen if staging waits +// for validation. +func TestSubmitWorkloadValidatesWhileStaging(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + uploadStarted := make(chan struct{}) + markUploadStarted := sync.OnceFunc(func() { close(uploadStarted) }) + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { + markUploadStarted() + return req.Workspace.WorkspaceFilesImportFile(req.Vars["path"], req.Body, true) + }) + server.Handle("POST", validateConfigPath, func(req testserver.Request) any { + select { + case <-uploadStarted: + return validateConfigResponse{} + case <-time.After(submissionValidationTimeout / 2): + return validateConfigResponse{Errors: []configFieldError{{Message: "staging did not overlap validation"}}} + } + }) + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + return jobs.SubmitRunResponse{RunId: 777} + }) + testserver.AddDefaultHandlers(server) + + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cfg, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + runID, _, err := submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) + require.NoError(t, err) + assert.Equal(t, int64(777), runID) +} + func TestSubmitWorkloadStagingErrorPreventsSubmit(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) diff --git a/cmd/air/validateconfig.go b/cmd/air/validateconfig.go index 62dad63f8ab..5033571bda2 100644 --- a/cmd/air/validateconfig.go +++ b/cmd/air/validateconfig.go @@ -12,7 +12,6 @@ import ( "time" "github.com/databricks/cli/libs/auth" - "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/client" @@ -86,22 +85,20 @@ func validationUnavailable(err error) (*validationUnavailableError, bool) { return errors.AsType[*validationUnavailableError](err) } -// preflightValidate checks the config against the backend before any upload, so -// a bad config fails fast with the server's own field-level errors. +// preflightValidate checks the config against the backend before submission, so +// a bad config fails with the server's own field-level errors. // -// Availability failures fail open, while caller failures and field errors block. -func preflightValidate(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, commandPath string, containers []submittedContainer, idempotencyToken string) error { +// Availability failures fail open: they report unavailable with a nil error, and +// the caller decides how to warn. Caller failures and field errors block. +func preflightValidate(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, commandPath string, containers []submittedContainer, idempotencyToken string) (unavailable bool, err error) { validationCtx, cancel := context.WithTimeout(ctx, submissionValidationTimeout) defer cancel() - err := validateConfigOnce(validationCtx, w, cfg, commandPath, containers, idempotencyToken) + err = validateConfigOnce(validationCtx, w, cfg, commandPath, containers, idempotencyToken) if unavailable, _ := classifyValidationFailure(err); unavailable { - if cmdio.HasIO(ctx) { - cmdio.LogString(ctx, "Warning: server-side config validation was unavailable; continuing with submission.") - } - return nil + return true, nil } - return err + return false, err } func validateConfigOnce(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, commandPath string, containers []submittedContainer, idempotencyToken string) error { diff --git a/cmd/air/validateconfig_test.go b/cmd/air/validateconfig_test.go index f0a0f3e38c8..dba9d60f420 100644 --- a/cmd/air/validateconfig_test.go +++ b/cmd/air/validateconfig_test.go @@ -13,7 +13,6 @@ import ( "testing" "time" - "github.com/databricks/cli/libs/cmdio" "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" sdkconfig "github.com/databricks/databricks-sdk-go/config" @@ -112,13 +111,14 @@ func TestPreflightValidateFailsOpenWhenBackendUnavailable(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { srv := validateServer(t, tt.status, tt.body, nil) - err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + unavailable, err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") assert.NoError(t, err) + assert.True(t, unavailable) }) } } -func TestPreflightValidationUnavailableWarnsWithoutRetry(t *testing.T) { +func TestPreflightValidationUnavailableDoesNotRetry(t *testing.T) { var requests atomic.Int32 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == validateConfigPath { @@ -130,12 +130,11 @@ func TestPreflightValidationUnavailableWarnsWithoutRetry(t *testing.T) { _, _ = w.Write([]byte(`{}`)) })) t.Cleanup(srv.Close) - ctx, stderr := cmdio.NewTestContextWithStderr(t.Context()) - err := preflightValidate(ctx, validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + unavailable, err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) + assert.True(t, unavailable) assert.Equal(t, int32(1), requests.Load()) - assert.Equal(t, "Warning: server-side config validation was unavailable; continuing with submission.\n", stderr.String()) } func TestPreflightValidationDoesNotRetryTransportFailure(t *testing.T) { @@ -147,8 +146,9 @@ func TestPreflightValidationDoesNotRetryTransportFailure(t *testing.T) { return nil, errors.New("connection failed") }) - err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + unavailable, err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) + assert.True(t, unavailable) assert.Equal(t, int32(1), requests.Load()) } @@ -178,8 +178,9 @@ func TestPreflightValidationFailsOpenOnInvalidSuccessResponse(t *testing.T) { }, nil }) - err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + unavailable, err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) + assert.True(t, unavailable) assert.Equal(t, int32(1), requests.Load()) }) } @@ -203,7 +204,7 @@ func TestPreflightValidationBodyReadFailureOnCallerErrorBlocks(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) defer cancel() - err := preflightValidate(ctx, w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(ctx, w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.Error(t, err) apiErr, ok := errors.AsType[*apierr.APIError](err) require.True(t, ok) @@ -239,9 +240,10 @@ func TestPreflightValidationOAuthRefreshUsesDeadline(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) defer cancel() started := time.Now() - err = preflightValidate(ctx, w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + unavailable, err := preflightValidate(ctx, w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) + assert.True(t, unavailable) assert.Less(t, time.Since(started), time.Second) select { case <-refreshStarted: @@ -285,7 +287,7 @@ func TestPreflightValidationM2MOAuthRejectionBlocks(t *testing.T) { }) require.NoError(t, err) - err = preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err = preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.ErrorIs(t, err, tt.wantErr) }) } @@ -305,9 +307,10 @@ func TestPreflightValidationTimeoutFailsOpen(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) defer cancel() started := time.Now() - err := preflightValidate(ctx, validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + unavailable, err := preflightValidate(ctx, validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") close(releaseRequest) require.NoError(t, err) + assert.True(t, unavailable) assert.Less(t, time.Since(started), time.Second) } @@ -323,7 +326,7 @@ func TestPreflightValidationPropagatesCancellation(t *testing.T) { ctx, cancel := context.WithCancel(t.Context()) cancel() - err := preflightValidate(ctx, validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(ctx, validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.ErrorIs(t, err, context.Canceled) } @@ -348,7 +351,7 @@ func TestPreflightValidationPropagatesCancellationWhileReadingErrorResponse(t *t <-validationStarted cancel() }() - err := preflightValidate(ctx, w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(ctx, w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.ErrorIs(t, err, context.Canceled) } @@ -356,7 +359,7 @@ func TestPreflightValidateBlocksOnCallerError(t *testing.T) { for _, status := range []int{http.StatusBadRequest, http.StatusUnauthorized, http.StatusForbidden} { t.Run(http.StatusText(status), func(t *testing.T) { srv := validateServer(t, status, `{"error_code":"INVALID_PARAMETER_VALUE","message":"request rejected"}`, nil) - err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.Error(t, err) }) } @@ -369,7 +372,7 @@ func TestPreflightValidateBlocksOnRequestVisitorError(t *testing.T) { return errors.New("failed to add request headers") } - err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.ErrorContains(t, err, "failed to add request headers") } @@ -412,7 +415,7 @@ func TestValidateConfigRequestShape(t *testing.T) { cfg.MaxRetries = new(3) cfg.EnvVariables = map[string]string{"HF_HOME": "/tmp/hf"} cfg.Environment = &environmentConfig{UnityCatalogImage: "main.air.training:v1"} - err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), cfg, "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), cfg, "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) task := gotReq["task"].(map[string]any) @@ -436,7 +439,7 @@ func TestValidateConfigRequestCarriesPriorityClass(t *testing.T) { cfg := baseRunConfig() cfg.Compute.PoolID = new("cap-8xh100-res") cfg.Compute.PriorityClass = new("CRITICAL") - err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), cfg, "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), cfg, "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) task := gotReq["task"].(map[string]any) @@ -450,7 +453,7 @@ func TestValidateConfigRequestOmitsUnsetOptions(t *testing.T) { var gotReq map[string]any srv := validateServer(t, http.StatusOK, `{}`, &gotReq) - err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) assert.Equal(t, map[string]any{"idempotency_token": "token"}, gotReq["run_options"]) From e034301e55a1e0c28b9523865e9a58efe82951a8 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Thu, 5 Nov 2026 00:10:00 +0000 Subject: [PATCH 2/7] Link PR in AIR changelog fragment Co-authored-by: Isaac --- .nextchanges/air/parallel-preflight-validation.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.nextchanges/air/parallel-preflight-validation.md b/.nextchanges/air/parallel-preflight-validation.md index e3fca898b67..adb46c09fc3 100644 --- a/.nextchanges/air/parallel-preflight-validation.md +++ b/.nextchanges/air/parallel-preflight-validation.md @@ -1 +1 @@ -* `databricks air run` now validates the config server-side while it uploads the run artifacts, instead of before, which shortens submission. +* `databricks air run` now validates the config server-side while it uploads the run artifacts, instead of before, which shortens submission. ([#7368](https://github.com/databricks/cli/pull/7368)) From 964c030e89a1f417ae721ffae82e05c2366de7d7 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Thu, 5 Nov 2026 04:31:59 +0000 Subject: [PATCH 3/7] Run preflight validation inline without moving submit code Start validation in a goroutine before staging and join it before submit, instead of extracting staging into a helper, so the diff doesn't move existing code. Co-authored-by: Isaac --- cmd/air/runsubmit.go | 92 +++++++++++++++++----------------- cmd/air/validateconfig_test.go | 15 ++---- 2 files changed, 52 insertions(+), 55 deletions(-) diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index b514574557e..3e77203937a 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -327,38 +327,73 @@ func stageRunArtifacts(ctx context.Context, launchWriter fileWriter, items []upl return snap, err } -// stageRun resolves the usage policy, ensures the experiment directory, and -// uploads the launch artifacts and optional code snapshot. -func stageRun(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, base, funcDir string, showProgress bool) (usagePolicyID string, snap snapshotResult, err error) { - // Resolve the usage policy to its id first, so a bad name fails fast with a +// submitWorkload runs the submit happy path: ensure the experiment directory, +// upload the launch artifacts, assemble the Jobs payload, and submit it. It +// returns the new run_id and its dashboard URL. showProgress enables the stderr +// staging spinner (text mode only). +func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (runID int64, dashboardURL string, err error) { + start := time.Now() + var snap snapshotResult + defer func() { logRunEvent(ctx, cfg, snap, runID, time.Since(start), err) }() + + // Compute and validate the actual submission path before creating artifacts. + base, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg) + if err != nil { + return 0, "", err + } + containers := submittedContainers(cfg, funcDir) + token := effectiveIdempotencyToken(idempotencyKey, cfg) + + // Validate server-side while the run artifacts stage below. The backend check + // does not read the uploaded files, so the two can overlap; submission waits + // for the result, so an invalid config may leave staged artifacts but is never + // submitted. + validationCtx, cancelValidation := context.WithCancel(ctx) + validated := make(chan struct{}) + var validationUnavailable bool + var validationErr error + go func() { + defer close(validated) + validationUnavailable, validationErr = preflightValidate(validationCtx, w, cfg, commandPath, containers, token) + }() + // Stop and wait for validation if staging fails first. + defer func() { + cancelValidation() + <-validated + }() + + // Resolve the usage policy to its id next, so a bad name fails fast with a // clear (caller-fixable) message before we upload any artifacts. Validation // guarantees name and id are mutually exclusive: a literal id is used as-is, a // name is resolved against the workspace. + usagePolicyID := "" if cfg.UsagePolicyID != nil { usagePolicyID = strings.TrimSpace(*cfg.UsagePolicyID) } if cfg.UsagePolicyName != nil { usagePolicyID, err = resolveUsagePolicyIDByName(ctx, w, *cfg.UsagePolicyName) if err != nil { - return "", snap, err + return 0, "", err } } + deps, _ := cfg.inlineDependencies() + experimentDir := "" if cfg.MLflowExperimentDirectory != nil { experimentDir = *cfg.MLflowExperimentDirectory } if err := ensureExperimentDirectory(ctx, w, experimentDir); err != nil { - return "", snap, err + return 0, "", err } fc, err := filer.NewWorkspaceFilesClient(w, funcDir) if err != nil { - return "", snap, err + return 0, "", err } items, err := buildArtifacts(cfg) if err != nil { - return "", snap, err + return 0, "", err } // Package and upload the code snapshot, if any; the remote code_source_path @@ -380,46 +415,15 @@ func stageRun(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig snap, stageErr = stageRunArtifacts(ctx, fc, items, stageSnapshot) return stageErr }) - return usagePolicyID, snap, err -} - -// submitWorkload runs the submit happy path: ensure the experiment directory, -// upload the launch artifacts, assemble the Jobs payload, and submit it. It -// returns the new run_id and its dashboard URL. showProgress enables the stderr -// staging spinner (text mode only). -func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (runID int64, dashboardURL string, err error) { - start := time.Now() - var snap snapshotResult - defer func() { logRunEvent(ctx, cfg, snap, runID, time.Since(start), err) }() - - // Compute and validate the actual submission path before creating artifacts. - base, funcDir, commandPath, err := prospectiveLaunchPaths(ctx, w, cfg) if err != nil { return 0, "", err } - containers := submittedContainers(cfg, funcDir) - token := effectiveIdempotencyToken(idempotencyKey, cfg) - // Validate server-side while the run artifacts stage. The backend check does - // not read the uploaded files, so the two can overlap. Submission waits for - // both, so an invalid config may leave staged artifacts but is never - // submitted. A failure on either side cancels the other. - group, groupCtx := errgroup.WithContext(ctx) - var validationUnavailable bool - group.Go(func() (err error) { - validationUnavailable, err = preflightValidate(groupCtx, w, cfg, commandPath, containers, token) - return err - }) - var usagePolicyID string - group.Go(func() (err error) { - usagePolicyID, snap, err = stageRun(groupCtx, w, cfg, configPath, base, funcDir, showProgress) - return err - }) - if err := group.Wait(); err != nil { - return 0, "", err + <-validated + if validationErr != nil { + return 0, "", validationErr } - // Warn only after staging finishes so the warning doesn't interleave with the - // staging spinner. + // Warn only after staging so the warning doesn't interleave with the spinner. if validationUnavailable && cmdio.HasIO(ctx) { cmdio.LogString(ctx, "Warning: server-side config validation was unavailable; continuing with submission.") } @@ -430,8 +434,6 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run return 0, "", err } - deps, _ := cfg.inlineDependencies() - runtimeVersion, _ := cfg.runtimeVersion() payload := buildSubmitPayload(cfg, commandPath, dlRuntimeImage(ctx, runtimeVersion), usagePolicyID, snap, deps) payload.IdempotencyToken = token diff --git a/cmd/air/validateconfig_test.go b/cmd/air/validateconfig_test.go index dba9d60f420..32878241503 100644 --- a/cmd/air/validateconfig_test.go +++ b/cmd/air/validateconfig_test.go @@ -111,9 +111,8 @@ func TestPreflightValidateFailsOpenWhenBackendUnavailable(t *testing.T) { for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { srv := validateServer(t, tt.status, tt.body, nil) - unavailable, err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(t.Context(), validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") assert.NoError(t, err) - assert.True(t, unavailable) }) } } @@ -146,9 +145,8 @@ func TestPreflightValidationDoesNotRetryTransportFailure(t *testing.T) { return nil, errors.New("connection failed") }) - unavailable, err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) - assert.True(t, unavailable) assert.Equal(t, int32(1), requests.Load()) } @@ -178,9 +176,8 @@ func TestPreflightValidationFailsOpenOnInvalidSuccessResponse(t *testing.T) { }, nil }) - unavailable, err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(t.Context(), w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) - assert.True(t, unavailable) assert.Equal(t, int32(1), requests.Load()) }) } @@ -240,10 +237,9 @@ func TestPreflightValidationOAuthRefreshUsesDeadline(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) defer cancel() started := time.Now() - unavailable, err := preflightValidate(ctx, w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err = preflightValidate(ctx, w, baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") require.NoError(t, err) - assert.True(t, unavailable) assert.Less(t, time.Since(started), time.Second) select { case <-refreshStarted: @@ -307,10 +303,9 @@ func TestPreflightValidationTimeoutFailsOpen(t *testing.T) { ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond) defer cancel() started := time.Now() - unavailable, err := preflightValidate(ctx, validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") + _, err := preflightValidate(ctx, validationTestWorkspaceClient(t, srv.URL), baseRunConfig(), "/Workspace/Users/me/cmd.sh", nil, "token") close(releaseRequest) require.NoError(t, err) - assert.True(t, unavailable) assert.Less(t, time.Since(started), time.Second) } From 464f79d0da0073efc7e5da70c99a7b01e7be9e07 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Thu, 5 Nov 2026 04:44:49 +0000 Subject: [PATCH 4/7] Cancel artifact staging when preflight validation fails Validation and staging now share a cancellable context: a validation failure cancels in-flight uploads and is reported instead of the cancellation, and a staging failure cancels in-flight validation. Co-authored-by: Isaac --- cmd/air/runsubmit.go | 24 +++++++++++++++--------- cmd/air/runsubmit_test.go | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 49 insertions(+), 9 deletions(-) diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index 3e77203937a..3375b058e3c 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -346,19 +346,25 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run // Validate server-side while the run artifacts stage below. The backend check // does not read the uploaded files, so the two can overlap; submission waits - // for the result, so an invalid config may leave staged artifacts but is never - // submitted. - validationCtx, cancelValidation := context.WithCancel(ctx) + // for the result, so an invalid config is never submitted. Both share + // stagingCtx, so a failure on either side cancels the other. + stagingCtx, cancelStaging := context.WithCancelCause(ctx) validated := make(chan struct{}) var validationUnavailable bool var validationErr error go func() { defer close(validated) - validationUnavailable, validationErr = preflightValidate(validationCtx, w, cfg, commandPath, containers, token) + validationUnavailable, validationErr = preflightValidate(stagingCtx, w, cfg, commandPath, containers, token) + if validationErr != nil { + cancelStaging(validationErr) + } }() - // Stop and wait for validation if staging fails first. defer func() { - cancelValidation() + // A failed validation cancels staging; report it, not the cancellation. + if cause := context.Cause(stagingCtx); err != nil && cause != nil { + err = cause + } + cancelStaging(nil) <-validated }() @@ -371,7 +377,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run usagePolicyID = strings.TrimSpace(*cfg.UsagePolicyID) } if cfg.UsagePolicyName != nil { - usagePolicyID, err = resolveUsagePolicyIDByName(ctx, w, *cfg.UsagePolicyName) + usagePolicyID, err = resolveUsagePolicyIDByName(stagingCtx, w, *cfg.UsagePolicyName) if err != nil { return 0, "", err } @@ -383,7 +389,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run if cfg.MLflowExperimentDirectory != nil { experimentDir = *cfg.MLflowExperimentDirectory } - if err := ensureExperimentDirectory(ctx, w, experimentDir); err != nil { + if err := ensureExperimentDirectory(stagingCtx, w, experimentDir); err != nil { return 0, "", err } @@ -412,7 +418,7 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run err = withSpinner(ctx, showProgress, "Staging run artifacts…", func() error { var stageErr error - snap, stageErr = stageRunArtifacts(ctx, fc, items, stageSnapshot) + snap, stageErr = stageRunArtifacts(stagingCtx, fc, items, stageSnapshot) return stageErr }) if err != nil { diff --git a/cmd/air/runsubmit_test.go b/cmd/air/runsubmit_test.go index c3703e6a321..6ca9e5292d1 100644 --- a/cmd/air/runsubmit_test.go +++ b/cmd/air/runsubmit_test.go @@ -505,6 +505,40 @@ func TestSubmitWorkloadValidatesWhileStaging(t *testing.T) { assert.Equal(t, int64(777), runID) } +// A validation failure must cancel in-flight staging and surface the validation +// error, not the cancellation. The upload handler stalls, so without the +// cancellation submitWorkload would block until the stall ends. +func TestSubmitWorkloadValidationFailureCancelsStaging(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + release := make(chan struct{}) + t.Cleanup(func() { close(release) }) + var submissions atomic.Int32 + server.Handle("POST", "/api/2.0/workspace-files/import-file/{path...}", func(req testserver.Request) any { + <-release + return req.Workspace.WorkspaceFilesImportFile(req.Vars["path"], req.Body, true) + }) + server.Handle("POST", validateConfigPath, func(req testserver.Request) any { + return validateConfigResponse{Errors: []configFieldError{{Path: "task.experiment", Message: "is invalid"}}} + }) + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + submissions.Add(1) + return jobs.SubmitRunResponse{RunId: 777} + }) + testserver.AddDefaultHandlers(server) + + w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + cfgPath := writeConfigFile(t, "run.yaml", minimalConfig) + cfg, err := loadRunConfig(cfgPath) + require.NoError(t, err) + + _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) + require.ErrorContains(t, err, "task.experiment: is invalid") + assert.Zero(t, submissions.Load()) +} + func TestSubmitWorkloadStagingErrorPreventsSubmit(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) From b4918869afd332b94b09b76e716aa1346fbcb6e5 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Thu, 5 Nov 2026 05:00:56 +0000 Subject: [PATCH 5/7] Overlap preflight validation with the artifact upload only Start validation right before the upload and run the usage policy lookup and experiment directory creation after it, so validation field errors keep precedence and a rejected config never creates the experiment directory, as before. Co-authored-by: Isaac --- cmd/air/runsubmit.go | 94 +++++++++++++++++++-------------------- cmd/air/runsubmit_test.go | 41 ++++++++++++----- 2 files changed, 76 insertions(+), 59 deletions(-) diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index 3375b058e3c..8756edca541 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -344,55 +344,8 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run containers := submittedContainers(cfg, funcDir) token := effectiveIdempotencyToken(idempotencyKey, cfg) - // Validate server-side while the run artifacts stage below. The backend check - // does not read the uploaded files, so the two can overlap; submission waits - // for the result, so an invalid config is never submitted. Both share - // stagingCtx, so a failure on either side cancels the other. - stagingCtx, cancelStaging := context.WithCancelCause(ctx) - validated := make(chan struct{}) - var validationUnavailable bool - var validationErr error - go func() { - defer close(validated) - validationUnavailable, validationErr = preflightValidate(stagingCtx, w, cfg, commandPath, containers, token) - if validationErr != nil { - cancelStaging(validationErr) - } - }() - defer func() { - // A failed validation cancels staging; report it, not the cancellation. - if cause := context.Cause(stagingCtx); err != nil && cause != nil { - err = cause - } - cancelStaging(nil) - <-validated - }() - - // Resolve the usage policy to its id next, so a bad name fails fast with a - // clear (caller-fixable) message before we upload any artifacts. Validation - // guarantees name and id are mutually exclusive: a literal id is used as-is, a - // name is resolved against the workspace. - usagePolicyID := "" - if cfg.UsagePolicyID != nil { - usagePolicyID = strings.TrimSpace(*cfg.UsagePolicyID) - } - if cfg.UsagePolicyName != nil { - usagePolicyID, err = resolveUsagePolicyIDByName(stagingCtx, w, *cfg.UsagePolicyName) - if err != nil { - return 0, "", err - } - } - deps, _ := cfg.inlineDependencies() - experimentDir := "" - if cfg.MLflowExperimentDirectory != nil { - experimentDir = *cfg.MLflowExperimentDirectory - } - if err := ensureExperimentDirectory(stagingCtx, w, experimentDir); err != nil { - return 0, "", err - } - fc, err := filer.NewWorkspaceFilesClient(w, funcDir) if err != nil { return 0, "", err @@ -416,6 +369,30 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run } } + // Validate server-side while the run artifacts upload. The backend check does + // not read the uploaded files, so the two can overlap; submission waits for + // the result, so an invalid config is never submitted. Both share stagingCtx, + // so a failure on either side cancels the other. + stagingCtx, cancelStaging := context.WithCancelCause(ctx) + validated := make(chan struct{}) + var validationUnavailable bool + var validationErr error + go func() { + defer close(validated) + validationUnavailable, validationErr = preflightValidate(stagingCtx, w, cfg, commandPath, containers, token) + if validationErr != nil { + cancelStaging(validationErr) + } + }() + defer func() { + // A failed validation cancels staging; report it, not the cancellation. + if cause := context.Cause(stagingCtx); err != nil && cause != nil { + err = cause + } + cancelStaging(nil) + <-validated + }() + err = withSpinner(ctx, showProgress, "Staging run artifacts…", func() error { var stageErr error snap, stageErr = stageRunArtifacts(stagingCtx, fc, items, stageSnapshot) @@ -440,6 +417,29 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run return 0, "", err } + // Resolve the usage policy to its id next, so a bad name fails with a clear + // (caller-fixable) message before submission. Validation guarantees name and + // id are mutually exclusive: a literal id is used as-is, a name is resolved + // against the workspace. + usagePolicyID := "" + if cfg.UsagePolicyID != nil { + usagePolicyID = strings.TrimSpace(*cfg.UsagePolicyID) + } + if cfg.UsagePolicyName != nil { + usagePolicyID, err = resolveUsagePolicyIDByName(ctx, w, *cfg.UsagePolicyName) + if err != nil { + return 0, "", err + } + } + + experimentDir := "" + if cfg.MLflowExperimentDirectory != nil { + experimentDir = *cfg.MLflowExperimentDirectory + } + if err := ensureExperimentDirectory(ctx, w, experimentDir); err != nil { + return 0, "", err + } + runtimeVersion, _ := cfg.runtimeVersion() payload := buildSubmitPayload(cfg, commandPath, dlRuntimeImage(ctx, runtimeVersion), usagePolicyID, snap, deps) payload.IdempotencyToken = token diff --git a/cmd/air/runsubmit_test.go b/cmd/air/runsubmit_test.go index 6ca9e5292d1..02fe0fe0003 100644 --- a/cmd/air/runsubmit_test.go +++ b/cmd/air/runsubmit_test.go @@ -1051,20 +1051,17 @@ func TestSubmitWorkloadGuards(t *testing.T) { base, err := loadRunConfig(cfgPath) require.NoError(t, err) - t.Run("unresolvable usage_policy_name fails before upload", func(t *testing.T) { - // An empty policy list makes the name unresolvable. Record every path the - // server sees so the "fails before any upload" ordering is asserted, not just - // asserted-by-comment: no import/mkdirs request may be made. + t.Run("unresolvable usage_policy_name blocks submission", func(t *testing.T) { + // An empty policy list makes the name unresolvable. server := testserver.New(t) t.Cleanup(server.Close) - var paths []string + var submissions atomic.Int32 server.Handle("GET", "/api/2.0/serverless-policies", func(req testserver.Request) any { - paths = append(paths, req.URL.Path) return usagePoliciesResponse{} }) - server.Handle("POST", "/api/2.0/workspace/{path...}", func(req testserver.Request) any { - paths = append(paths, req.URL.Path) - return testserver.Response{StatusCode: 200} + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + submissions.Add(1) + return jobs.SubmitRunResponse{RunId: 777} }) stubValidateConfig(server) testserver.AddDefaultHandlers(server) @@ -1075,9 +1072,29 @@ func TestSubmitWorkloadGuards(t *testing.T) { cfg.UsagePolicyName = new("nope") _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false) require.ErrorContains(t, err, `no usage policy named "nope"`) - for _, p := range paths { - assert.NotContains(t, p, "/workspace/", "no workspace write may precede policy resolution") - } + assert.Zero(t, submissions.Load()) + }) + + t.Run("validation field errors take precedence over usage policy errors", func(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + var policyLookups atomic.Int32 + server.Handle("GET", "/api/2.0/serverless-policies", func(req testserver.Request) any { + policyLookups.Add(1) + return usagePoliciesResponse{} + }) + server.Handle("POST", validateConfigPath, func(req testserver.Request) any { + return validateConfigResponse{Errors: []configFieldError{{Path: "task.experiment", Message: "is invalid"}}} + }) + testserver.AddDefaultHandlers(server) + pw, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"}) + require.NoError(t, err) + + cfg := *base + cfg.UsagePolicyName = new("nope") + _, _, err = submitWorkload(t.Context(), pw, &cfg, cfgPath, "", false) + require.ErrorContains(t, err, "task.experiment: is invalid") + assert.Zero(t, policyLookups.Load()) }) } From d0c7079f4082c10e70fdcd90f8de2fde52a27600 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Thu, 5 Nov 2026 05:57:45 +0000 Subject: [PATCH 6/7] Use errgroup to coordinate preflight validation and the upload errgroup already returns the first error and cancels the other side, so the manual goroutine, channel, and deferred error override are not needed. Co-authored-by: Isaac --- cmd/air/runsubmit.go | 45 ++++++++++++++------------------------------ 1 file changed, 14 insertions(+), 31 deletions(-) diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index 8756edca541..00c0ec87054 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -371,41 +371,24 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run // Validate server-side while the run artifacts upload. The backend check does // not read the uploaded files, so the two can overlap; submission waits for - // the result, so an invalid config is never submitted. Both share stagingCtx, - // so a failure on either side cancels the other. - stagingCtx, cancelStaging := context.WithCancelCause(ctx) - validated := make(chan struct{}) + // both, so an invalid config is never submitted. The first failure cancels + // the other side and is the one reported. + group, stagingCtx := errgroup.WithContext(ctx) var validationUnavailable bool - var validationErr error - go func() { - defer close(validated) - validationUnavailable, validationErr = preflightValidate(stagingCtx, w, cfg, commandPath, containers, token) - if validationErr != nil { - cancelStaging(validationErr) - } - }() - defer func() { - // A failed validation cancels staging; report it, not the cancellation. - if cause := context.Cause(stagingCtx); err != nil && cause != nil { - err = cause - } - cancelStaging(nil) - <-validated - }() - - err = withSpinner(ctx, showProgress, "Staging run artifacts…", func() error { - var stageErr error - snap, stageErr = stageRunArtifacts(stagingCtx, fc, items, stageSnapshot) - return stageErr + group.Go(func() (err error) { + validationUnavailable, err = preflightValidate(stagingCtx, w, cfg, commandPath, containers, token) + return err }) - if err != nil { + group.Go(func() error { + return withSpinner(ctx, showProgress, "Staging run artifacts…", func() error { + var stageErr error + snap, stageErr = stageRunArtifacts(stagingCtx, fc, items, stageSnapshot) + return stageErr + }) + }) + if err := group.Wait(); err != nil { return 0, "", err } - - <-validated - if validationErr != nil { - return 0, "", validationErr - } // Warn only after staging so the warning doesn't interleave with the spinner. if validationUnavailable && cmdio.HasIO(ctx) { cmdio.LogString(ctx, "Warning: server-side config validation was unavailable; continuing with submission.") From c6b050cdb0cda67337131b3312f82914cb5316d0 Mon Sep 17 00:00:00 2001 From: Steven Chen Date: Fri, 6 Nov 2026 17:45:43 +0000 Subject: [PATCH 7/7] Stop plain-tar snapshot packaging when the submission is cancelled A failed config validation now cancels artifact staging, but plain-tar packaging ignored its context and archived the whole tree before the validation error surfaced. Fail tar writes once the context is done. Also bound the cancellation test so a regression fails instead of hanging, and update the submitWorkload doc comment for the new order. Co-authored-by: Isaac --- cmd/air/runsubmit.go | 5 +++-- cmd/air/runsubmit_test.go | 9 ++++++--- cmd/air/snapshot_package.go | 18 +++++++++++++++++- cmd/air/snapshot_package_test.go | 13 +++++++++++++ 4 files changed, 39 insertions(+), 6 deletions(-) diff --git a/cmd/air/runsubmit.go b/cmd/air/runsubmit.go index 00c0ec87054..f01198d2b7a 100644 --- a/cmd/air/runsubmit.go +++ b/cmd/air/runsubmit.go @@ -327,8 +327,9 @@ func stageRunArtifacts(ctx context.Context, launchWriter fileWriter, items []upl return snap, err } -// submitWorkload runs the submit happy path: ensure the experiment directory, -// upload the launch artifacts, assemble the Jobs payload, and submit it. It +// submitWorkload runs the submit happy path: validate the config while uploading +// the launch artifacts, ensure the experiment directory, assemble the Jobs +// payload, and submit it. It // returns the new run_id and its dashboard URL. showProgress enables the stderr // staging spinner (text mode only). func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string, showProgress bool) (runID int64, dashboardURL string, err error) { diff --git a/cmd/air/runsubmit_test.go b/cmd/air/runsubmit_test.go index ffebebbe05a..441a1c12d26 100644 --- a/cmd/air/runsubmit_test.go +++ b/cmd/air/runsubmit_test.go @@ -505,8 +505,8 @@ func TestSubmitWorkloadValidatesWhileStaging(t *testing.T) { } // A validation failure must cancel in-flight staging and surface the validation -// error, not the cancellation. The upload handler stalls, so without the -// cancellation submitWorkload would block until the stall ends. +// error, not the cancellation. The upload handler stalls until cleanup, so +// without the cancellation submitWorkload would hit the test deadline instead. func TestSubmitWorkloadValidationFailureCancelsStaging(t *testing.T) { server := testserver.New(t) t.Cleanup(server.Close) @@ -533,8 +533,11 @@ func TestSubmitWorkloadValidationFailureCancelsStaging(t *testing.T) { cfg, err := loadRunConfig(cfgPath) require.NoError(t, err) - _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key", false) + ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second) + defer cancel() + _, _, err = submitWorkload(ctx, w, cfg, cfgPath, "idem-key", false) require.ErrorContains(t, err, "task.experiment: is invalid") + require.NoError(t, ctx.Err(), "staging was not cancelled by the validation failure") assert.Zero(t, submissions.Load()) } diff --git a/cmd/air/snapshot_package.go b/cmd/air/snapshot_package.go index e4027083335..ca8a6942f5f 100644 --- a/cmd/air/snapshot_package.go +++ b/cmd/air/snapshot_package.go @@ -5,6 +5,7 @@ import ( "context" "errors" "fmt" + "io" "os" "os/exec" "path/filepath" @@ -79,7 +80,7 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, fil out.Close() return err } - if err := tarpack.Write(gz, entries); err != nil { + if err := tarpack.Write(ctxWriter{ctx, gz}, entries); err != nil { gz.Close() out.Close() return fmt.Errorf("failed to create plain tarball: %w", err) @@ -91,6 +92,21 @@ func createPlainTarball(ctx context.Context, repoPath, outputTarball string, fil return out.Close() } +// ctxWriter fails writes once ctx is done, so a cancelled submission (for +// example, one whose config validation failed) stops packaging promptly instead +// of archiving the whole tree. +type ctxWriter struct { + ctx context.Context + w io.Writer +} + +func (w ctxWriter) Write(p []byte) (int, error) { + if err := w.ctx.Err(); err != nil { + return 0, err + } + return w.w.Write(p) +} + // snapshotFile is a file selected for the snapshot: its repo-relative path (native // separators) plus the metadata used to content-address the plain_tar upload. type snapshotFile struct { diff --git a/cmd/air/snapshot_package_test.go b/cmd/air/snapshot_package_test.go index 24249021a4f..db42f634356 100644 --- a/cmd/air/snapshot_package_test.go +++ b/cmd/air/snapshot_package_test.go @@ -3,6 +3,7 @@ package aircmd import ( "archive/tar" "compress/gzip" + "context" "os" "path/filepath" "slices" @@ -146,6 +147,18 @@ func TestCreatePlainTarball(t *testing.T) { } } +func TestCreatePlainTarballStopsWhenCancelled(t *testing.T) { + repo := t.TempDir() + writeRepoFile(t, repo, "a.txt", "1") + files, err := snapshotFiles(t.Context(), repo, nil, false) + require.NoError(t, err) + + ctx, cancel := context.WithCancel(t.Context()) + cancel() + err = createPlainTarball(ctx, repo, filepath.Join(t.TempDir(), "snap.tar.gz"), files) + require.ErrorIs(t, err, context.Canceled) +} + func TestCreatePlainTarball_HonorsGitignore(t *testing.T) { ctx := t.Context() repo := t.TempDir()