[AIR Validation to Backend] Pre-flight config against ValidateConfig before submit

acceptance/experimental/air/run-submit-deps/test.toml

@@ -7,6 +7,10 @@ RecordRequests = true
 Pattern = "HEAD /"
 Response.Body = ''
 
+[[Server]]
+Pattern = "POST /api/2.0/ai-training/config:validate"
+Response.Body = '{}'
+
 [[Server]]
 Pattern = "POST /api/2.2/jobs/runs/submit"
 Response.Body = '''

acceptance/experimental/air/run-submit/test.toml

@@ -11,6 +11,10 @@ Ignore = ["run.yaml"]
 Pattern = "HEAD /"
 Response.Body = ''
 
+[[Server]]
+Pattern = "POST /api/2.0/ai-training/config:validate"
+Response.Body = '{}'
+
 [[Server]]
 Pattern = "POST /api/2.2/jobs/runs/submit"
 Response.Body = '''

experimental/air/cmd/runsubmit.go

@@ -145,6 +145,12 @@ func submitToken(flag string, cfg *runConfig) (string, error) {
 // upload the launch artifacts, assemble the Jobs payload, and submit it. It
 // returns the new run_id and its dashboard URL.
 func submitWorkload(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig, configPath, idempotencyKey string) (int64, string, error) {
+	// Pre-flight the config server-side before touching the workspace, so a bad
+	// config fails with the backend's field-level errors and no orphaned uploads.
+	if err := preflightValidate(ctx, w, cfg); err != nil {
+		return 0, "", err
+	}
+
 	// Resolve the idempotency token first so a bad key fails before any upload,
 	// and before the policy lookup below spends a round trip on it.
 	token, err := submitToken(idempotencyKey, cfg)

experimental/air/cmd/runsubmit_test.go

@@ -18,6 +18,14 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
+// stubValidateConfig registers an OK ValidateConfig response so submitWorkload's
+// pre-flight passes. Register before AddDefaultHandlers (the router is first-wins).
+func stubValidateConfig(server *testserver.Server) {
+	server.Handle("POST", "/api/2.0/ai-training/config:validate", func(req testserver.Request) any {
+		return validateConfigResponse{}
+	})
+}
+
 func TestDlRuntimeImage(t *testing.T) {
 	ctx := t.Context()
 	// A config runtime version wins and is used bare.
@@ -216,6 +224,7 @@ func TestSubmitWorkload(t *testing.T) {
 		require.NoError(t, json.Unmarshal(req.Body, &got))
 		return jobs.SubmitRunResponse{RunId: 777}
 	})
+	stubValidateConfig(server)
 	testserver.AddDefaultHandlers(server)
 
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
@@ -260,6 +269,7 @@ func TestSubmitWorkloadHonorsOverride(t *testing.T) {
 		require.NoError(t, json.Unmarshal(req.Body, &got))
 		return jobs.SubmitRunResponse{RunId: 777}
 	})
+	stubValidateConfig(server)
 	testserver.AddDefaultHandlers(server)
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 	require.NoError(t, err)
@@ -290,6 +300,7 @@ func TestSubmitWorkloadWithCodeSource(t *testing.T) {
 		require.NoError(t, json.Unmarshal(req.Body, &got))
 		return jobs.SubmitRunResponse{RunId: 555}
 	})
+	stubValidateConfig(server)
 	testserver.AddDefaultHandlers(server)
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 	require.NoError(t, err)
@@ -331,6 +342,7 @@ func TestSubmitWorkloadWithGitPinnedCodeSource(t *testing.T) {
 		require.NoError(t, json.Unmarshal(req.Body, &got))
 		return jobs.SubmitRunResponse{RunId: 555}
 	})
+	stubValidateConfig(server)
 	testserver.AddDefaultHandlers(server)
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 	require.NoError(t, err)
@@ -380,6 +392,7 @@ func TestSubmitWorkloadPlainTarNameIsUnique(t *testing.T) {
 	server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {
 		return jobs.SubmitRunResponse{RunId: 555}
 	})
