[air] Add training ETA for running runs in air get and air list

acceptance/experimental/air/get-ai-runtime/output.txt

@@ -47,6 +47,7 @@ MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1
     "attempt_number": 0,
     "experiment_name": "my-exp",
     "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[NUMID]",
-    "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/artifacts/logs/node_0"
+    "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/artifacts/logs/node_0",
+    "est_remaining_seconds": null
   }
 }

acceptance/experimental/air/get/output.txt

@@ -47,7 +47,8 @@ MLflow URL: [DATABRICKS_URL]/ml/experiments/exp1/runs/run1
     "attempt_number": 0,
     "experiment_name": "my-exp",
     "dashboard_url": "[DATABRICKS_URL]/jobs/runs/123?o=[NUMID]",
-    "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/artifacts/logs/node_0"
+    "mlflow_url": "[DATABRICKS_URL]/ml/experiments/exp1/runs/run1/artifacts/logs/node_0",
+    "est_remaining_seconds": null
   }
 }
 

acceptance/experimental/air/list/output.txt

@@ -1,8 +1,8 @@
 
 === list (text)
 >>> [CLI] experimental air list
-   Run ID           Experiment  Status     Started              Duration  MLflow                  User                   Accelerators
-   [NUMID]  qwen-train  ● SUCCESS  [TIMESTAMP]       12s  qwen-train-001          [USERNAME]  8x H100
+   Run ID           Experiment  Status     Started              Duration  ETA  MLflow                  User                   Accelerators
+   [NUMID]  qwen-train  ● SUCCESS  [TIMESTAMP]       12s    -  qwen-train-001          [USERNAME]  8x H100
 
 === list (json)
 >>> [CLI] experimental air list -o json
@@ -25,8 +25,8 @@
 
 === list --all-status (text, via AiTrainingService index)
 >>> [CLI] experimental air list --all-status
-   Run ID           Experiment  Status     Started              Duration  MLflow                  User                   Accelerators
-   [NUMID]  qwen-train  ● SUCCESS  [TIMESTAMP]       12s  qwen-train-001          [USERNAME]  8x H100
+   Run ID           Experiment  Status     Started              Duration  ETA  MLflow                  User                   Accelerators
+   [NUMID]  qwen-train  ● SUCCESS  [TIMESTAMP]       12s    -  qwen-train-001          [USERNAME]  8x H100
 
 === list --all-status (json)
 >>> [CLI] experimental air list --all-status -o json

experimental/air/cmd/eta.go

