diff --git a/.nextchanges/air/accelerator-field-compute-help.md b/.nextchanges/air/accelerator-field-compute-help.md new file mode 100644 index 00000000000..bdc6b13b9e0 --- /dev/null +++ b/.nextchanges/air/accelerator-field-compute-help.md @@ -0,0 +1 @@ +* Show workspace-supported accelerator types in `databricks air run -h config.compute.accelerator_type`. ([#7275](https://github.com/databricks/cli/pull/7275)) diff --git a/acceptance/air/config-help/output.txt b/acceptance/air/config-help/output.txt index 16fedb9ffa9..4df01151206 100644 --- a/acceptance/air/config-help/output.txt +++ b/acceptance/air/config-help/output.txt @@ -75,7 +75,7 @@ config.compute Which accelerators to run on and how many. Supported accelerator types: - GPU_1xH100 - 1 accelerator per node [Private Preview] + GPU_1xH100 - 1 accelerator per node GPU_8xH100 - 8 accelerators per node GPU_1xRTX6000 - 1 accelerator per node @@ -128,6 +128,11 @@ config.environment.version config.compute.accelerator_type Which accelerator to run on, e.g. GPU_1xA10. See https://docs.databricks.com/aws/en/machine-learning/ai-runtime/cli/yaml-config#reference for the current list of supported GPU types. Matched case-sensitively. + Supported accelerator types: + GPU_1xH100 - 1 accelerator per node + GPU_8xH100 - 8 accelerators per node + GPU_1xRTX6000 - 1 accelerator per node + Type: string Required: no diff --git a/acceptance/air/config-help/test.toml b/acceptance/air/config-help/test.toml index 79cb5772cdf..6759bc32089 100644 --- a/acceptance/air/config-help/test.toml +++ b/acceptance/air/config-help/test.toml @@ -9,7 +9,7 @@ Pattern = "GET /api/2.0/ai-training/compute-options" Response.Body = ''' { "compute_options": [ - {"hardware_accelerator": "GPU_1xH100", "per_node_accelerator_count": 1, "launch_stage": "PRIVATE_PREVIEW"}, + {"hardware_accelerator": "GPU_1xH100", "per_node_accelerator_count": 1, "launch_stage": "GA"}, {"hardware_accelerator": "GPU_8xH100", "per_node_accelerator_count": 8, "launch_stage": "GA"}, {"hardware_accelerator": "GPU_1xRTX6000", "per_node_accelerator_count": 1, "launch_stage": "GA"} ] diff --git a/cmd/air/run.go b/cmd/air/run.go index d63193e5888..332f1aa1218 100644 --- a/cmd/air/run.go +++ b/cmd/air/run.go @@ -122,7 +122,7 @@ The path must be a separate argument: cobra reserves -h as a boolean, so } var computeOptions []computeOption - if field.path == "config.compute" { + if field.path == "config.compute" || field.path == "config.compute.accelerator_type" { computeOptions = loadComputeOptionsForHelp(c, args) } renderConfigField(c.OutOrStdout(), field, computeOptions) diff --git a/cmd/air/runconfig.go b/cmd/air/runconfig.go index 31c701cdc12..e4b8fb5411e 100644 --- a/cmd/air/runconfig.go +++ b/cmd/air/runconfig.go @@ -862,7 +862,7 @@ func renderConfigField(w io.Writer, f configField, computeOptions []computeOptio if f.help != "" { fmt.Fprintf(w, " %s\n", f.help) } - if f.path == "config.compute" && computeOptions != nil { + if (f.path == "config.compute" || f.path == "config.compute.accelerator_type") && computeOptions != nil { fmt.Fprintln(w, "\n Supported accelerator types:") if len(computeOptions) == 0 { fmt.Fprintln(w, " No accelerator types are currently available for AI Runtime in this workspace.") diff --git a/cmd/air/runconfig_test.go b/cmd/air/runconfig_test.go index e3c2f34520a..e2d4b47c934 100644 --- a/cmd/air/runconfig_test.go +++ b/cmd/air/runconfig_test.go @@ -984,12 +984,19 @@ func TestRunCommandHelp_NoConfigPath(t *testing.T) { // -h with a config path documents the field instead of the command, and does so // without the otherwise-required --file. func TestRunCommandHelp_ConfigPath(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _, _ = w.Write([]byte(`{"compute_options":[{"hardware_accelerator":"GPU_FUTURE","per_node_accelerator_count":16}]}`)) + })) + t.Cleanup(srv.Close) + var out strings.Builder cmd := newRunCommand() + cmd.SetContext(cmdctx.SetWorkspaceClient(t.Context(), newTestWorkspaceClient(t, srv.URL))) cmd.SetOut(&out) cmd.SetArgs([]string{"-h", "config.compute.accelerator_type"}) require.NoError(t, cmd.Execute()) assert.Contains(t, out.String(), "config.compute.accelerator_type") + assert.Contains(t, out.String(), "GPU_FUTURE - 16 accelerators per node") assert.Contains(t, out.String(), "Type: string") assert.NotContains(t, out.String(), "Usage:") }