+	stubValidateConfig(server)
 	testserver.AddDefaultHandlers(server)
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 	require.NoError(t, err)
@@ -431,6 +444,7 @@ func TestSubmitWorkloadGitArchiveCaching(t *testing.T) {
 		}
 		return req.Workspace.WorkspaceFilesImportFile(p, req.Body, req.URL.Query().Get("overwrite") == "true")
 	})
+	stubValidateConfig(server)
 	testserver.AddDefaultHandlers(server)
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 	require.NoError(t, err)
@@ -474,6 +488,7 @@ func TestSubmitWorkloadUploadsGitSidecars(t *testing.T) {
 	server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any {
 		return jobs.SubmitRunResponse{RunId: 555}
 	})
+	stubValidateConfig(server)
 	testserver.AddDefaultHandlers(server)
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 	require.NoError(t, err)
@@ -531,6 +546,7 @@ func TestSubmitWorkloadWithRemoteVolumeCodeSource(t *testing.T) {
 	server.Handle("PUT", "/api/2.0/fs/files/Volumes/{path...}", func(req testserver.Request) any {
 		return testserver.Response{StatusCode: 204}
 	})
+	stubValidateConfig(server)
 	testserver.AddDefaultHandlers(server)
 	w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 	require.NoError(t, err)
@@ -578,6 +594,7 @@ func TestSubmitWorkloadGuards(t *testing.T) {
 			paths = append(paths, req.URL.Path)
 			return testserver.Response{StatusCode: 200}
 		})
+		stubValidateConfig(server)
 		testserver.AddDefaultHandlers(server)
 		pw, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 		require.NoError(t, err)
@@ -599,6 +616,7 @@ func TestSubmitWorkloadGuards(t *testing.T) {
 			uploaded = true
 			return nil
 		})
+		stubValidateConfig(server)
 		testserver.AddDefaultHandlers(server)
 		tw, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 		require.NoError(t, err)