@@ -0,0 +1,197 @@
+package aircmd
+
+import (
+	"cmp"
+	"context"
+	"fmt"
+	"math"
+	"slices"
+	"strconv"
+
+	"github.com/databricks/cli/libs/log"
+	"github.com/databricks/databricks-sdk-go"
+	"github.com/databricks/databricks-sdk-go/service/ml"
+)
+
+// trainingETA is a best-effort estimate of the wall-clock time remaining for a
+// running training job, derived from MLflow progress signals.
+//
+// It is only produced when MLflow carries both a progress metric and a known
+// total. That combination is logged by HuggingFace Trainer's MLflow integration
+// (a metric per training log carrying the global step and fractional epoch, plus
+// `max_steps` / `num_train_epochs` as params); AIR itself logs only system
+// metrics, which are a wall-clock heartbeat, not training progress. For a run
+// without a known total there is no reliable denominator, so no estimate is
+// shown rather than a misleading one.
+type trainingETA struct {
+	// RemainingSeconds is the projected time to completion.
+	RemainingSeconds int64
+	// Progress is a short human breadcrumb, e.g. "step 4120/10000" or "epoch 1.8/3".
+	Progress string
+}
+
+// etaProgressMetric is the MLflow metric HuggingFace Trainer logs on every
+// training log call. Each point carries the fractional epoch as its value and
+// the global step as its MLflow step, so a single metric history drives both the
+// step-based and epoch-based estimates.
+const etaProgressMetric = "epoch"
+
+// etaWindowPoints bounds how many trailing history points feed the rate, so the
+// estimate reflects recent throughput rather than a warmup-skewed whole-run
+// average.
+const etaWindowPoints = 10
+
+// detailed renders the ETA for the single-run view: "~48m 20s · step 4120/10000".
+func (e *trainingETA) detailed() string {
+	return fmt.Sprintf("~%s · %s", formatDuration(e.RemainingSeconds), e.Progress)
+}
+
+// compact renders the ETA for the list table's ETA column: "~48m 20s".
+func (e *trainingETA) compact() string {
+	return "~" + formatDuration(e.RemainingSeconds)
+}
+
+// estimateTrainingETA fetches a running run's MLflow params and progress-metric
+// history and projects the remaining time. It returns nil whenever an estimate
+// can't be made (no known total, too little history, an API error): the ETA is a
+// convenience, so any failure is logged and treated as "no estimate" rather than
+// failing the command.
+func estimateTrainingETA(ctx context.Context, w *databricks.WorkspaceClient, mlflowRunID string) *trainingETA {
+	if mlflowRunID == "" {
+		return nil
+	}
+
+	resp, err := w.Experiments.GetRun(ctx, ml.GetRunRequest{RunId: mlflowRunID})
+	if err != nil {
+		log.Debugf(ctx, "air: could not fetch MLflow run for ETA: %v", err)
+		return nil
+	}
+	if resp.Run == nil || resp.Run.Data == nil {
+		return nil
+	}
+
+	// The total is only present for Trainer runs; skip the metric-history fetch
+	// entirely when it is absent, so a run we can't estimate costs one call, not two.
+	params := paramMap(resp.Run.Data.Params)
+	if _, hasSteps := positiveInt(params["max_steps"]); !hasSteps {
+		if _, hasEpochs := positiveFloat(params["num_train_epochs"]); !hasEpochs {
+			return nil
+		}
+	}
+
+	history, err := w.Experiments.GetHistoryAll(ctx, ml.GetHistoryRequest{
+		RunId:     mlflowRunID,
+		MetricKey: etaProgressMetric,
+	})
+	if err != nil {
+		log.Debugf(ctx, "air: could not fetch %q metric history for ETA: %v", etaProgressMetric, err)
+		return nil
+	}
+	return computeETA(params, history)
+}
+
+// computeETA is the pure projection: given a run's params and the ascending
+// history of the progress metric, it returns the remaining-time estimate, or nil
+// when it can't be computed. Split from estimateTrainingETA so it can be tested
+// without an API client.
+//
+// When max_steps is set (> 0) HuggingFace Trainer trains by step and overrides
+// num_train_epochs, so the step total wins when both are present.
+func computeETA(params map[string]string, history []ml.Metric) *trainingETA {
+	if len(history) < 2 {
+		return nil
+	}
+
+	// History should already be ordered by step; sort by timestamp defensively so
+	// the rate is measured over increasing wall-clock time.
+	points := append([]ml.Metric(nil), history...)
+	slices.SortStableFunc(points, func(a, b ml.Metric) int { return cmp.Compare(a.Timestamp, b.Timestamp) })
+	if len(points) > etaWindowPoints {
+		points = points[len(points)-etaWindowPoints:]
+	}
+	first, last := points[0], points[len(points)-1]
+
+	elapsedSec := float64(last.Timestamp-first.Timestamp) / 1000.0
+	if elapsedSec <= 0 {
+		return nil
+	}
+
+	if maxSteps, ok := positiveInt(params["max_steps"]); ok {
+		return stepETA(first, last, elapsedSec, maxSteps)
+	}
+	if totalEpochs, ok := positiveFloat(params["num_train_epochs"]); ok {
+		return epochETA(first, last, elapsedSec, totalEpochs)
+	}
+	return nil
+}
+
+// stepETA projects remaining time from the global step (the metric's MLflow
+// step) against a known max_steps.
+func stepETA(first, last ml.Metric, elapsedSec float64, maxSteps int64) *trainingETA {
+	current := last.Step
+	if current <= 0 || current >= maxSteps {
+		return nil
+	}
+	done := last.Step - first.Step
+	if done <= 0 {
+		return nil
+	}
+	perSec := float64(done) / elapsedSec
+	remaining := float64(maxSteps-current) / perSec
+	return &trainingETA{
+		RemainingSeconds: int64(math.Round(remaining)),
+		Progress:         fmt.Sprintf("step %d/%d", current, maxSteps),
+	}
+}
+
+// epochETA projects remaining time from the fractional epoch (the metric's
+// value) against a known num_train_epochs.
+func epochETA(first, last ml.Metric, elapsedSec, totalEpochs float64) *trainingETA {
+	current := last.Value
+	if current <= 0 || current >= totalEpochs {
+		return nil
+	}
+	done := last.Value - first.Value
+	if done <= 0 {
+		return nil
+	}
+	perSec := done / elapsedSec
+	remaining := (totalEpochs - current) / perSec
+	return &trainingETA{
+		RemainingSeconds: int64(math.Round(remaining)),
+		Progress:         fmt.Sprintf("epoch %.1f/%s", current, trimFloat(totalEpochs)),
+	}
+}
+
+// paramMap indexes MLflow params by key.
+func paramMap(params []ml.Param) map[string]string {
+	m := make(map[string]string, len(params))
+	for _, p := range params {
+		m[p.Key] = p.Value
+	}
+	return m
+}
+
+// positiveInt parses s as an integer and reports it only when it is > 0. Trainer
+// logs max_steps=-1 when training by epochs, which this rejects.
+func positiveInt(s string) (int64, bool) {
+	n, err := strconv.ParseInt(s, 10, 64)
+	if err != nil || n <= 0 {
+		return 0, false
+	}
+	return n, true
+}
+
+// positiveFloat parses s as a float and reports it only when it is > 0.
+func positiveFloat(s string) (float64, bool) {
+	f, err := strconv.ParseFloat(s, 64)
+	if err != nil || f <= 0 {
+		return 0, false
+	}
+	return f, true
+}
+
+// trimFloat renders a float without trailing zeros: 3.0 -> "3", 2.5 -> "2.5".
+func trimFloat(f float64) string {
+	return strconv.FormatFloat(f, 'f', -1, 64)
+}

