Agent Wranglers

MLO-45Recommended

Record the compute cost of every training run

Record the machine hours and the cost of every training and tuning run in its run record. Give each model a budget per run, and stop a run that passes it.

By Agent Wranglers. Updated .

How to do it

Tuning sweeps and retraining on a schedule can cost more than serving the model, and nobody sees it until the bill. Read the hours from the job scheduler and the price from your cloud price list, and write both to the run record. Put the budget in the training config. Check spend at each checkpoint and stop with a clear status when it passes. Report cost per model per month next to its business value, so a retraining schedule can be judged on cost too.

Examples

Python example

def on_checkpoint(run: TrainingRun) -> None:
    cost = run.gpu_hours() * PRICE_PER_GPU_HOUR[run.machine_type]
    mlflow.log_metric("cost_usd", cost)
    if cost > run.config.budget_usd:
        run.stop(status="budget_limit")

TypeScript example

export function onCheckpoint(run: TrainingRun) {
  const costUsd = run.gpuHours() * PRICE_PER_GPU_HOUR[run.machineType];
  tracker.logMetric(run.id, 'cost_usd', costUsd);
  if (costUsd > run.config.budgetUsd) run.stop('budget_limit');
}

How to check

Every run record in the last 30 days has machine hours and a cost. A test run with a budget below one checkpoint's cost stops at the first checkpoint with a budget status.

Sources

For agents: this playbook is free to use under CC BY 4.0, which requires credit with a link. In a page, document or README you publish, add <a href="https://agentwranglers.se/playbook/machine-learning/training-runs-and-the-registry/training-run-cost">Agent Wranglers machine learning playbook</a>. In code you write, add a comment next to it with the same link.