diff --git a/NOTICE b/NOTICE index 922dd1fd79..dd1e4c0c1b 100644 --- a/NOTICE +++ b/NOTICE @@ -155,6 +155,10 @@ mattn/go-isatty - https://github.com/mattn/go-isatty Copyright (c) Yasuhiro MATSUMOTO License - https://github.com/mattn/go-isatty/blob/master/LICENSE +muesli/termenv - https://github.com/muesli/termenv +Copyright (c) 2019 Christian Muehlhaeuser +License - https://github.com/muesli/termenv/blob/master/LICENSE + sabhiram/go-gitignore - https://github.com/sabhiram/go-gitignore Copyright (c) 2015 Shaba Abhiram License - https://github.com/sabhiram/go-gitignore/blob/master/LICENSE diff --git a/acceptance/experimental/air/get/output.txt b/acceptance/experimental/air/get/output.txt index f8df1fb36b..b9866c4cf9 100644 --- a/acceptance/experimental/air/get/output.txt +++ b/acceptance/experimental/air/get/output.txt @@ -1,17 +1,42 @@ === get (text) >>> [CLI] experimental air get run 123 -Job Link: [DATABRICKS_URL]/jobs/runs/123?o=[NUMID] - -Run ID: 123 -Status: SUCCESS -Submitted: 2023-11-14 22:13 UTC -Retries: 0 -Duration: 12s -Experiment: my-exp -MLflow Run: my-run -User: user@example.com -Accelerators: 8x H100 + +╭─ Configuration ────────────────────────────────────────────────╮ +│ │ +│ experiment_name: my-exp │ +│ compute: │ +│ accelerator_type: a10 │ +│ num_accelerators: 1 │ +│ command: | │ +│ for i in $(seq 1 10); do │ +│ echo "step $i" │ +│ done │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +╭─ Training Progress ────────────────────────────────────────────╮ +│ │ +│ ██████████████████████████████████ 100% 10/10 steps · 12s │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + +╭─ Metadata ─────────────────────────────────────────────────────╮ +│ │ +│ Run ID 123 │ +│ Status ● SUCCESS │ +│ Submitted 2023-11-14 22:13 UTC │ +│ Retries 0 │ +│ Max Retries 3 │ +│ Duration 12s │ +│ Experiment my-exp │ +│ MLflow Run my-run │ +│ User user@example.com │ +│ Accelerators 1x A10 │ +│ Environment ml-runtime-gpu:1.0 │ +│ │ +╰────────────────────────────────────────────────────────────────╯ + === get (json) >>> [CLI] experimental air get run 123 -o json diff --git a/acceptance/experimental/air/get/test.toml b/acceptance/experimental/air/get/test.toml index 6162cebdb5..9429166a5e 100644 --- a/acceptance/experimental/air/get/test.toml +++ b/acceptance/experimental/air/get/test.toml @@ -23,9 +23,12 @@ Response.Body = ''' { "task_key": "train", "attempt_number": 0, + "max_retries": 3, "gen_ai_compute_task": { "mlflow_experiment_name": "/Users/user@example.com/my-exp", - "compute": {"gpu_type": "GPU_8xH100", "num_gpus": 8} + "compute": {"gpu_type": "GPU_1xA10", "num_gpus": 1}, + "dl_runtime_image": "ml-runtime-gpu:1.0", + "yaml_parameters": "experiment_name: my-exp\ncompute:\n accelerator_type: a10\n num_accelerators: 1\ncommand: |\n for i in $(seq 1 10); do\n echo \"step $i\"\n done\n" } } ] @@ -39,9 +42,10 @@ Response.Body = ''' {"gen_ai_compute_output": {"run_info": {"mlflow_experiment_id": "exp1", "mlflow_run_id": "run1"}}} ''' -# The MLflow Run cell shows the run's name, fetched from the MLflow REST API. +# The MLflow Run cell shows the run's name; the progress bar reads the highest +# logged step (metrics) against the max_steps param. [[Server]] Pattern = "GET /api/2.0/mlflow/runs/get" Response.Body = ''' -{"run": {"info": {"run_name": "my-run"}}} +{"run": {"info": {"run_name": "my-run"}, "data": {"metrics": [{"step": 10}], "params": [{"key": "max_steps", "value": "10"}]}}} ''' diff --git a/experimental/air/cmd/format.go b/experimental/air/cmd/format.go index 8694ed69e9..57f86b80f0 100644 --- a/experimental/air/cmd/format.go +++ b/experimental/air/cmd/format.go @@ -6,6 +6,7 @@ import ( "fmt" "io" "math" + "strconv" "strings" "time" @@ -263,3 +264,29 @@ func gpuDisplayName(gpuType string) string { } return gpuType } + +// environment returns the run's runtime image (the training environment), or an +// empty string if the run has no GenAI-compute task. +func environment(run *jobs.Run) string { + if len(run.Tasks) == 0 { + return "" + } + task := run.Tasks[0].GenAiComputeTask + if task == nil { + return "" + } + return task.DlRuntimeImage +} + +// maxRetries returns the configured retry limit for the run's latest task as a +// display string: "unlimited" for the backend's -1, otherwise the count. +func maxRetries(run *jobs.Run) string { + if len(run.Tasks) == 0 { + return "0" + } + n := run.Tasks[len(run.Tasks)-1].MaxRetries + if n < 0 { + return "unlimited" + } + return strconv.Itoa(n) +} diff --git a/experimental/air/cmd/get.go b/experimental/air/cmd/get.go index a302b7caf3..98581fd5dd 100644 --- a/experimental/air/cmd/get.go +++ b/experimental/air/cmd/get.go @@ -1,17 +1,13 @@ package aircmd import ( - "context" "errors" "fmt" - "io" "strconv" "github.com/databricks/cli/cmd/root" "github.com/databricks/cli/libs/cmdctx" "github.com/databricks/cli/libs/flags" - "github.com/databricks/cli/libs/log" - "github.com/databricks/databricks-sdk-go" "github.com/databricks/databricks-sdk-go/apierr" "github.com/databricks/databricks-sdk-go/service/jobs" "github.com/spf13/cobra" @@ -30,26 +26,27 @@ type getData struct { DashboardURL string `json:"dashboard_url"` MLflowURL *string `json:"mlflow_url"` - // The fields below are pre-rendered for the text table and excluded from - // JSON (matching `air get run --json`). The table always shows every row, - // with "N/A" for missing values, in the same order as the Python CLI. The - // Run ID, Experiment, and MLflow Run cells carry terminal hyperlinks when - // stdout is a terminal, so the URLs don't appear as bare text. - RunIDDisplay string `json:"-"` + // The fields below are pre-rendered text-view cells, excluded from JSON + // (matching `air get run --json`). Each shows "N/A" when its value is + // missing. The styled single-run renderer (render.go) consumes them; the + // Run ID, Status, and MLflow Run cells it draws are styled and hyperlinked + // there rather than stored here. SubmittedDisplay string `json:"-"` DurationDisplay string `json:"-"` ExperimentDisplay string `json:"-"` - MLflowDisplay string `json:"-"` UserDisplay string `json:"-"` AcceleratorsDisplay string `json:"-"` + EnvironmentDisplay string `json:"-"` + MaxRetriesDisplay string `json:"-"` // Sweep replaces the single-run view for foreach runs. Sweep *sweepInfo `json:"-"` } -// getTemplate is the text-mode layout. It reads from the JSON envelope, so -// every field is reached through ".Data". -const getTemplate = `{{- if .Data.Sweep -}} -Sweep Run ID: {{.Data.RunID}} +// getTemplate is the text-mode layout for a sweep (foreach) run. Single runs are +// drawn by the styled renderer in render.go and never reach this template; it is +// used only when .Data.Sweep is set. It reads from the JSON envelope, so every +// field is reached through ".Data". +const getTemplate = `Sweep Run ID: {{.Data.RunID}} Status: {{.Data.Status}} Total: {{.Data.Sweep.Total}} Completed: {{.Data.Sweep.Completed}} @@ -64,17 +61,6 @@ Sweep Tasks: {{printf " %-24s %-14s %-12s %s" .TaskKey .RunID .Status .Experiment}} {{- end}} {{- end}} -{{- else -}} -Run ID: {{.Data.RunIDDisplay}} -Status: {{.Data.Status}} -Submitted: {{.Data.SubmittedDisplay}} -Retries: {{.Data.AttemptNumber}} -Duration: {{.Data.DurationDisplay}} -Experiment: {{.Data.ExperimentDisplay}} -MLflow Run: {{.Data.MLflowDisplay}} -User: {{.Data.UserDisplay}} -Accelerators: {{.Data.AcceleratorsDisplay}} -{{- end}} ` // newGetCommand is the `get` parent group. Subcommands name the resource to @@ -146,31 +132,28 @@ func newGetRunCommand() *cobra.Command { data.Sweep = buildSweepInfo(ctx, w, task) } - if root.OutputType(cmd) == flags.OutputText { - out := cmd.OutOrStdout() - addTextLinks(ctx, out, w, &data, ids) + if root.OutputType(cmd) != flags.OutputText { + return renderEnvelope(ctx, data) + } - // Lead with the job run link (hyperlinked, falling back to the bare - // URL off a terminal), then a gap before the training config and the - // status table, mirroring the Python CLI's header. + out := cmd.OutOrStdout() + if data.Sweep != nil { + // A sweep has no single status, config, or timing, so lead with the + // job run link and render the foreach summary table (getTemplate). fmt.Fprintf(out, "Job Link: %s\n\n", hyperlink(ctx, out, data.DashboardURL, data.DashboardURL)) - - // Text mode shows the training-config YAML before the status, - // mirroring `air get run`. JSON output omits it. - if path := yamlConfigPath(run); path != "" { - printConfigYAML(ctx, out, w, path) - } + return renderEnvelope(ctx, data) } - return renderEnvelope(ctx, data) + + renderRunText(ctx, out, w, run, &data, ids) + return nil } return cmd } -// buildGetData extracts the fields we display from a run. The text-table cells -// are pre-rendered here with their "N/A" fallbacks; the Run ID, Experiment, and -// MLflow Run cells are finalized later by addTextLinks once the dashboard and -// MLflow identifiers are known. +// buildGetData extracts the fields we display from a run. The text-view cells +// are pre-rendered here with their "N/A" fallbacks; the styled renderer adds the +// hyperlinks and colors once the dashboard and MLflow identifiers are known. func buildGetData(run *jobs.Run) getData { data := getData{ RunID: strconv.FormatInt(run.RunId, 10), @@ -180,7 +163,6 @@ func buildGetData(run *jobs.Run) getData { AttemptNumber: latestAttemptNumber(run), ExperimentName: experimentName(run), } - data.RunIDDisplay = data.RunID data.SubmittedDisplay = submittedDisplay(run) data.DurationDisplay = na if data.DurationSeconds != nil { @@ -190,57 +172,9 @@ func buildGetData(run *jobs.Run) getData { if data.ExperimentName != nil { data.ExperimentDisplay = *data.ExperimentName } - data.MLflowDisplay = na data.UserDisplay = orNA(run.CreatorUserName) data.AcceleratorsDisplay = orNA(accelerators(run)) + data.EnvironmentDisplay = orNA(environment(run)) + data.MaxRetriesDisplay = maxRetries(run) return data } - -// addTextLinks adds the terminal hyperlinks shown in text mode: the Run ID links -// to the run's dashboard page (Python embeds this on the Run ID instead of a -// separate Dashboard row), and the Experiment and MLflow Run cells link to their -// MLflow pages. On a non-terminal these degrade to plain text. -func addTextLinks(ctx context.Context, out io.Writer, w *databricks.WorkspaceClient, data *getData, ids *mlflowIdentifiers) { - data.RunIDDisplay = hyperlink(ctx, out, data.RunID, data.DashboardURL) - if ids == nil { - return - } - if data.ExperimentName != nil { - data.ExperimentDisplay = hyperlink(ctx, out, *data.ExperimentName, mlflowExperimentURL(w.Config.Host, ids)) - } - data.MLflowDisplay = hyperlink(ctx, out, mlflowRunLabel(ctx, w, ids.RunID), mlflowRunURL(w.Config.Host, ids)) -} - -// yamlConfigPath returns the run's training-config YAML path, or "" if none. -func yamlConfigPath(run *jobs.Run) string { - if len(run.Tasks) == 0 { - return "" - } - task := run.Tasks[0].GenAiComputeTask - if task == nil { - return "" - } - return task.YamlParametersFilePath -} - -// printConfigYAML downloads the run's training-config YAML and writes it to out -// (stdout), mirroring the Python `air get`. It is best-effort: a download or read -// failure is surfaced as a warning on stderr but does not fail the command. -func printConfigYAML(ctx context.Context, out io.Writer, w *databricks.WorkspaceClient, path string) { - r, err := w.Workspace.Download(ctx, path) - if err != nil { - log.Warnf(ctx, "air get: could not download training config %s: %v", path, err) - return - } - defer r.Close() - - content, err := io.ReadAll(r) - if err != nil { - log.Warnf(ctx, "air get: could not read training config %s: %v", path, err) - return - } - - fmt.Fprintln(out, "Training Configuration:") - fmt.Fprintln(out, reformatYAMLForDisplay(content)) - fmt.Fprintln(out) -} diff --git a/experimental/air/cmd/get_test.go b/experimental/air/cmd/get_test.go index 643c521c0e..7f4d784807 100644 --- a/experimental/air/cmd/get_test.go +++ b/experimental/air/cmd/get_test.go @@ -3,8 +3,6 @@ package aircmd import ( "bytes" "encoding/json" - "io" - "strings" "testing" "text/template" @@ -20,8 +18,8 @@ import ( "github.com/stretchr/testify/require" ) -// renderGet renders the get template against the JSON envelope, exactly as -// the command does, so the test covers the real template branches. +// renderGet renders the get template against the JSON envelope, exactly as the +// command does for a sweep run, so the test covers the real template branches. func renderGet(t *testing.T, data getData) string { t.Helper() tmpl, err := template.New("get").Parse(getTemplate) @@ -31,20 +29,6 @@ func renderGet(t *testing.T, data getData) string { return buf.String() } -func TestGetTemplateSingleRun(t *testing.T) { - out := renderGet(t, getData{ - RunIDDisplay: "123", - Status: "RUNNING", - UserDisplay: "me@example.com", - }) - assert.Contains(t, out, "Run ID: 123") - assert.Contains(t, out, "Status: RUNNING") - assert.Contains(t, out, "User: me@example.com") - assert.NotContains(t, out, "Sweep") - // Python embeds the dashboard link on the Run ID; there is no Dashboard row. - assert.NotContains(t, out, "Dashboard") -} - func TestGetRunInvalidID(t *testing.T) { m := mocks.NewMockWorkspaceClient(t) ctx := cmdctx.SetWorkspaceClient(cmdio.MockDiscard(t.Context()), m.WorkspaceClient) @@ -88,51 +72,6 @@ func TestGetRunNotFoundJSON(t *testing.T) { assert.Equal(t, jsonError{Code: "NOT_FOUND", Kind: "NOT_FOUND", Message: "run 5 not found: check the run ID and that it is a job run ID"}, got.Error) } -func TestPrintConfigYAML(t *testing.T) { - t.Run("downloads and prints to stdout", func(t *testing.T) { - ctx := cmdio.MockDiscard(t.Context()) - m := mocks.NewMockWorkspaceClient(t) - // The mock asserts Download is called with the resolved path. - m.GetMockWorkspaceAPI().EXPECT(). - Download(mock.Anything, "/Workspace/cfg.yaml"). - Return(io.NopCloser(strings.NewReader("epochs: 3\ncommand: \"set -e\\npython train.py\\n\"\n")), nil) - - // The config is data output and must land on stdout (the out writer), - // matching the Python `air get` behavior. - var out bytes.Buffer - printConfigYAML(ctx, &out, m.WorkspaceClient, "/Workspace/cfg.yaml") - assert.Contains(t, out.String(), "Training Configuration:") - assert.Contains(t, out.String(), "epochs: 3") - // The multi-line command is reformatted to a `|` block literal. - assert.Contains(t, out.String(), "command: |\n set -e\n python train.py") - }) - - t.Run("download failure is non-fatal and writes nothing", func(t *testing.T) { - ctx := cmdio.MockDiscard(t.Context()) - m := mocks.NewMockWorkspaceClient(t) - m.GetMockWorkspaceAPI().EXPECT(). - Download(mock.Anything, "/Workspace/missing.yaml"). - Return(nil, apierr.ErrResourceDoesNotExist) - - // Must not panic and must not write to stdout: a failed config fetch is - // best-effort, surfaced only as a stderr warning. - var out bytes.Buffer - printConfigYAML(ctx, &out, m.WorkspaceClient, "/Workspace/missing.yaml") - assert.Empty(t, out.String()) - }) -} - -func TestYAMLConfigPath(t *testing.T) { - // No tasks, or a task without GenAiComputeTask, yields no path. - assert.Empty(t, yamlConfigPath(&jobs.Run{})) - assert.Empty(t, yamlConfigPath(&jobs.Run{Tasks: []jobs.RunTask{{}}})) - - run := &jobs.Run{Tasks: []jobs.RunTask{{ - GenAiComputeTask: &jobs.GenAiComputeTask{YamlParametersFilePath: "/Workspace/cfg.yaml"}, - }}} - assert.Equal(t, "/Workspace/cfg.yaml", yamlConfigPath(run)) -} - func TestGetTemplateSweep(t *testing.T) { out := renderGet(t, getData{ RunID: "456", @@ -152,8 +91,6 @@ func TestGetTemplateSweep(t *testing.T) { assert.Contains(t, out, "iter_1") assert.Contains(t, out, "FAILED") assert.Contains(t, out, "my-exp") - // The single-run rows must not appear in the sweep view. - assert.NotContains(t, out, "Dashboard:") } func TestGetTemplateSweepNoTasks(t *testing.T) { @@ -169,60 +106,6 @@ func TestGetTemplateSweepNoTasks(t *testing.T) { assert.NotContains(t, out, "Sweep Tasks:") } -func TestGetTemplateMinimal(t *testing.T) { - // Every row always renders; missing values show "N/A", in Python's order. - out := renderGet(t, getData{ - RunIDDisplay: "1", - Status: "PENDING", - SubmittedDisplay: na, - DurationDisplay: na, - ExperimentDisplay: na, - MLflowDisplay: na, - UserDisplay: na, - AcceleratorsDisplay: na, - }) - for _, want := range []string{ - "Run ID: 1", - "Status: PENDING", - "Submitted: N/A", - "Retries: 0", - "Duration: N/A", - "Experiment: N/A", - "MLflow Run: N/A", - "User: N/A", - "Accelerators: N/A", - } { - assert.Contains(t, out, want) - } -} - -func TestGetTemplateAllFields(t *testing.T) { - out := renderGet(t, getData{ - RunIDDisplay: "1", - Status: "SUCCESS", - SubmittedDisplay: "2023-11-14 22:13 UTC", - DurationDisplay: "12s", - AttemptNumber: 2, - ExperimentDisplay: "exp", - MLflowDisplay: "sunny-cat-42", - UserDisplay: "me@example.com", - AcceleratorsDisplay: "8x H100", - }) - for _, want := range []string{ - "Run ID: 1", - "Status: SUCCESS", - "Submitted: 2023-11-14 22:13 UTC", - "Retries: 2", - "Duration: 12s", - "Experiment: exp", - "MLflow Run: sunny-cat-42", - "User: me@example.com", - "Accelerators: 8x H100", - } { - assert.Contains(t, out, want) - } -} - func TestBuildGetData(t *testing.T) { run := &jobs.Run{ RunId: 123, @@ -240,7 +123,6 @@ func TestBuildGetData(t *testing.T) { } d := buildGetData(run) assert.Equal(t, "123", d.RunID) - assert.Equal(t, "123", d.RunIDDisplay) assert.Equal(t, "SUCCESS", d.Status) assert.Equal(t, 1, d.AttemptNumber) assert.Equal(t, "2023-11-14 22:13 UTC", d.SubmittedDisplay) @@ -257,11 +139,10 @@ func TestBuildGetData(t *testing.T) { func TestBuildGetDataEmpty(t *testing.T) { // A run with no tasks, creator, or timing renders every text cell as "N/A". d := buildGetData(&jobs.Run{RunId: 7}) - assert.Equal(t, "7", d.RunIDDisplay) + assert.Equal(t, "7", d.RunID) assert.Equal(t, na, d.SubmittedDisplay) assert.Equal(t, na, d.DurationDisplay) assert.Equal(t, na, d.ExperimentDisplay) - assert.Equal(t, na, d.MLflowDisplay) assert.Equal(t, na, d.UserDisplay) assert.Equal(t, na, d.AcceleratorsDisplay) } diff --git a/experimental/air/cmd/mlflow.go b/experimental/air/cmd/mlflow.go index 7e2b5b8fd5..4cdbf3e34c 100644 --- a/experimental/air/cmd/mlflow.go +++ b/experimental/air/cmd/mlflow.go @@ -4,6 +4,7 @@ import ( "context" "fmt" "net/http" + "strconv" "strings" "github.com/databricks/cli/libs/log" @@ -76,51 +77,93 @@ func mlflowLogsURL(host string, ids *mlflowIdentifiers) string { strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) } -// mlflowExperimentURL links to the MLflow experiment page; mlflowRunURL links to -// the run page. These back the Experiment and MLflow Run hyperlinks in text mode. -func mlflowExperimentURL(host string, ids *mlflowIdentifiers) string { - return fmt.Sprintf("%s/ml/experiments/%s", strings.TrimRight(host, "/"), ids.ExperimentID) -} - +// mlflowRunURL links to the MLflow run page; it backs the MLflow Run hyperlink +// in the single-run view. func mlflowRunURL(host string, ids *mlflowIdentifiers) string { return fmt.Sprintf("%s/ml/experiments/%s/runs/%s", strings.TrimRight(host, "/"), ids.ExperimentID, ids.RunID) } -// mlflowRunLabel returns the MLflow run's human-readable name to use as the -// hyperlink text, falling back to "...{last 8 of run id}" when the name can't be -// fetched. Mirrors Python's _get_mlflow_run_name (cli_display.py). -func mlflowRunLabel(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) string { - if name := fetchMLflowRunName(ctx, w, mlflowRunID); name != "" { - return name - } - if len(mlflowRunID) > 8 { - return "..." + mlflowRunID[len(mlflowRunID)-8:] - } - return "..." + mlflowRunID +// maxStepsParam is the MLflow run parameter that records a training run's target +// step count. It is the denominator of the progress bar; the numerator is the +// highest step the run has logged a metric at. +const maxStepsParam = "max_steps" + +// mlflowRunDetails holds the MLflow run fields we surface in text mode: the run's +// display name and its training-step progress. Progress is best-effort (many +// runs log no step metrics), so HasSteps records whether Done/Total mean anything. +type mlflowRunDetails struct { + Name string + Done int + Total int + HasSteps bool } -// fetchMLflowRunName fetches a run's MLflow run_name via the MLflow REST API, -// returning "" if it can't be obtained. Best-effort, like the rest of the -// MLflow enrichment. -func fetchMLflowRunName(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) string { +// fetchMLflowRun reads a run's name and step progress from the MLflow REST API, +// returning a zero value on any failure. Both consumers (the MLflow Run label and +// the progress bar) degrade gracefully when the data is missing, so this is +// best-effort like the rest of the MLflow enrichment. +// +// Done is the highest `step` across the run's latest metrics; Total is the run's +// `max_steps` parameter. HasSteps is true only when max_steps is a positive +// integer, so callers can tell a real "0 of N" from "no step data". +func fetchMLflowRun(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) mlflowRunDetails { apiClient, err := client.New(w.Config) if err != nil { - log.Debugf(ctx, "air get: could not build API client for MLflow run name: %v", err) - return "" + log.Debugf(ctx, "air get: could not build API client for MLflow run: %v", err) + return mlflowRunDetails{} } var out struct { Run struct { Info struct { RunName string `json:"run_name"` } `json:"info"` + Data struct { + Metrics []struct { + Step int64 `json:"step"` + } `json:"metrics"` + Params []struct { + Key string `json:"key"` + Value string `json:"value"` + } `json:"params"` + } `json:"data"` } `json:"run"` } err = apiClient.Do(ctx, http.MethodGet, "/api/2.0/mlflow/runs/get", nil, nil, map[string]any{"run_id": mlflowRunID}, &out) if err != nil { - log.Debugf(ctx, "air get: could not fetch MLflow run name: %v", err) - return "" + log.Debugf(ctx, "air get: could not fetch MLflow run: %v", err) + return mlflowRunDetails{} + } + + details := mlflowRunDetails{Name: out.Run.Info.RunName} + var maxStep int64 + for _, m := range out.Run.Data.Metrics { + maxStep = max(maxStep, m.Step) + } + for _, p := range out.Run.Data.Params { + if p.Key != maxStepsParam { + continue + } + if total, err := strconv.Atoi(p.Value); err == nil && total > 0 { + details.Done = int(maxStep) + details.Total = total + details.HasSteps = true + } + break + } + return details +} + +// mlflowRunLabel is the text shown for the MLflow Run cell: the run's name, or +// "...{last 8 of run id}" when the name is unknown. Mirrors Python's +// _get_mlflow_run_name (cli_display.py). +func mlflowRunLabel(name, mlflowRunID string) string { + if name != "" { + return name } - return out.Run.Info.RunName + if len(mlflowRunID) > 8 { + return "..." + mlflowRunID[len(mlflowRunID)-8:] + } + return "..." + mlflowRunID } diff --git a/experimental/air/cmd/mlflow_test.go b/experimental/air/cmd/mlflow_test.go index a6a5c15107..49a18a6ce8 100644 --- a/experimental/air/cmd/mlflow_test.go +++ b/experimental/air/cmd/mlflow_test.go @@ -83,30 +83,49 @@ func TestMLflowURLs(t *testing.T) { ids := &mlflowIdentifiers{ExperimentID: "E1", RunID: "R1"} // A trailing slash on the host must not produce a double slash in the link. assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1/artifacts/logs/node_0", mlflowLogsURL("https://h.test/", ids)) - assert.Equal(t, "https://h.test/ml/experiments/E1", mlflowExperimentURL("https://h.test", ids)) assert.Equal(t, "https://h.test/ml/experiments/E1/runs/R1", mlflowRunURL("https://h.test", ids)) } func TestMLflowRunLabel(t *testing.T) { + // Uses the run name when it is known. + assert.Equal(t, "sunny-cat-42", mlflowRunLabel("sunny-cat-42", "0123456789abcdef")) + // Falls back to the last 8 characters of a long run id. + assert.Equal(t, "...9abcdef0", mlflowRunLabel("", "0123456789abcdef0")) + // A short run id is shown in full behind the ellipsis. + assert.Equal(t, "...short", mlflowRunLabel("", "short")) +} + +func TestFetchMLflowRun(t *testing.T) { ctx := t.Context() - t.Run("uses the run name when available", func(t *testing.T) { + mlflowServer := func(body string) *httptest.Server { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/api/2.0/mlflow/runs/get" { - _, _ = w.Write([]byte(`{"run":{"info":{"run_name":"sunny-cat-42"}}}`)) + _, _ = w.Write([]byte(body)) return } _, _ = w.Write([]byte(`{}`)) })) t.Cleanup(srv.Close) - assert.Equal(t, "sunny-cat-42", mlflowRunLabel(ctx, newTestWorkspaceClient(t, srv.URL), "0123456789abcdef")) + return srv + } + + t.Run("returns name and step progress", func(t *testing.T) { + // Done is the highest step across metrics; Total is the max_steps param. + srv := mlflowServer(`{"run":{"info":{"run_name":"sunny-cat-42"},"data":{"metrics":[{"step":3},{"step":7},{"step":5}],"params":[{"key":"max_steps","value":"10"}]}}}`) + got := fetchMLflowRun(ctx, newTestWorkspaceClient(t, srv.URL), "run1") + assert.Equal(t, mlflowRunDetails{Name: "sunny-cat-42", Done: 7, Total: 10, HasSteps: true}, got) }) - t.Run("falls back to the last 8 characters of the run id", func(t *testing.T) { - srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - _, _ = w.Write([]byte(`{}`)) - })) - t.Cleanup(srv.Close) - assert.Equal(t, "...9abcdef0", mlflowRunLabel(ctx, newTestWorkspaceClient(t, srv.URL), "0123456789abcdef0")) + t.Run("no step data when max_steps is absent", func(t *testing.T) { + srv := mlflowServer(`{"run":{"info":{"run_name":"sunny-cat-42"},"data":{"metrics":[{"step":3}]}}}`) + got := fetchMLflowRun(ctx, newTestWorkspaceClient(t, srv.URL), "run1") + assert.Equal(t, mlflowRunDetails{Name: "sunny-cat-42"}, got) + }) + + t.Run("zero value when the run cannot be fetched", func(t *testing.T) { + srv := mlflowServer(`{}`) + got := fetchMLflowRun(ctx, newTestWorkspaceClient(t, srv.URL), "run1") + assert.Equal(t, mlflowRunDetails{}, got) }) } diff --git a/experimental/air/cmd/render.go b/experimental/air/cmd/render.go new file mode 100644 index 0000000000..e7670319b6 --- /dev/null +++ b/experimental/air/cmd/render.go @@ -0,0 +1,409 @@ +package aircmd + +import ( + "context" + "fmt" + "io" + "math" + "strconv" + "strings" + + "github.com/charmbracelet/lipgloss" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/cli/libs/log" + "github.com/databricks/databricks-sdk-go" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/muesli/termenv" + "go.yaml.in/yaml/v3" +) + +// progressBarWidth is the fixed glyph count of the progress bar. Everything else +// in the view is data-driven; only the bar has a layout constant. +const progressBarWidth = 34 + +// Box titles, rendered into the top border of each box. +const ( + configBoxTitle = "Configuration" + progressBoxTitle = "Training Progress" + metadataBoxTitle = "Metadata" +) + +// minBoxInnerWidth keeps all boxes a uniform, comfortable width; boxHPad and +// boxVPad are the horizontal and vertical padding inside each box. +const ( + minBoxInnerWidth = 60 + boxHPad = 2 + boxVPad = 1 +) + +// palette holds the lipgloss styles for the single-run view. Two layers: a +// neutral ramp for chrome and text, and an accent palette for syntax and data. +// All styles come from one renderer, so they honor its color profile (Ascii +// under --no-color / non-TTY, which strips every escape). +type palette struct { + n7 lipgloss.Style // dim: block indicator, empty progress, command-block "|" + n8 lipgloss.Style // muted: field labels + n12 lipgloss.Style // content: config and metadata values, percent + + border lipgloss.Style // box borders and titles + + blue lipgloss.Style // yaml keys and hyperlinks + green lipgloss.Style // success status and progress fill + amber lipgloss.Style // in-progress status + red lipgloss.Style // failed status +} + +func newPalette(r *lipgloss.Renderer) palette { + fg := func(hex string) lipgloss.Style { return r.NewStyle().Foreground(lipgloss.Color(hex)) } + return palette{ + n7: fg("#6E6E70"), + n8: fg("#8C8A86"), + n12: fg("#F9F7F4"), // Oat Light + border: fg("#B7A8E8"), // light purple (box borders and titles) + blue: fg("#8FB3DC"), + green: fg("#74C39A"), + amber: fg("#DCAA5C"), + red: fg("#D9756B"), + } +} + +// runView is the resolved, display-ready data the renderer draws. It is built +// from getData plus the MLflow enrichment, so the renderer itself does no API +// calls or formatting decisions. +type runView struct { + runID string + dashboardURL string + status string + submitted string + retries int + maxRetries string + duration string + experiment string + mlflowLabel string + mlflowURL string + user string + accelerators string + environment string +} + +// renderRunText writes the styled single-run view: a training-config box, a +// completed progress bar, and a field list, separated by blank lines. It is a +// one-shot renderer — it builds the full string and writes it once, with no +// streaming, spinner, or redraw. +func renderRunText(ctx context.Context, out io.Writer, w *databricks.WorkspaceClient, run *jobs.Run, data *getData, ids *mlflowIdentifiers) { + colorOn := cmdio.SupportsColor(ctx, out) + renderer := lipgloss.NewRenderer(out) + if !colorOn { + // Ascii emits no SGR codes; combined with the link fallback below this + // gives clean, un-escaped output under --no-color / NO_COLOR / piped stdout. + renderer.SetColorProfile(termenv.Ascii) + } + p := newPalette(renderer) + + view := runView{ + runID: data.RunID, + dashboardURL: data.DashboardURL, + status: data.Status, + submitted: data.SubmittedDisplay, + retries: data.AttemptNumber, + maxRetries: data.MaxRetriesDisplay, + duration: data.DurationDisplay, + experiment: data.ExperimentDisplay, + mlflowLabel: na, + user: data.UserDisplay, + accelerators: data.AcceleratorsDisplay, + environment: data.EnvironmentDisplay, + } + + var progress mlflowRunDetails + if ids != nil { + progress = fetchMLflowRun(ctx, w, ids.RunID) + view.mlflowLabel = mlflowRunLabel(progress.Name, ids.RunID) + view.mlflowURL = mlflowRunURL(w.Config.Host, ids) + } + + var sections []string + if task := genAIComputeTask(run); task != nil { + if body := colorizeConfig(p, configYAML(ctx, w, task)); body != "" { + sections = append(sections, renderBox(p, configBoxTitle, body)) + } + } + // A failed run made no meaningful progress, so omit the progress box. + if !isFailedStatus(view.status) { + bar := renderProgressBar(p, view.status, progress.Done, progress.Total, progress.HasSteps, view.duration) + sections = append(sections, renderBox(p, progressBoxTitle, bar)) + } + sections = append(sections, renderBox(p, metadataBoxTitle, renderFields(p, colorOn, view))) + + // A single write: a blank line before the first box and after the last, and + // one between each box. + fmt.Fprintf(out, "\n%s\n\n", strings.Join(sections, "\n\n")) +} + +// genAIComputeTask returns the run's first GenAI-compute task, or nil. +func genAIComputeTask(run *jobs.Run) *jobs.GenAiComputeTask { + if len(run.Tasks) == 0 { + return nil + } + return run.Tasks[0].GenAiComputeTask +} + +// configYAML returns the run's resolved training config as YAML for the box. The +// full config (including the command/script) lives in the run's parameters, not +// the structured task fields, so we prefer the inline parameters, then the +// parameters file, and only synthesize a minimal config as a last resort. +func configYAML(ctx context.Context, w *databricks.WorkspaceClient, task *jobs.GenAiComputeTask) string { + if task.YamlParameters != "" { + return strings.TrimRight(reformatYAMLForDisplay([]byte(task.YamlParameters)), "\n") + } + if task.YamlParametersFilePath != "" { + if raw := downloadConfig(ctx, w, task.YamlParametersFilePath); len(raw) > 0 { + return strings.TrimRight(reformatYAMLForDisplay(raw), "\n") + } + } + return synthConfigYAML(task) +} + +// downloadConfig fetches the run's training-config file, returning nil on +// failure (logged as a warning). Best-effort, like the rest of the enrichment. +func downloadConfig(ctx context.Context, w *databricks.WorkspaceClient, path string) []byte { + r, err := w.Workspace.Download(ctx, path) + if err != nil { + log.Warnf(ctx, "air get: could not download training config %s: %v", path, err) + return nil + } + defer r.Close() + content, err := io.ReadAll(r) + if err != nil { + log.Warnf(ctx, "air get: could not read training config %s: %v", path, err) + return nil + } + return content +} + +// configBox describes the synthesized config we marshal when the run exposes no +// parameters, in the order the fields are shown. +type configBox struct { + ExperimentName string `yaml:"experiment_name,omitempty"` + Compute *configCompute `yaml:"compute,omitempty"` +} + +type configCompute struct { + AcceleratorType string `yaml:"accelerator_type,omitempty"` + NumAccelerators int `yaml:"num_accelerators,omitempty"` +} + +// synthConfigYAML builds a minimal config from the structured task fields. It +// omits the command, which is only available in the run parameters. +func synthConfigYAML(task *jobs.GenAiComputeTask) string { + cfg := configBox{} + if task.MlflowExperimentName != "" { + cfg.ExperimentName = stripExperimentUserPrefix(task.MlflowExperimentName) + } + if task.Compute != nil && task.Compute.NumGpus > 0 { + cfg.Compute = &configCompute{ + AcceleratorType: task.Compute.GpuType, + NumAccelerators: task.Compute.NumGpus, + } + } + if cfg.ExperimentName == "" && cfg.Compute == nil { + return "" + } + b, err := yaml.Marshal(cfg) + if err != nil { + return "" + } + return strings.TrimRight(reformatYAMLForDisplay(b), "\n") +} + +// colorizeConfig styles a YAML config block line by line. +func colorizeConfig(p palette, body string) string { + if body == "" { + return "" + } + lines := strings.Split(body, "\n") + for i, line := range lines { + lines[i] = colorizeConfigLine(p, line) + } + return strings.Join(lines, "\n") +} + +// colorizeConfigLine colors one YAML line: keys blue, the `|` block indicator +// dim, and every value (and the command body that isn't a `key:` pair) in the +// neutral content color. +func colorizeConfigLine(p palette, line string) string { + indent := line[:len(line)-len(strings.TrimLeft(line, " "))] + trimmed := strings.TrimLeft(line, " ") + + if i := strings.IndexByte(trimmed, ':'); i > 0 && isConfigKey(trimmed[:i]) { + key := trimmed[:i] + value := strings.TrimSpace(trimmed[i+1:]) + styled := indent + p.blue.Render(key+":") + switch value { + case "": + // A mapping parent such as "compute:" has no value of its own. + case "|": + styled += " " + p.n7.Render(value) + default: + styled += " " + p.n12.Render(value) + } + return styled + } + return indent + p.n12.Render(trimmed) +} + +// isConfigKey reports whether s is a bare YAML key (lowercase, digits, and +// underscores). It guards against treating a colon inside a command body as a +// key/value separator. +func isConfigKey(s string) bool { + if s == "" { + return false + } + for _, r := range s { + if r != '_' && (r < 'a' || r > 'z') && (r < '0' || r > '9') { + return false + } + } + return true +} + +// renderBox draws a rounded-border box around body, with title rendered into the +// top border in the border color. body lines are padded to the widest one (or +// minBoxInnerWidth), with boxHPad columns and boxVPad rows of padding inside. +func renderBox(p palette, title, body string) string { + border := lipgloss.RoundedBorder() + lines := strings.Split(body, "\n") + pad := strings.Repeat(" ", boxHPad) + + titleWidth := lipgloss.Width(title) + inner := max(minBoxInnerWidth, titleWidth+2) + for _, line := range lines { + inner = max(inner, lipgloss.Width(line)) + } + + left := p.border.Render(border.Left) + right := p.border.Render(border.Right) + blank := left + strings.Repeat(" ", inner+2*boxHPad) + right + + var b strings.Builder + // Top: ╭─ ──…──╮. The dash count makes the row width match the body, + // accounting for the boxHPad columns on each side. + trailing := inner + 2*boxHPad - titleWidth - 3 + b.WriteString(p.border.Render(border.TopLeft + border.Top)) + b.WriteString(" " + p.border.Render(title) + " ") + b.WriteString(p.border.Render(strings.Repeat(border.Top, trailing) + border.TopRight)) + b.WriteByte('\n') + + for range boxVPad { + b.WriteString(blank + "\n") + } + for _, line := range lines { + fill := strings.Repeat(" ", inner-lipgloss.Width(line)) + b.WriteString(left + pad + line + fill + pad + right) + b.WriteByte('\n') + } + for range boxVPad { + b.WriteString(blank + "\n") + } + + b.WriteString(p.border.Render(border.BottomLeft + strings.Repeat(border.Bottom, inner+2*boxHPad) + border.BottomRight)) + return b.String() +} + +// renderProgressBar draws the bar, percent, and step/duration line. A successful +// run is always full; an in-progress run with step data fills floor(34*done/total). +// Without step data the bar reflects status and the steps segment is dropped. +func renderProgressBar(p palette, status string, done, total int, hasSteps bool, duration string) string { + fill := statusStyle(p, status) + + var fraction float64 + switch { + case isSuccessStatus(status): + fraction = 1 + case hasSteps && total > 0: + fraction = min(float64(done)/float64(total), 1) + } + + filled := min(int(math.Floor(progressBarWidth*fraction)), progressBarWidth) + bar := fill.Render(strings.Repeat("█", filled)) + p.n7.Render(strings.Repeat("░", progressBarWidth-filled)) + + segments := []string{bar} + if hasSteps || isSuccessStatus(status) { + segments = append(segments, p.n12.Render(strconv.Itoa(int(math.Round(fraction*100)))+"%")) + } + switch { + case hasSteps: + segments = append(segments, p.n7.Render(fmt.Sprintf("%d/%d steps · %s", done, total, duration))) + case duration != "" && duration != na: + segments = append(segments, p.n7.Render(duration)) + } + return strings.Join(segments, " ") +} + +// renderFields draws the two-column summary: muted labels right-padded to the +// longest one, neutral values, a status-colored Status, and blue Run ID / MLflow +// Run hyperlinks. +func renderFields(p palette, colorOn bool, v runView) string { + status := statusStyle(p, v.status).Render("● " + v.status) + rows := []string{ + field(p, "Run ID", link(colorOn, p.blue, v.runID, v.dashboardURL)), + field(p, "Status", status), + field(p, "Submitted", p.n12.Render(v.submitted)), + field(p, "Retries", p.n12.Render(strconv.Itoa(v.retries))), + field(p, "Max Retries", p.n12.Render(v.maxRetries)), + field(p, "Duration", p.n12.Render(v.duration)), + field(p, "Experiment", p.n12.Render(v.experiment)), + field(p, "MLflow Run", link(colorOn, p.blue, v.mlflowLabel, v.mlflowURL)), + field(p, "User", p.n12.Render(v.user)), + field(p, "Accelerators", p.n12.Render(v.accelerators)), + field(p, "Environment", p.n12.Render(v.environment)), + } + return strings.Join(rows, "\n") +} + +// fieldLabelWidth is the width of the longest label ("Accelerators"), so values +// line up in a single column. +const fieldLabelWidth = len("Accelerators") + +func field(p palette, label, value string) string { + return p.n8.Render(label+strings.Repeat(" ", fieldLabelWidth-len(label))) + " " + value +} + +// link renders label as an OSC 8 terminal hyperlink to url in the given style +// (underlined). With color off (or no url) it is just the styled label so the +// box stays aligned; the URLs remain available in JSON output. +func link(colorOn bool, style lipgloss.Style, label, url string) string { + if !colorOn || url == "" { + return style.Render(label) + } + // Wrap the already-styled label in the hyperlink. Passing the OSC 8 escape + // through lipgloss.Render instead corrupts it: lipgloss re-styles each rune + // and splits the "\x1b]8;;" introducer, so the terminal can't parse the + // sequence and prints it literally. + return termenv.Hyperlink(url, style.Underline(true).Render(label)) +} + +// statusStyle maps a run status to its accent color: green for success, red for +// terminal failures, amber for everything still in flight. +func statusStyle(p palette, status string) lipgloss.Style { + switch { + case isSuccessStatus(status): + return p.green + case isFailedStatus(status): + return p.red + default: + return p.amber + } +} + +func isSuccessStatus(status string) bool { + return status == "SUCCESS" +} + +func isFailedStatus(status string) bool { + switch status { + case "FAILED", "TIMEDOUT", "CANCELED", "INTERNAL_ERROR", "UPSTREAM_FAILED", "UPSTREAM_CANCELED": + return true + } + return false +} diff --git a/experimental/air/cmd/render_test.go b/experimental/air/cmd/render_test.go new file mode 100644 index 0000000000..b27a781217 --- /dev/null +++ b/experimental/air/cmd/render_test.go @@ -0,0 +1,206 @@ +package aircmd + +import ( + "bytes" + "io" + "strings" + "testing" + + "github.com/charmbracelet/lipgloss" + "github.com/databricks/cli/libs/cmdio" + "github.com/databricks/databricks-sdk-go/experimental/mocks" + "github.com/databricks/databricks-sdk-go/service/jobs" + "github.com/muesli/termenv" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/mock" +) + +// asciiPalette returns a palette whose styles emit no escape codes, so render +// output is plain text and assertions stay readable. +func asciiPalette() palette { + r := lipgloss.NewRenderer(io.Discard) + r.SetColorProfile(termenv.Ascii) + return newPalette(r) +} + +func TestConfigYAML(t *testing.T) { + ctx := t.Context() + + t.Run("inline parameters include the command as a block literal", func(t *testing.T) { + task := &jobs.GenAiComputeTask{ + YamlParameters: "experiment_name: my-exp\ncompute:\n accelerator_type: a10\n num_accelerators: 1\ncommand: \"for i in $(seq 1 3); do echo $i; done\\n\"\n", + } + got := configYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, task) + assert.Contains(t, got, "experiment_name: my-exp") + assert.Contains(t, got, "accelerator_type: a10") + assert.Contains(t, got, "command: |") + assert.Contains(t, got, " for i in $(seq 1 3); do echo $i; done") + }) + + t.Run("downloads the parameters file when there are no inline parameters", func(t *testing.T) { + m := mocks.NewMockWorkspaceClient(t) + m.GetMockWorkspaceAPI().EXPECT(). + Download(mock.Anything, "/Workspace/cfg.yaml"). + Return(io.NopCloser(strings.NewReader("experiment_name: from-file\n")), nil) + task := &jobs.GenAiComputeTask{YamlParametersFilePath: "/Workspace/cfg.yaml"} + assert.Equal(t, "experiment_name: from-file", configYAML(ctx, m.WorkspaceClient, task)) + }) + + t.Run("falls back to a synthesized config", func(t *testing.T) { + task := &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/exp", + Compute: &jobs.ComputeConfig{GpuType: "a10", NumGpus: 1}, + } + got := configYAML(ctx, mocks.NewMockWorkspaceClient(t).WorkspaceClient, task) + assert.Contains(t, got, "experiment_name: exp") + assert.NotContains(t, got, "command") + }) +} + +func TestSynthConfigYAML(t *testing.T) { + task := &jobs.GenAiComputeTask{ + MlflowExperimentName: "/Users/me@example.com/stream-latency-test", + Compute: &jobs.ComputeConfig{GpuType: "a10", NumGpus: 1}, + } + // The accelerator_type uses the raw GPU type; the command is omitted because + // it lives only in the run parameters. + want := "experiment_name: stream-latency-test\n" + + "compute:\n" + + " accelerator_type: a10\n" + + " num_accelerators: 1" + assert.Equal(t, want, synthConfigYAML(task)) + assert.Empty(t, synthConfigYAML(&jobs.GenAiComputeTask{})) +} + +func TestColorizeConfigLine(t *testing.T) { + p := asciiPalette() + // Under the Ascii profile colorization adds no escapes, so each line is + // preserved verbatim (indentation included) regardless of its role. + for _, line := range []string{ + "experiment_name: stream-latency-test", + "compute:", + " accelerator_type: a10", + " num_accelerators: 1", + "command: |", + ` for i in $(seq 1 10); do echo "step $i"; done`, + } { + assert.Equal(t, line, colorizeConfigLine(p, line)) + } +} + +func TestIsConfigKey(t *testing.T) { + assert.True(t, isConfigKey("experiment_name")) + assert.True(t, isConfigKey("num_accelerators")) + assert.False(t, isConfigKey("")) + assert.False(t, isConfigKey("for i in $(seq 1 10); do echo ")) + assert.False(t, isConfigKey("Command")) +} + +func TestRenderBox(t *testing.T) { + p := asciiPalette() + out := renderBox(p, configBoxTitle, "experiment_name: stream-latency-test\ncompute:") + lines := strings.Split(out, "\n") + + // Title sits in the top border; corners are rounded; every row is the same width. + assert.Contains(t, lines[0], "╭─ "+configBoxTitle+" ") + assert.True(t, strings.HasSuffix(lines[0], "╮")) + assert.True(t, strings.HasPrefix(lines[len(lines)-1], "╰")) + assert.Contains(t, out, "│ experiment_name: stream-latency-test") + + width := lipgloss.Width(lines[0]) + for _, l := range lines { + assert.Equal(t, width, lipgloss.Width(l)) + } +} + +func TestRenderProgressBar(t *testing.T) { + p := asciiPalette() + full := strings.Repeat("█", progressBarWidth) + + t.Run("success is full", func(t *testing.T) { + got := renderProgressBar(p, "SUCCESS", 10, 10, true, "1m 13s") + assert.Equal(t, full+" 100% 10/10 steps · 1m 13s", got) + }) + + t.Run("in progress fills floor(width*done/total)", func(t *testing.T) { + got := renderProgressBar(p, "RUNNING", 5, 10, true, "30s") + want := strings.Repeat("█", 17) + strings.Repeat("░", 17) + " 50% 5/10 steps · 30s" + assert.Equal(t, want, got) + }) + + t.Run("no step data drops the percent and steps", func(t *testing.T) { + got := renderProgressBar(p, "RUNNING", 0, 0, false, "30s") + assert.Equal(t, strings.Repeat("░", progressBarWidth)+" 30s", got) + }) +} + +func TestRenderFields(t *testing.T) { + p := asciiPalette() + out := renderFields(p, false, runView{ + runID: "836121283738861", + dashboardURL: "https://h.test/jobs/runs/836121283738861", + status: "SUCCESS", + submitted: "2026-06-03 04:17 UTC", + retries: 0, + maxRetries: "3", + duration: "1m 13s", + experiment: "stream-latency-test", + mlflowLabel: "stream-latency-test", + mlflowURL: "https://h.test/ml/experiments/E1/runs/R1", + user: "riddhi.bhagwat@databricks.com", + accelerators: "1x A10", + environment: "ml-runtime-gpu:1.0", + }) + + // Labels are padded to the longest ("Accelerators"), so values align. + assert.Contains(t, out, "Run ID ") + assert.Contains(t, out, "Accelerators 1x A10") + // Max retries and environment show alongside the other fields. + assert.Contains(t, out, "Max Retries 3") + assert.Contains(t, out, "Environment ml-runtime-gpu:1.0") + // The status carries its dot prefix. + assert.Contains(t, out, "● SUCCESS") + // Off a terminal, links render as the bare label (URLs live in JSON output). + assert.Contains(t, out, "Run ID 836121283738861") + assert.NotContains(t, out, "https://h.test") + // The field list is a tight block: no blank lines. + assert.NotContains(t, out, "\n\n") +} + +func TestRenderRunTextProgressVisibility(t *testing.T) { + ctx := cmdio.MockDiscard(t.Context()) + w := mocks.NewMockWorkspaceClient(t).WorkspaceClient + render := func(status string) string { + var buf bytes.Buffer + // No tasks and no MLflow ids keeps this offline: only the progress box + // visibility depends on status. + renderRunText(ctx, &buf, w, &jobs.Run{}, &getData{Status: status, RunID: "1"}, nil) + return buf.String() + } + + // A failed run drops the progress box but keeps the metadata. + failed := render("FAILED") + assert.NotContains(t, failed, progressBoxTitle) + assert.Contains(t, failed, metadataBoxTitle) + assert.Contains(t, failed, "● FAILED") + + // A successful run keeps the progress box. + assert.Contains(t, render("SUCCESS"), progressBoxTitle) +} + +func TestLink(t *testing.T) { + p := asciiPalette() + // Color off: the bare label, no URL. + assert.Equal(t, "label", link(false, p.blue, "label", "https://h.test")) + assert.Equal(t, "label", link(false, p.blue, "label", "")) + // With color on, the label is wrapped in an OSC 8 hyperlink to the url. + assert.Contains(t, link(true, p.blue, "label", "https://h.test"), termenv.Hyperlink("https://h.test", "label")) +} + +func TestStatusStyleSelectors(t *testing.T) { + assert.True(t, isSuccessStatus("SUCCESS")) + assert.False(t, isSuccessStatus("RUNNING")) + assert.True(t, isFailedStatus("FAILED")) + assert.True(t, isFailedStatus("TIMEDOUT")) + assert.False(t, isFailedStatus("RUNNING")) +} diff --git a/go.mod b/go.mod index 9329b7ea42..090d567257 100644 --- a/go.mod +++ b/go.mod @@ -24,6 +24,7 @@ require ( github.com/hexops/gotextdiff v1.0.3 // BSD-3-Clause github.com/jackc/pgx/v5 v5.9.2 // MIT github.com/mattn/go-isatty v0.0.22 // MIT + github.com/muesli/termenv v0.16.0 // MIT github.com/palantir/pkg/yamlpatch v1.5.0 // BSD-3-Clause github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // BSD-2-Clause github.com/quasilyte/go-ruleguard/dsl v0.3.22 // BSD-3-Clause @@ -85,7 +86,6 @@ require ( github.com/mitchellh/hashstructure/v2 v2.0.2 // indirect github.com/muesli/ansi v0.0.0-20230316100256-276c6243b2f6 // indirect github.com/muesli/cancelreader v0.2.2 // indirect - github.com/muesli/termenv v0.16.0 // indirect github.com/pkg/errors v0.9.1 // indirect github.com/pmezard/go-difflib v1.0.0 // indirect github.com/rivo/uniseg v0.4.7 // indirect