experimental/air/cmd/eta_test.go

@@ -0,0 +1,94 @@
+package aircmd
+
+import (
+	"testing"
+
+	"github.com/databricks/databricks-sdk-go/service/ml"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// metric builds one epoch-metric history point: fractional epoch as the value,
+// global step as the MLflow step, wall-clock as the timestamp (ms).
+func metric(tsMillis, step int64, epoch float64) ml.Metric {
+	return ml.Metric{Key: "epoch", Timestamp: tsMillis, Step: step, Value: epoch}
+}
+
+func TestComputeETA_StepMode(t *testing.T) {
+	// max_steps wins over num_train_epochs. 4000 steps in 100s => 40 steps/s;
+	// 5000 steps remain => 125s.
+	eta := computeETA(
+		map[string]string{"max_steps": "10000", "num_train_epochs": "3"},
+		[]ml.Metric{metric(0, 1000, 0.3), metric(100_000, 5000, 1.5)},
+	)
+	require.NotNil(t, eta)
+	assert.Equal(t, int64(125), eta.RemainingSeconds)
+	assert.Equal(t, "step 5000/10000", eta.Progress)
+}
+
+func TestComputeETA_EpochMode(t *testing.T) {
+	// No usable max_steps (Trainer logs -1 when training by epochs), so the
+	// fractional epoch drives it: 1.0 epoch in 100s => 0.01 epoch/s; 1.5 remain
+	// of 3 => 150s.
+	eta := computeETA(
+		map[string]string{"max_steps": "-1", "num_train_epochs": "3"},
+		[]ml.Metric{metric(0, 100, 0.5), metric(100_000, 900, 1.5)},
+	)
+	require.NotNil(t, eta)
+	assert.Equal(t, int64(150), eta.RemainingSeconds)
+	assert.Equal(t, "epoch 1.5/3", eta.Progress)
+}
+
+func TestComputeETA_UsesTrailingWindow(t *testing.T) {
+	// Early points are slow (10 steps over the first 100s), recent points fast
+	// (100 steps/10s). The window should reflect the recent rate, not the whole
+	// run's average. Build 12 points; only the last etaWindowPoints (10) count.
+	var pts []ml.Metric
+	pts = append(pts, metric(0, 0, 0), metric(100_000, 10, 0.1))
+	// Ten fast points: +100 steps every 10s starting at step 10, t=100s.
+	for i := int64(1); i <= 10; i++ {
+		pts = append(pts, metric(100_000+i*10_000, 10+i*100, 0.1+float64(i)*0.1))
+	}
+	eta := computeETA(map[string]string{"max_steps": "2000"}, pts)
+	require.NotNil(t, eta)
+	// Windowed rate is 10 steps/s; current step 1010, so 990 remain => 99s.
+	// (The whole-run average would be far slower and give a larger estimate.)
+	assert.Equal(t, int64(99), eta.RemainingSeconds)
+	assert.Equal(t, "step 1010/2000", eta.Progress)
+}
+
+func TestComputeETA_NoEstimate(t *testing.T) {
+	base := []ml.Metric{metric(0, 100, 0.3), metric(100_000, 5000, 1.5)}
+	cases := []struct {
+		name    string
+		params  map[string]string
+		history []ml.Metric
+	}{
+		{"no total param", map[string]string{}, base},
+		{"max_steps not positive", map[string]string{"max_steps": "-1"}, base},
+		{"num_train_epochs zero", map[string]string{"num_train_epochs": "0"}, base},
+		{"total param unparseable", map[string]string{"max_steps": "lots"}, base},
+		{"too few points", map[string]string{"max_steps": "10000"}, base[:1]},
+		{"no elapsed time", map[string]string{"max_steps": "10000"}, []ml.Metric{metric(500, 100, 0.3), metric(500, 5000, 1.5)}},
+		{"already at max steps", map[string]string{"max_steps": "5000"}, base},
+		{"epoch already complete", map[string]string{"num_train_epochs": "1.5"}, base},
+		{"step not advancing", map[string]string{"max_steps": "10000"}, []ml.Metric{metric(0, 5000, 1.5), metric(100_000, 5000, 1.5)}},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Nil(t, computeETA(tc.params, tc.history))
+		})
+	}
+}
+
+func TestTrainingETADisplay(t *testing.T) {
+	eta := &trainingETA{RemainingSeconds: 2900, Progress: "step 5000/10000"}
+	assert.Equal(t, "~48m 20s · step 5000/10000", eta.detailed())
+	assert.Equal(t, "~48m 20s", eta.compact())
+}
+
+func TestTrimFloat(t *testing.T) {
+	assert.Equal(t, "3", trimFloat(3.0))
+	assert.Equal(t, "2.5", trimFloat(2.5))
+	assert.Equal(t, "0.1", trimFloat(0.1))
+}

