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

[vinchenzo-db]

Surface an estimated time-remaining for in-flight air run GPU jobs, derived from MLflow progress signals.

MLflow has no native ETA and AIR itself logs only system heartbeat metrics, so a genuine estimate is only computable for HuggingFace Trainer runs: its MLflow integration logs max_steps / num_train_epochs params and an epoch metric carrying the global step and fractional epoch. The new eta.go gates on those params, then projects the remaining time from a trailing window of the epoch metric's history (step-based when max_steps is set, else epoch-based). Everything is best-effort: any missing signal or API error yields no estimate rather than failing the command, so non-Trainer runs simply show nothing.

air get shows an "ETA" row for a running run (and est_remaining_seconds in JSON); air list adds an "ETA" column computed per running row in the existing parallel MLflow enrichment.

Changes

Why

Tests

Repro Steps:

Env:

$ cd ~/databricks-cli                                                                            
git checkout air-training-eta        # the branch with the ETA feature
go build -o /tmp/aircli .
alias air='/tmp/aircli experimental air'

Dummy Training Script:

$ mkdir -p ~/air-eta-e2e                                                                         

cat > ~/air-eta-e2e/eta_trainer.py <<'EOF'
import torch
from torch.utils.data import Dataset
from transformers import GPT2Config, GPT2ForSequenceClassification, Trainer, TrainingArguments

VOCAB, SEQ, NUM_LABELS = 512, 512, 2


class SyntheticDataset(Dataset):
    def __init__(self, n):
        self.n = n

    def __len__(self):
        return self.n

    def __getitem__(self, i):
        g = torch.Generator().manual_seed(i)
        return {
            "input_ids": torch.randint(0, VOCAB, (SEQ,), generator=g),
            "labels": int(torch.randint(0, NUM_LABELS, (1,), generator=g).item()),
        }


def main():
    config = GPT2Config(
        vocab_size=VOCAB,
        n_positions=SEQ,
        n_embd=1024,
        n_layer=24,
        n_head=16,
        num_labels=NUM_LABELS,
        pad_token_id=0,
        bos_token_id=0,
        eos_token_id=0,
    )
    model = GPT2ForSequenceClassification(config)
    args = TrainingArguments(
        output_dir="/tmp/eta-out",
        max_steps=600,
        logging_steps=5,
        per_device_train_batch_size=8,
        report_to=["mlflow"],
        save_strategy="no",
        eval_strategy="no",
        disable_tqdm=True,
    )
    # Do NOT start an mlflow run here: HF's MLflowCallback resumes the run named
    # by the MLFLOW_RUN_ID env var that AIR already sets. Starting one ourselves
    # double-activates it and crashes at on_train_begin.
    Trainer(model=model, args=args, train_dataset=SyntheticDataset(1_000_000)).train()


if __name__ == "__main__":
    main()
EOF

Yaml:

$ cat > ~/air-eta-e2e/eta_run.yaml <<'EOF'                                                    
experiment_name: air-eta-demo

environment:
  docker_image:
    url: databricksruntime/environment-test:sgcdcs-test-devel-aws-pytorch

compute:
  num_accelerators: 1
  accelerator_type: GPU_1xA10

max_retries: 0
timeout_minutes: 60

code_source:
  type: snapshot
  snapshot:
    root_path: .
        
command: |-
  set -uxo pipefail
  cd "$CODE_SOURCE_PATH"
  # AIR nodes can't reach the internal pypi proxy, but public PyPI works.
  pip install --no-cache-dir transformers "accelerate>=1.1.0" "mlflow-skinny>=2.15"
  python eta_trainer.py
EOF

Run:

$ air register-image databricksruntime/environment-test:sgcdcs-test-devel-aws-pytorch -p $PROFILE
air run -f eta_run.yaml -p $PROFILE
Image already up to date: sha256:32c9914c5...

To use this image in your training config:
  environment:
    docker_image:
      url: databricksruntime/environment-test:sgcdcs-test-devel-aws-pytorch
Submitting experiment: air-eta-demo
Uploading air-eta-e2e_20260825_191509.tar.gz...
Submitted workload with Job Run ID: 629555509101647
View job run at: https://e2-dogfood.staging.cloud.databricks.com/jobs/runs/629555509101647
View MLflow run at: https://e2-dogfood.staging.cloud.databricks.com/ml/experiments/2492474306876561/runs/672ed4a0bcda4383b35e79d6986b2c1b
View MLflow experiment at: https://e2-dogfood.staging.cloud.databricks.com/ml/experiments/2492474306876561

Model hasn't started training yet so don't display ETA:

Model starts training:

[github-actions]

Approval status: pending

/acceptance/experimental/air/ - needs approval

Files: acceptance/experimental/air/get-ai-runtime/output.txt, acceptance/experimental/air/get/output.txt, acceptance/experimental/air/list/output.txt
Suggested: @riddhibhagwat-db
Also eligible: @apeforest, @bfontain, @lu-wang-dl, @panchalhp-db, @maggiewang-db, @ben-hansen-db, @pardis-beikzadeh-db

/experimental/air/ - needs approval

8 files changed
Suggested: @riddhibhagwat-db
Also eligible: @apeforest, @bfontain, @lu-wang-dl, @panchalhp-db, @maggiewang-db, @ben-hansen-db, @pardis-beikzadeh-db

_{Any maintainer (@andrewnester, @anton-107, @denik, @pietern, @shreyas-goenka, @simonfaltum, @renaudhartert-db, @janniklasrose, @lennartkats-db, @rugpanov, @rclarey) can approve all areas.
See OWNERS for ownership rules.}

[eng-dev-ecosystem-bot]

Integration test report

Commit: 4ee8293

Run: 32884792580

	Env	💚​RECOVERED	🙈​SKIP	✅​pass	🙈​skip	Time
💚​	aws linux	1	4	274	1184	5:09
💚​	aws windows	1	4	276	1182	3:59
💚​	azure linux	1	4	273	1184	4:45
💚​	azure windows	1	4	275	1182	3:17
💚​	gcp linux	1	4	274	1184	5:13
💚​	gcp windows	1	4	276	1182	4:44

	Test Name	aws linux	aws windows	azure linux	azure windows	gcp linux	gcp windows
💚​	TestAccept	💚​R	💚​R	💚​R	💚​R	💚​R	💚​R
🙈​	TestAccept/bundle/invariant/no_drift	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S
🙈​	TestAccept/bundle/resources/vector_search_endpoints/drift/recreated_same_name	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S
🙈​	TestAccept/bundle/resources/vector_search_indexes/recreate/embedding_dimension	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S
🙈​	TestAccept/ssh/connection	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S	🙈​S

Top 3 slowest tests (at least 2 minutes):

duration	env	testname
3:14	aws windows	TestAccept
3:14	azure windows	TestAccept
3:11	gcp windows	TestAccept