Run AIR preflight validation in parallel with artifact staging

.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. ([#7368](https://github.com/databricks/cli/pull/7368))

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) {
@@ -344,42 +345,8 @@ func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *run
 	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
-	// 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
-		}
-	}
-
 	deps, _ := cfg.inlineDependencies()
 
-	experimentDir := ""
-	if cfg.MLflowExperimentDirectory != nil {
-		experimentDir = *cfg.MLflowExperimentDirectory
-	}
-	if err := ensureExperimentDirectory(ctx, w, experimentDir); err != nil {
-		return 0, "", err
-	}
-
 	fc, err := filer.NewWorkspaceFilesClient(w, funcDir)
 	if err != nil {
 		return 0, "", err
@@ -403,12 +370,57 @@ 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)
-		return stageErr
+	// 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
+	// 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
+	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
+	}
+	// 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.")
+	}
+
+	// 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
+	}
+
+	// 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
 	}
 

cmd/air/runsubmit_test.go

@@ -408,6 +408,7 @@ func TestSubmitWorkloadValidationOutcome(t *testing.T) {
 		response   any
 		wantError  string
 		wantSubmit bool
+		wantStderr string
 	}{
 		{
 			name: "field error",
@@ -420,6 +421,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",
 		},
 	}
 
@@ -445,23 +447,100 @@ 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)
+}
+
+// A validation failure must cancel in-flight staging and surface the validation
+// 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)
+
+	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)
+
+	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())
+}
+
 func TestSubmitWorkloadStagingErrorPreventsSubmit(t *testing.T) {
 	server := testserver.New(t)
 	t.Cleanup(server.Close)
@@ -974,20 +1053,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)
@@ -998,9 +1074,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())
 	})
 }
 

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 {

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