@@ -627,6 +645,7 @@ func TestSubmitWorkloadSendsUsagePolicy(t *testing.T) {
 		server.Handle("GET", "/api/2.0/serverless-policies", func(req testserver.Request) any {
 			return usagePoliciesResponse{Policies: []usagePolicy{{PolicyID: policyID, PolicyName: "team-a"}}}
 		})
+		stubValidateConfig(server)
 		testserver.AddDefaultHandlers(server)
 		w, err := databricks.NewWorkspaceClient(&databricks.Config{Host: server.URL, Token: "token"})
 		require.NoError(t, err)

experimental/air/cmd/validateconfig.go

@@ -0,0 +1,140 @@
+package aircmd
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+	"strings"
+
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/apierr"
+	"github.com/databricks/databricks-sdk-go/client"
+)
+
+// validateConfigPath is AiTrainingService's pre-flight: it checks a training
+// config server-side and returns the problems, without submitting. Called with a
+// raw client.Do because the SDK does not model AiTrainingService.
+const validateConfigPath = "/api/2.0/ai-training/config:validate"
+
+// configFieldError is one problem the server found, addressed to the config
+// field that caused it. Mirrors the proto FieldError.
+type configFieldError struct {
+	Path    string `json:"path"`
+	Message string `json:"message"`
+	Code    string `json:"code"`
+}
+
+type validateConfigResponse struct {
+	Errors []configFieldError `json:"errors"`
+}
+
+// preflightValidate checks the config against the backend before any upload, so
+// a bad config fails fast with the server's own field-level errors.
+//
+// It fails open: the endpoint is behind a SAFE flag and older workspaces do not
+// have it, so a disabled or missing endpoint skips the check and lets submission
+// proceed (where the config is validated again, authoritatively). Only a
+// populated error list — a config the server actively rejected — blocks.
+func preflightValidate(ctx context.Context, w *databricks.WorkspaceClient, cfg *runConfig) error {
+	apiClient, err := client.New(w.Config)
+	if err != nil {
+		return fmt.Errorf("failed to create API client: %w", err)
+	}
+
+	var resp validateConfigResponse
+	err = apiClient.Do(ctx, http.MethodPost, validateConfigPath, nil, nil, validateConfigRequest(cfg), &resp)
+	if err != nil {
+		if endpointUnavailable(err) {
+			return nil
+		}
+		return fmt.Errorf("failed to validate config: %w", err)
+	}
+	if len(resp.Errors) == 0 {
+		return nil
+	}
+	return errors.New(formatConfigErrors(resp.Errors))
+}
+
+// validateConfigRequest builds the {task, run_options} body from the user's
+// config. It carries what the user wrote — the fields set only at submit time
+// (command_path, code_source_path) are absent, which the server treats as
+// optional; the pre-flight's job is the config-level rules. Absent optional
+// fields are omitted so the server doesn't validate values the user never set.
+func validateConfigRequest(cfg *runConfig) map[string]any {
+	compute := map[string]any{}
+	if cfg.Compute != nil {
+		compute["accelerator_type"] = cfg.Compute.AcceleratorType
+		compute["accelerator_count"] = cfg.Compute.NumAccelerators
+	}
+	task := map[string]any{
+		"experiment":  cfg.ExperimentName,
+		"deployments": []any{map[string]any{"compute": compute}},
+	}
+	putOpt(task, "mlflow_run", cfg.MLflowRunName)
+	putOpt(task, "mlflow_experiment_directory", cfg.MLflowExperimentDirectory)
+	if len(cfg.Parameters) > 0 {
+		task["parameters"] = cfg.Parameters
+	}
+
+	req := map[string]any{"task": task}
+	if runOptions := validateConfigRunOptions(cfg); len(runOptions) > 0 {
+		req["run_options"] = runOptions
+	}
+	return req
+}
+
+// validateConfigRunOptions gathers the run-level fields into run_options,
+// omitting any the user didn't set.
+func validateConfigRunOptions(cfg *runConfig) map[string]any {
+	runOptions := map[string]any{}
+	putOpt(runOptions, "max_retries", cfg.MaxRetries)
+	putOpt(runOptions, "timeout_minutes", cfg.TimeoutMinutes)
+	putOpt(runOptions, "idempotency_token", cfg.IdempotencyToken)
+	putOpt(runOptions, "usage_policy_name", cfg.UsagePolicyName)
+	putOpt(runOptions, "usage_policy_id", cfg.UsagePolicyID)
+	if len(cfg.EnvVariables) > 0 {
+		runOptions["env_variables"] = cfg.EnvVariables
+	}
+	if len(cfg.Secrets) > 0 {
+		runOptions["secrets"] = cfg.Secrets
+	}
+	return runOptions
+}
+
+// putOpt sets key to the pointer's value only when it is non-nil, so an unset
+// config field is left out of the request rather than sent as a zero value.
+func putOpt[T any](m map[string]any, key string, value *T) {
+	if value != nil {
+		m[key] = *value
+	}
+}
+
+// endpointUnavailable reports whether the failure means the endpoint isn't there
+// to answer — the flag is off, or the workspace predates it — as opposed to the
+// config being rejected. Those cases fail open.
+func endpointUnavailable(err error) bool {
+	var apiErr *apierr.APIError
+	if errors.As(err, &apiErr) {
+		return apiErr.ErrorCode == "FEATURE_DISABLED" ||
+			apiErr.StatusCode == http.StatusNotFound ||
+			apiErr.StatusCode == http.StatusNotImplemented
+	}
+	return false
+}
+
+// formatConfigErrors renders the field errors as one message, one problem per
+// line, each pointing at the config field the user wrote.
+func formatConfigErrors(fieldErrors []configFieldError) string {
+	var b strings.Builder
+	b.WriteString("config validation failed:")
+	for _, e := range fieldErrors {
+		b.WriteString("\n  ")
+		if e.Path != "" {
+			b.WriteString(e.Path)
+			b.WriteString(": ")
+		}
+		b.WriteString(e.Message)
+	}
+	return b.String()
+}

experimental/air/cmd/validateconfig_test.go

