diff --git a/acceptance/experimental/air/run/output.txt b/acceptance/experimental/air/run/output.txt index a753eabd19..e733bca091 100644 --- a/acceptance/experimental/air/run/output.txt +++ b/acceptance/experimental/air/run/output.txt @@ -14,9 +14,22 @@ Dry run: configuration for "smoke-test" is valid; not submitting. } } -=== override not yet supported ->>> [CLI] experimental air run -f valid.yaml --dry-run --override a=b -Error: --override is not yet supported +=== override applies and logs the change +>>> [CLI] experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=2 --override timeout_minutes=45 +Override: changing compute.num_accelerators from 1 to 2 +Override: setting timeout_minutes to 45 +Dry run: configuration for "smoke-test" is valid; not submitting. + +=== override of an unknown field is rejected +>>> [CLI] experimental air run -f valid.yaml --dry-run --override bogus=1 +Error: invalid --override "bogus": "bogus" is not a known field; available fields are: code_source, command, compute, env_variables, environment, experiment_name, idempotency_token, max_retries, mlflow_experiment_directory, mlflow_run_name, parameters, permissions, secrets, timeout_minutes, usage_policy_id, usage_policy_name + +Exit code: 1 + +=== override still runs schema validation +>>> [CLI] experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0 +Override: changing compute.num_accelerators from 1 to 0 +Error: compute.num_accelerators must be positive, got 0 Exit code: 1 diff --git a/acceptance/experimental/air/run/script b/acceptance/experimental/air/run/script index 312b2f6fec..3308b5bc0f 100644 --- a/acceptance/experimental/air/run/script +++ b/acceptance/experimental/air/run/script @@ -4,8 +4,14 @@ trace $CLI experimental air run -f valid.yaml --dry-run title "dry-run (json)" trace $CLI experimental air run -f valid.yaml --dry-run -o json -title "override not yet supported" -errcode trace $CLI experimental air run -f valid.yaml --dry-run --override a=b +title "override applies and logs the change" +trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=2 --override timeout_minutes=45 + +title "override of an unknown field is rejected" +errcode trace $CLI experimental air run -f valid.yaml --dry-run --override bogus=1 + +title "override still runs schema validation" +errcode trace $CLI experimental air run -f valid.yaml --dry-run --override compute.num_accelerators=0 title "watch not yet supported" errcode trace $CLI experimental air run -f valid.yaml --dry-run --watch diff --git a/experimental/air/cmd/run.go b/experimental/air/cmd/run.go index bd32810e9b..327c4de167 100644 --- a/experimental/air/cmd/run.go +++ b/experimental/air/cmd/run.go @@ -57,16 +57,12 @@ The workload is described by a YAML config file (see --file).`, cmd.RunE = func(cmd *cobra.Command, args []string) error { ctx := cmd.Context() - // These flags' pipelines are not ported yet; reject rather than silently - // ignore them. - if len(overrides) > 0 { - return errors.New("--override is not yet supported") - } + // --watch's pipeline is not ported yet; reject rather than silently ignore it. if watch { return errors.New("--watch is not yet supported") } - cfg, err := loadRunConfig(file) + cfg, err := loadRunConfigWithOverrides(ctx, file, overrides) if err != nil { return err } diff --git a/experimental/air/cmd/runconfig_load.go b/experimental/air/cmd/runconfig_load.go index 81b07d3ca5..7f6ad8b5e1 100644 --- a/experimental/air/cmd/runconfig_load.go +++ b/experimental/air/cmd/runconfig_load.go @@ -1,6 +1,8 @@ package aircmd import ( + "bytes" + "context" "errors" "fmt" "io" @@ -10,10 +12,11 @@ import ( ) // decodeRunConfig reads and decodes the run YAML into the schema. Unknown keys -// are rejected (KnownFields), mirroring the Python schema's extra="forbid". +// are rejected (KnownFields). // -// The `_bases_` composition feature and CLI `--override` handling are not yet -// ported; a config using `_bases_` is currently rejected as an unknown field. +// The `_bases_` composition feature is not yet ported; a config using `_bases_` +// is currently rejected as an unknown field. CLI `--override` handling lives in +// runconfig_override.go and is applied to the parsed map before this decode. func decodeRunConfig(path string) (*runConfig, error) { f, err := os.Open(path) if err != nil { @@ -21,7 +24,13 @@ func decodeRunConfig(path string) (*runConfig, error) { } defer f.Close() - dec := yaml.NewDecoder(f) + return decodeRunConfigReader(f, path) +} + +// decodeRunConfigReader decodes and unknown-key-checks a run YAML from r. path is +// used only for error messages. +func decodeRunConfigReader(r io.Reader, path string) (*runConfig, error) { + dec := yaml.NewDecoder(r) dec.KnownFields(true) var cfg runConfig @@ -50,3 +59,53 @@ func loadRunConfig(path string) (*runConfig, error) { } return cfg, nil } + +// loadRunConfigWithOverrides decodes a run YAML config, applies any +// --override KEY=VALUE entries to the parsed map, then re-decodes (with unknown +// keys rejected) and structurally validates the result. Applying overrides to +// the map — rather than the typed config — lets the single decode+validate +// pipeline enforce path existence, type coercion, and the semantic rules at +// once. ctx is used only to log applied overrides. +func loadRunConfigWithOverrides(ctx context.Context, path string, overrides []string) (*runConfig, error) { + if len(overrides) == 0 { + return loadRunConfig(path) + } + + entries, err := parseOverrides(overrides) + if err != nil { + return nil, err + } + if err := validateOverridePaths(entries); err != nil { + return nil, err + } + + raw, err := os.ReadFile(path) + if err != nil { + return nil, err + } + var m map[string]any + if err := yaml.Unmarshal(raw, &m); err != nil { + return nil, fmt.Errorf("invalid config %s: %w", path, err) + } + if m == nil { + // An empty file decodes to a nil map; start from an empty one so overrides + // can populate it (the re-decode still enforces required fields). + m = map[string]any{} + } + if err := applyOverrides(ctx, m, entries); err != nil { + return nil, err + } + + merged, err := yaml.Marshal(m) + if err != nil { + return nil, err + } + cfg, err := decodeRunConfigReader(bytes.NewReader(merged), path) + if err != nil { + return nil, err + } + if err := validateRunConfig(cfg); err != nil { + return nil, err + } + return cfg, nil +} diff --git a/experimental/air/cmd/runconfig_override.go b/experimental/air/cmd/runconfig_override.go new file mode 100644 index 0000000000..a00a178603 --- /dev/null +++ b/experimental/air/cmd/runconfig_override.go @@ -0,0 +1,162 @@ +package aircmd + +import ( + "context" + "fmt" + "maps" + "reflect" + "slices" + "strings" + + "github.com/databricks/cli/libs/cmdio" + "go.yaml.in/yaml/v3" +) + +// This file implements the `--override KEY=VALUE` flag. Overrides are applied to +// the parsed YAML map (not the typed runConfig) before re-decode, so one pipeline +// covers path existence, type coercion, and the semantic validate() rules. + +// freeFormFields hold free-form maps, so path validation stops at them: any +// sub-path is valid. +var freeFormFields = map[string]bool{ + "parameters": true, + "env_variables": true, + "secrets": true, +} + +// parseOverrides parses --override KEY=VALUE arguments, preserving order. +func parseOverrides(overrides []string) ([]overrideEntry, error) { + entries := make([]overrideEntry, 0, len(overrides)) + for _, item := range overrides { + key, value, found := strings.Cut(item, "=") + if !found { + // --override is repeatable, so a config path meant for -f can be + // swallowed here; point at the real fix. + hint := "" + if strings.HasSuffix(item, ".yaml") || strings.HasSuffix(item, ".yml") { + hint = fmt.Sprintf("; %q looks like a config file — pass it with -f/--file", item) + } + return nil, fmt.Errorf("invalid --override %q: expected KEY=VALUE (e.g. compute.num_accelerators=32)%s", item, hint) + } + key = strings.TrimSpace(key) + if key == "" { + return nil, fmt.Errorf("invalid --override %q: empty key", item) + } + entries = append(entries, overrideEntry{path: key, raw: value}) + } + return entries, nil +} + +// overrideEntry is one parsed --override: its dotted path and the raw RHS string. +type overrideEntry struct { + path string + raw string +} + +// validateOverridePaths checks every dotted path against the runConfig schema +// before mutation, so an error names the exact --override key rather than the +// re-decode's Go-type language. +func validateOverridePaths(entries []overrideEntry) error { + for _, e := range entries { + if err := checkOverridePath(strings.Split(e.path, "."), reflect.TypeFor[runConfig](), e.path); err != nil { + return err + } + } + return nil +} + +// checkOverridePath recursively validates one dotted path against a struct type +// whose fields carry `yaml:` tags. +func checkOverridePath(parts []string, t reflect.Type, fullPath string) error { + field := parts[0] + fields := yamlFields(t) + sub, ok := fields[field] + if !ok { + return fmt.Errorf("invalid --override %q: %q is not a known field; available fields are: %s", + fullPath, field, strings.Join(slices.Sorted(maps.Keys(fields)), ", ")) + } + if len(parts) == 1 { + return nil + } + if freeFormFields[field] { + return nil + } + subStruct := underlyingStruct(sub) + if subStruct == nil { + return fmt.Errorf("invalid --override %q: %q is not a nested object; cannot address sub-field %q", + fullPath, field, strings.Join(parts[1:], ".")) + } + return checkOverridePath(parts[1:], subStruct, fullPath) +} + +// yamlFields maps a struct's yaml tag names to their field types, skipping +// fields without a yaml tag (or tagged "-"). +func yamlFields(t reflect.Type) map[string]reflect.Type { + out := map[string]reflect.Type{} + for f := range t.Fields() { + tag := f.Tag.Get("yaml") + if tag == "" || tag == "-" { + continue + } + name, _, _ := strings.Cut(tag, ",") + if name == "" || name == "-" { + continue + } + out[name] = f.Type + } + return out +} + +// underlyingStruct unwraps pointer/slice indirection and returns the struct type +// a field decodes into, or nil if the field is not a struct (a scalar/map/etc.). +func underlyingStruct(t reflect.Type) reflect.Type { + for t.Kind() == reflect.Pointer || t.Kind() == reflect.Slice { + t = t.Elem() + } + if t.Kind() == reflect.Struct { + return t + } + return nil +} + +// applyOverrides walks each dotted path into the parsed YAML map and sets the +// leaf to the RHS parsed as a YAML scalar. Intermediate maps are auto-created so +// an override can add a field the YAML omits; the later re-decode rejects paths +// absent from the schema. Changes are logged to stderr to keep JSON stdout clean. +func applyOverrides(ctx context.Context, m map[string]any, entries []overrideEntry) error { + for _, e := range entries { + var value any + if err := yaml.Unmarshal([]byte(e.raw), &value); err != nil { + return fmt.Errorf("invalid --override %q: cannot parse value %q: %w", e.path, e.raw, err) + } + + parts := strings.Split(e.path, ".") + current := m + for _, part := range parts[:len(parts)-1] { + next, ok := current[part].(map[string]any) + if !ok { + next = map[string]any{} + current[part] = next + } + current = next + } + + leaf := parts[len(parts)-1] + old, had := current[leaf] + current[leaf] = value + if had { + logOverride(ctx, fmt.Sprintf("Override: changing %s from %v to %v", e.path, old, value)) + } else { + logOverride(ctx, fmt.Sprintf("Override: setting %s to %v", e.path, value)) + } + } + return nil +} + +// logOverride writes to stderr only when a cmdIO is present; cmdio.LogString +// panics without one, as in non-command callers such as unit tests. +func logOverride(ctx context.Context, msg string) { + if cmdio.HasIO(ctx) { + cmdio.LogString(ctx, msg) + } +} diff --git a/experimental/air/cmd/runconfig_override_test.go b/experimental/air/cmd/runconfig_override_test.go new file mode 100644 index 0000000000..99eab7dd88 --- /dev/null +++ b/experimental/air/cmd/runconfig_override_test.go @@ -0,0 +1,168 @@ +package aircmd + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestParseOverrides(t *testing.T) { + tests := []struct { + name string + in []string + want []overrideEntry + wantErr string + }{ + { + name: "key=value pairs preserve order", + in: []string{"compute.num_accelerators=8", "timeout_minutes=45"}, + want: []overrideEntry{ + {path: "compute.num_accelerators", raw: "8"}, + {path: "timeout_minutes", raw: "45"}, + }, + }, + { + name: "value may contain =", + in: []string{"env_variables.EXPR=a=b"}, + want: []overrideEntry{{path: "env_variables.EXPR", raw: "a=b"}}, + }, + { + name: "key is trimmed", + in: []string{" timeout_minutes = 45"}, + want: []overrideEntry{{path: "timeout_minutes", raw: " 45"}}, + }, + { + name: "missing = is rejected", + in: []string{"compute.num_accelerators"}, + wantErr: `expected KEY=VALUE`, + }, + { + name: "a .yaml token hints at -f", + in: []string{"train.yaml"}, + wantErr: `looks like a config file`, + }, + { + name: "empty key is rejected", + in: []string{"=5"}, + wantErr: `empty key`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + got, err := parseOverrides(tt.in) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + assert.Equal(t, tt.want, got) + }) + } +} + +func TestValidateOverridePaths(t *testing.T) { + tests := []struct { + name string + path string + wantErr string + }{ + {name: "known top-level field", path: "experiment_name"}, + {name: "known nested field", path: "compute.num_accelerators"}, + {name: "free-form sub-path", path: "env_variables.MY_VAR"}, + {name: "deep free-form sub-path", path: "parameters.model.layers"}, + { + name: "unknown top-level field", + path: "bogus", + wantErr: `"bogus" is not a known field`, + }, + { + name: "unknown nested field", + path: "compute.bogus", + wantErr: `"bogus" is not a known field`, + }, + { + name: "sub-field of a scalar", + path: "command.sub", + wantErr: `"command" is not a nested object`, + }, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := validateOverridePaths([]overrideEntry{{path: tt.path, raw: "x"}}) + if tt.wantErr != "" { + require.ErrorContains(t, err, tt.wantErr) + return + } + require.NoError(t, err) + }) + } +} + +// overrideBaseConfig is a valid 8-GPU config the override tests mutate. +const overrideBaseConfig = `experiment_name: smoke +command: python train.py +compute: + accelerator_type: GPU_8xH100 + num_accelerators: 8 +env_variables: + EXISTING: hello +` + +func TestLoadRunConfigWithOverrides(t *testing.T) { + t.Run("no overrides matches loadRunConfig", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), nil) + require.NoError(t, err) + assert.Equal(t, 8, cfg.Compute.NumAccelerators) + }) + + t.Run("typed scalar override is coerced", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=16"}) + require.NoError(t, err) + assert.Equal(t, 16, cfg.Compute.NumAccelerators) + }) + + t.Run("multiple overrides all apply", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=16", "timeout_minutes=45"}) + require.NoError(t, err) + assert.Equal(t, 16, cfg.Compute.NumAccelerators) + require.NotNil(t, cfg.TimeoutMinutes) + assert.Equal(t, 45, *cfg.TimeoutMinutes) + }) + + t.Run("free-form env var adds a key as a string", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"env_variables.RANK=0"}) + require.NoError(t, err) + // A numeric-looking value stays a string because env_variables is map[string]string. + assert.Equal(t, "0", cfg.EnvVariables["RANK"]) + }) + + t.Run("intermediate maps are auto-created", func(t *testing.T) { + cfg, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"environment.docker_image.url=my/img:1"}) + require.NoError(t, err) + require.NotNil(t, cfg.Environment) + require.NotNil(t, cfg.Environment.DockerImage) + assert.Equal(t, "my/img:1", cfg.Environment.DockerImage.URL) + }) + + t.Run("unknown path errors before mutation", func(t *testing.T) { + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"bogus=1"}) + require.ErrorContains(t, err, `"bogus" is not a known field`) + }) + + t.Run("semantic validation runs after override", func(t *testing.T) { + // 3 is a known field with a valid type, so only validate() can reject it. + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=3"}) + require.ErrorContains(t, err, "must be a multiple of 8") + }) + + t.Run("type mismatch is rejected on re-decode", func(t *testing.T) { + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators=abc"}) + require.Error(t, err) + }) + + t.Run("malformed override is rejected", func(t *testing.T) { + _, err := loadRunConfigWithOverrides(t.Context(), writeConfig(t, overrideBaseConfig), []string{"compute.num_accelerators"}) + require.ErrorContains(t, err, "expected KEY=VALUE") + }) +} diff --git a/experimental/air/cmd/runsubmit_test.go b/experimental/air/cmd/runsubmit_test.go index fd5103599d..dcd9dd8eec 100644 --- a/experimental/air/cmd/runsubmit_test.go +++ b/experimental/air/cmd/runsubmit_test.go @@ -156,6 +156,39 @@ func TestSubmitWorkload(t *testing.T) { assert.Equal(t, jobs.ComputeSpec{AcceleratorType: jobs.ComputeSpecAcceleratorTypeGpu1xH100, AcceleratorCount: 1}, d.Compute) } +// TestSubmitWorkloadHonorsOverride proves a --override reaches the actual +// runs/submit payload on a real submit, not just dry-run validation: the config +// pins num_accelerators=1, the override bumps it to 4, and the recorded request +// body must carry 4. +func TestSubmitWorkloadHonorsOverride(t *testing.T) { + server := testserver.New(t) + t.Cleanup(server.Close) + + // Register before AddDefaultHandlers: the router is first-wins, so this must + // claim the route ahead of the default jobs/runs/submit handler. + var got jobs.SubmitRun + server.Handle("POST", "/api/2.2/jobs/runs/submit", func(req testserver.Request) any { + require.NoError(t, json.Unmarshal(req.Body, &got)) + 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 := loadRunConfigWithOverrides(t.Context(), cfgPath, []string{"compute.num_accelerators=4"}) + require.NoError(t, err) + + _, _, err = submitWorkload(t.Context(), w, cfg, cfgPath, "idem-key") + require.NoError(t, err) + + require.Len(t, got.Tasks, 1) + at := got.Tasks[0].AiRuntimeTask + require.NotNil(t, at) + require.Len(t, at.Deployments, 1) + assert.Equal(t, 4, at.Deployments[0].Compute.AcceleratorCount) +} + // TestSubmitWorkloadWithCodeSource exercises the snapshot path end to end: a // git-pinned code_source is packaged, uploaded, and its paths attached to the task. func TestSubmitWorkloadWithCodeSource(t *testing.T) {