experimental/air/cmd/format.go

@@ -84,6 +84,13 @@ var gpuDisplayNames = map[string]string{
 	"GPU_1xH100": "H100",
 }
 
+// isRunning reports whether a run is currently executing (lifecycle RUNNING), as
+// opposed to pending, terminal, or state-unknown. Only a running run has a
+// meaningful remaining-time estimate.
+func isRunning(run *jobs.Run) bool {
+	return run.State != nil && run.State.LifeCycleState == jobs.RunLifeCycleStateRunning
+}
+
 // runStatus returns the single status word to show for a run. The backend
 // reports two values: a lifecycle state (e.g. PENDING, RUNNING) and, once the
 // run has finished, a result state (e.g. SUCCESS, FAILED). The result state is

experimental/air/cmd/get.go

@@ -29,6 +29,9 @@ type getData struct {
 	ExperimentName  *string `json:"experiment_name"`
 	DashboardURL    string  `json:"dashboard_url"`
 	MLflowURL       *string `json:"mlflow_url"`
+	// ESTRemainingSeconds is a best-effort estimate of the time left for a
+	// running run, or null when one can't be made (see estimateTrainingETA).
+	ESTRemainingSeconds *int64 `json:"est_remaining_seconds"`
 
 	// The fields below are pre-rendered text-view cells, excluded from JSON
 	// (matching `air get --json`). Each shows "N/A" when its value is
@@ -42,6 +45,9 @@ type getData struct {
 	AcceleratorsDisplay string `json:"-"`
 	EnvironmentDisplay  string `json:"-"`
 	MaxRetriesDisplay   string `json:"-"`
+	// ETADisplay is the pre-rendered "ETA" cell ("~48m 20s · step 4120/10000"),
+	// set only for a running run with an estimable remaining time.
+	ETADisplay string `json:"-"`
 	// TrainingConfigPath is the run's config file, downloaded for the config box.
 	TrainingConfigPath string `json:"-"`
 	// Sweep replaces the single-run view for foreach runs.
@@ -155,6 +161,15 @@ func newGetCommand() *cobra.Command {
 		if ids != nil {
 			url := mlflowLogsURL(w.Config.Host, ids)
 			data.MLflowURL = &url
+			// A remaining-time estimate only makes sense while the run is still
+			// training; a terminal run is either done or stopped mid-progress.
+			if isRunning(run) {
+				if eta := estimateTrainingETA(ctx, w, ids.RunID); eta != nil {
+					secs := eta.RemainingSeconds
+					data.ESTRemainingSeconds = &secs
+					data.ETADisplay = eta.detailed()
+				}
+			}
 		}
 		if task := findForEachTask(run); task != nil {
 			data.Sweep = buildSweepInfo(ctx, w, task)