From 3fa0ddb40931b44aafff168e2c67bc9959c6c34c Mon Sep 17 00:00:00 2001 From: riddhibhagwat-db Date: Fri, 7 Aug 2026 16:58:51 +0000 Subject: [PATCH] air: pre-flight the config against ValidateConfig before submit `air run` now checks the config server-side before uploading anything, so a bad config fails fast with the backend's field-level errors instead of after the code snapshot is packaged and uploaded. The same rules back the submit gate, so the pre-flight can't disagree with what submission enforces. Fails open: the endpoint is behind a SAFE flag and older workspaces don't 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. --dry-run stays local-only and needs no workspace; the pre-flight is on the submit path, where it saves the wasted upload. --- .../air/run-submit-deps/test.toml | 4 + .../experimental/air/run-submit/test.toml | 4 + experimental/air/cmd/runsubmit.go | 6 + experimental/air/cmd/runsubmit_test.go | 19 +++ experimental/air/cmd/validateconfig.go | 140 ++++++++++++++++++ experimental/air/cmd/validateconfig_test.go | 111 ++++++++++++++ 6 files changed, 284 insertions(+) create mode 100644 experimental/air/cmd/validateconfig.go create mode 100644 experimental/air/cmd/validateconfig_test.go diff --git a/acceptance/experimental/air/run-submit-deps/test.toml b/acceptance/experimental/air/run-submit-deps/test.toml index 590c2918a25..097eb39d8c9 100644 --- a/acceptance/experimental/air/run-submit-deps/test.toml +++ b/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 = ''' diff --git a/acceptance/experimental/air/run-submit/test.toml b/acceptance/experimental/air/run-submit/test.toml index 3dc9ff81b99..2e641379092 100644 --- a/acceptance/experimental/air/run-submit/test.toml +++ b/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 = ''' diff --git a/experimental/air/cmd/runsubmit.go b/experimental/air/cmd/runsubmit.go index 3aa2436b827..389422edadc 100644 --- a/experimental/air/cmd/runsubmit.go +++ b/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) diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index 76ab3aeb3ac..e9c12fe7a7c 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/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) diff --git a/experimental/air/cmd/validateconfig.go b/experimental/air/cmd/validateconfig.go new file mode 100644 index 00000000000..4cc8709d98d --- /dev/null +++ b/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() +} diff --git a/experimental/air/cmd/validateconfig_test.go b/experimental/air/cmd/validateconfig_test.go new file mode 100644 index 00000000000..5fddcf3df10 --- /dev/null +++ b/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 }