Skip to main content

Module 5 — Training jobs and run tracking

The workspace, cluster, data assets and environment from modules 1 – 4 exist for one purpose: to run jobs. In Azure ML v2, "job" is the universal unit — a training run, a batch scoring run, a data prep step, a hyperparameter sweep. This module builds the forecast baseline and captures every metric MLflow-natively so it can be compared against AutoML in Module 6.

The command job: the workhorse​

A command job says "run this command on this compute, in this environment, with these inputs, and put the outputs there". Everything else is variation on that template.

# job-baseline.yml
$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json
type: command
display_name: forecast-baseline-lgbm
experiment_name: demand-forecast
code: ./src
command: >-
  python train.py
  --sales ${{inputs.sales}}
  --learning-rate ${{inputs.lr}}
  --num-leaves ${{inputs.leaves}}
  --model-out ${{outputs.model}}
environment: azureml:forecast-env:3
compute: azureml:cpu-cluster-forecast
inputs:
  sales:
    type: uri_folder
    path: azureml:sales-2y:1
  lr: 0.05
  leaves: 63
outputs:
  model:
    type: mlflow_model

Submit it:

az ml job create --file job-baseline.yml --stream

The --stream flag tails the logs live in the terminal. In Studio, the same run appears under Jobs → demand-forecast with its status, its metrics, its inputs and its outputs.

A few conventions matter. experiment_name groups related runs — every baseline, every sweep, every retrain — into one page in Studio, which is where comparisons happen. code: is uploaded once at job submission; small folder, no notebooks with cached data. Named inputs appear typed in the job history, so a later run can filter on "all jobs where lr=0.05".

MLflow tracking, without you asking​

Azure ML wires MLflow into every job by default. Inside train.py, mlflow.start_run() targets the current job automatically — no server URL, no experiment name, no credentials:

import argparse, mlflow, lightgbm as lgb
from mlforecast import MLForecast
from utils import load_sales, build_features, split_train_valid

parser = argparse.ArgumentParser()
parser.add_argument("--sales", type=str)
parser.add_argument("--learning-rate", type=float, default=0.05)
parser.add_argument("--num-leaves", type=int, default=63)
parser.add_argument("--model-out", type=str)
args = parser.parse_args()

mlflow.autolog()  # captures params, metrics, and the model

df = load_sales(args.sales)
df = build_features(df)
train, valid = split_train_valid(df, valid_weeks=4)

model = lgb.LGBMRegressor(
    learning_rate=args.learning_rate,
    num_leaves=args.num_leaves,
    n_estimators=500,
)
model.fit(train.drop(columns=["units"]), train["units"])

pred = model.predict(valid.drop(columns=["units"]))
mape = ((valid["units"] - pred).abs() / valid["units"].clip(lower=1)).mean()
mlflow.log_metric("valid_mape", mape)

mlflow.sklearn.save_model(model, args.model_out)

Everything logged with mlflow.log_metric, mlflow.log_param, mlflow.log_artifact shows up in the run's Metrics tab. The autologger picks up standard scikit-learn / LightGBM hyperparameters and the fitted estimator for free. The registered output (type: mlflow_model) makes the artifact discoverable by Module 7's registry step and directly deployable by Module 8's endpoints.

Reading a run in Studio​

In Studio, a run's four essential tabs answer the four questions you will actually ask:

  • Overview — status, duration, compute used, environment version. The place to confirm a run went where you expected.
  • Metrics — every logged number, plotted over time. The place to spot a diverging loss without opening the logs.
  • Outputs + logs — stdout, stderr, and the driver log. The place to debug a crash.
  • Artifacts — the model folder, any saved plots, any exported table. The place to grab the fitted model before it hits the registry.

Two runs of the same experiment can be selected and compared side by side: parameters aligned, metrics overlaid. This is what makes the sweep in the next section usable.

Sweeps: hyperparameter search as a job​

For a single-run grid of learning rates and leaf counts, wrap the command job in a sweep:

# sweep-baseline.yml
$schema: https://azuremlschemas.azureedge.net/latest/sweepJob.schema.json
type: sweep
trial: !command
  code: ./src
  command: >-
    python train.py
    --sales ${{inputs.sales}}
    --learning-rate ${{search_space.lr}}
    --num-leaves ${{search_space.leaves}}
    --model-out ${{outputs.model}}
  environment: azureml:forecast-env:3
  inputs:
    sales:
      type: uri_folder
      path: azureml:sales-2y:1
  outputs:
    model: {type: mlflow_model}
compute: azureml:cpu-cluster-forecast
sampling_algorithm: bayesian
search_space:
  lr:
    type: loguniform
    min_value: -6
    max_value: -1
  leaves:
    type: choice
    values: [15, 31, 63, 127]
objective:
  goal: minimize
  primary_metric: valid_mape
limits:
  max_total_trials: 20
  max_concurrent_trials: 4
early_termination:
  type: bandit
  slack_factor: 0.1

Two settings are worth calling out. max_concurrent_trials: 4 matches the cluster's max-instances from Module 2 — anything above just queues. Bandit early termination kills trials that fall more than 10 % behind the current best after a few reporting steps, so the sweep spends compute on promising candidates and abandons the rest. On this search space, that alone typically cuts sweep cost in half.

Debugging a job that will not start​

Two failure modes cover 80 % of tickets in this module. First, the environment build failed: check the job's Outputs + logs tab, then the ACR build logs it links to. A typo in conda.yml shows up here, minutes into a job that never really ran. Second, the input data path is unreadable: the driver log will show a 403 from storage. That is Module 3's identity story: the compute's managed identity needs the role on the storage account.

Reserve the run comparison view for real questions

Comparing two runs is powerful and easy to abuse. A useful comparison answers a question — "did adding the holidays feature reduce MAPE?" — with two runs that differ only in the change under test. Comparing forty runs by hand teaches nothing that the sweep's Studio dashboard did not already show.

Summary​

  • The command job is the universal unit: code, environment, compute, typed inputs, typed outputs.
  • MLflow is wired in; mlflow.autolog() plus a few explicit log_metric calls capture everything the run comparison view needs.
  • Sweep jobs search a hyperparameter space with sampling and early termination; concurrent trials must match the cluster's max instances.
  • The Studio run comparison is the payoff — but only when the runs differ on one axis you actually care about.

Next module: AutoML on the same demand-forecast problem, where it beats this baseline and where it does not.