@@ -0,0 +1,111 @@
+package aircmd
+
+import (
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func baseRunConfig() *runConfig {
+	return &runConfig{
+		ExperimentName: "llama-fine-tune",
+		Compute:        &computeConfig{NumAccelerators: 16, AcceleratorType: "GPU_8xH100"},
+	}
+}
+
+// validateServer serves one ValidateConfig response with the given status and
+// body, and records the request body it received.
+func validateServer(t *testing.T, status int, body string, gotReq *map[string]any) *httptest.Server {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if r.URL.Path == validateConfigPath {
+			if gotReq != nil {
+				_ = json.NewDecoder(r.Body).Decode(gotReq)
+			}
+			w.WriteHeader(status)
+			_, _ = w.Write([]byte(body))
+			return
+		}
+		_, _ = w.Write([]byte(`{}`))
+	}))
+	t.Cleanup(srv.Close)
+	return srv
+}
+
+func TestPreflightValidatePasses(t *testing.T) {
+	srv := validateServer(t, http.StatusOK, `{}`, nil)
+	err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig())
+	assert.NoError(t, err)
+}
+
+func TestPreflightValidateReportsErrors(t *testing.T) {
+	body := `{"errors":[
+		{"path":"experiment","message":"only letters, digits, hyphens, underscores","code":"DISALLOWED_CHARACTERS"},
+		{"path":"deployments[0].compute.accelerator_count","message":"must be a multiple of 8","code":"COUNT_NOT_MULTIPLE"}
+	]}`
+	srv := validateServer(t, http.StatusOK, body, nil)
+	err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig())
+	require.Error(t, err)
+	// Every problem is surfaced, each pointing at its config field.
+	assert.Contains(t, err.Error(), "experiment: only letters")
+	assert.Contains(t, err.Error(), "deployments[0].compute.accelerator_count: must be a multiple of 8")
+}
+
+func TestPreflightValidateFailsOpenWhenDisabled(t *testing.T) {
+	// The endpoint is behind a SAFE flag; a disabled endpoint must not block the run.
+	srv := validateServer(t, http.StatusBadRequest,
+		`{"error_code":"FEATURE_DISABLED","message":"ValidateConfig is not yet enabled."}`, nil)
+	err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig())
+	assert.NoError(t, err)
+}
+
+func TestPreflightValidateFailsOpenWhenNotFound(t *testing.T) {
+	// A workspace that predates the endpoint returns 404; skip and let submit proceed.
+	srv := validateServer(t, http.StatusNotFound, `{"error_code":"ENDPOINT_NOT_FOUND","message":"not found"}`, nil)
+	err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig())
+	assert.NoError(t, err)
+}
+
+func TestValidateConfigRequestShape(t *testing.T) {
+	var gotReq map[string]any
+	srv := validateServer(t, http.StatusOK, `{}`, &gotReq)
+
+	cfg := baseRunConfig()
+	cfg.MaxRetries = intPtr(3)
+	cfg.EnvVariables = map[string]string{"HF_HOME": "/tmp/hf"}
+	err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), cfg)
+	require.NoError(t, err)
+
+	task := gotReq["task"].(map[string]any)
+	assert.Equal(t, "llama-fine-tune", task["experiment"])
+	deployment := task["deployments"].([]any)[0].(map[string]any)
+	compute := deployment["compute"].(map[string]any)
+	assert.Equal(t, "GPU_8xH100", compute["accelerator_type"])
+	assert.EqualValues(t, 16, compute["accelerator_count"])
+
+	runOptions := gotReq["run_options"].(map[string]any)
+	assert.EqualValues(t, 3, runOptions["max_retries"])
+	assert.Equal(t, map[string]any{"HF_HOME": "/tmp/hf"}, runOptions["env_variables"])
+}
+
+func TestValidateConfigRequestOmitsUnsetOptions(t *testing.T) {
+	// A minimal config carries no run_options and only the fields it set, so the
+	// server never validates values the user didn't provide.
+	var gotReq map[string]any
+	srv := validateServer(t, http.StatusOK, `{}`, &gotReq)
+
+	err := preflightValidate(t.Context(), newTestWorkspaceClient(t, srv.URL), baseRunConfig())
+	require.NoError(t, err)
+
+	_, hasRunOptions := gotReq["run_options"]
+	assert.False(t, hasRunOptions)
+	task := gotReq["task"].(map[string]any)
+	_, hasMlflowRun := task["mlflow_run"]
+	assert.False(t, hasMlflowRun)
+}
+
+func intPtr(v int) *int { return &v }