AZURE ML: Automate ML training with GitHub Actions using OIDC federation, az ml CLI v2, environment promotion, and metric-gated registration

Picture a Leeds mortgage lender pushing roughly 4,200 affordability decisions an hour through a gradient-boosted credit risk model. The model gets retrained every Tuesday morning because the BoE base rate moves, broker mix shifts, and the data science team keeps tweaking the feature set. For nine months the retrain was a human ritual: someone on the data team SSHed into a dusty VM, ran a notebook, copied a pickle file into blob storage, then pinged the MLOps lead on Teams to register the new version. Two months ago that ritual failed silently. A stale model served bad scores for six business days before anyone noticed. The team got hauled in front of the model risk committee, and the verdict was loud and clear: the retrain pipeline ships through GitHub Actions, federated identity, no secrets in CI, and a metric gate decides if the new model gets to wear the production tag. That is exactly what we are going to build.

Tools used

• Azure ML CLI v2 2.32.0 on the runner (az extension add -n ml)
• Azure CLI 2.65.0 (azure/login@v2 action handles install)
• Azure ML Python SDK v2 1.21.0 for the registration gate script
• MLflow 2.16.2 for run tracking and metric reads
• GitHub Actions runner: ubuntu-22.04 with OIDC token endpoint enabled
• azure/login@v2 GitHub Action with federated credentials (no client secret)
• azure/setup-python@v5 pinned to Python 3.11
• LightGBM 4.5.0 as the training framework
• Azure ML workspace aml-credit-prod in uksouth, compute cluster cpu-cluster-prod (Standard_D8s_v5, min 0 / max 4 nodes)
• Resource group rg-ml-credit-prod-uks, ACR acrmlcreditproduks, key vault kv-credit-prod-uks
• A registered Azure ML environment lgbm-train-env:7 (curated base + LightGBM 4.5.0)

Prerequisites

• Owner or User Access Administrator on rg-ml-credit-prod-uks so you can create the federated credential.
• azureml-runner user-assigned managed identity already exists in rg-ml-credit-prod-uks and holds AzureML Data Scientist plus Contributor on the workspace scope.
• GitHub repo leeds-mortgage/credit-risk-model with branch protection on main and the Settings -> Secrets and variables -> Actions page reachable.
• Local az CLI 2.65+ and az extension add -n ml --version 2.32.0 for the one-off federation setup.
• Workspace aml-credit-prod already provisioned, dataset affordability-features registered as a versioned mltable asset, and an existing baseline model credit-risk-lgbm at version 11 in the workspace registry (this is the model the gate compares against).
• A trained model lineage that already writes auc_pr, auc_roc, and ks_stat to MLflow on every run. If you not tracking those today, stop and add them, the gate is meaningless otherwise.

Project Architecture

The repo holds three things: the training code under src/train.py, the Azure ML job spec at aml/job-train.yml, and the GitHub Actions workflow at .github/workflows/retrain.yml. When a PR opens against main, Actions runs a smoke job against a 5% sample on the same compute cluster. When a merge lands on main, Actions runs the full retrain against the latest version of the affordability-features mltable. The runner does not hold any Azure secret. It exchanges its OIDC token at https://token.actions.githubusercontent.com for a short-lived Azure AD access token, scoped only to the azureml-runner UAMI via a federated credential bound to the exact repo, branch, and environment.

The training job runs inside the workspace on cpu-cluster-prod. MLflow auto-logging captures the run. Once the job finishes, a follow-up Actions step pulls the metrics off the completed run, compares against the currently-deployed model's metrics held as workspace tags, and only if the new run beats the threshold does it call az ml model create to register a new version. Nothing on the runner ever sees a long-lived credential, ever. No client secret, no SP password, no SAS token sitting in a repo secret pretending to be safe.

Step 1. Wire up OIDC federation between the repo and the managed identity.

This is the one-off setup. Run it locally, signed in as someone with permission to write to the UAMI.

SUBSCRIPTION_ID="8c1a4f5e-6b22-4f8d-9c0a-3e1b5d4f8a7c"
RG="rg-ml-credit-prod-uks"
UAMI_NAME="azureml-runner"
REPO="leeds-mortgage/credit-risk-model"

az account set --subscription "$SUBSCRIPTION_ID"

UAMI_OBJECT_ID=$(az identity show -g "$RG" -n "$UAMI_NAME" --query principalId -o tsv)
UAMI_CLIENT_ID=$(az identity show -g "$RG" -n "$UAMI_NAME" --query clientId -o tsv)
TENANT_ID=$(az account show --query tenantId -o tsv)

# Federated cred for pushes to main
az identity federated-credential create \
  --name gh-main-branch \
  --identity-name "$UAMI_NAME" \
  --resource-group "$RG" \
  --issuer "https://token.actions.githubusercontent.com" \
  --subject "repo:${REPO}:ref:refs/heads/main" \
  --audiences "api://AzureADTokenExchange"

# Federated cred for PRs (so smoke job can run)
az identity federated-credential create \
  --name gh-pull-request \
  --identity-name "$UAMI_NAME" \
  --resource-group "$RG" \
  --issuer "https://token.actions.githubusercontent.com" \
  --subject "repo:${REPO}:pull_request" \
  --audiences "api://AzureADTokenExchange"

echo "AZURE_CLIENT_ID=$UAMI_CLIENT_ID"
echo "AZURE_TENANT_ID=$TENANT_ID"
echo "AZURE_SUBSCRIPTION_ID=$SUBSCRIPTION_ID"

The subject claim is the bit that matters. Azure AD will only mint a token for the UAMI when the incoming OIDC token's sub field exactly matches. That binds the credential to one repo and one ref. A fork cannot mint a token. A different branch cannot mint a token. The PR credential uses pull_request so any PR head ref in this repo can run the smoke job, which is what we want.

Take the three echoed values and stash them as repository variables (not secrets, they are not secrets) under Settings -> Secrets and variables -> Actions -> Variables.

Step 2. Build the Azure ML environment the training job will run in.

We register a curated-base + LightGBM environment so the job always runs against pinned versions. Drop this into aml/env.yml.

$schema: https://azuremlschemas.azureedge.net/latest/environment.schema.json
name: lgbm-train-env
version: 7
image: mcr.microsoft.com/azureml/curated/sklearn-1.5:18
conda_file:
  channels:
    - conda-forge
  dependencies:
    - python=3.11
    - pip=24.2
    - pip:
        - lightgbm==4.5.0
        - mlflow==2.16.2
        - azureml-mlflow==1.59.0
        - pandas==2.2.3
        - scikit-learn==1.5.2
        - pyarrow==17.0.0
description: LightGBM 4.5 training env for credit-risk retrain pipeline

Register it once from your laptop so it has version 7 before the workflow ever runs:

az ml environment create \
  -f aml/env.yml \
  --resource-group rg-ml-credit-prod-uks \
  --workspace-name aml-credit-prod

The reason we pin the image to a specific MCR tag rather than :latest is that "latest" floats and breaks reproducibility. When the model risk committee asks "what scikit-learn version trained the model in production on 14 May" you want a deterministic answer.

Step 3. Write the training job spec.

This is the YAML the runner will submit to the workspace. Save as aml/job-train.yml.

$schema: https://azuremlschemas.azureedge.net/latest/commandJob.schema.json
type: command
display_name: credit-risk-retrain
experiment_name: credit-risk-retrain
description: Weekly retrain of the LightGBM affordability model
code: ../src
command: >-
  python train.py
  --input ${{inputs.training_data}}
  --learning-rate ${{inputs.learning_rate}}
  --num-leaves ${{inputs.num_leaves}}
  --sample-fraction ${{inputs.sample_fraction}}
  --output-dir ${{outputs.model_dir}}
inputs:
  training_data:
    type: mltable
    path: azureml:affordability-features@latest
    mode: ro_mount
  learning_rate: 0.04
  num_leaves: 127
  sample_fraction: 1.0
outputs:
  model_dir:
    type: uri_folder
    mode: rw_mount
environment: azureml:lgbm-train-env:7
compute: azureml:cpu-cluster-prod
resources:
  instance_count: 1
services:
  Studio:
    type: jupyter_lab
tags:
  triggered_by: github-actions
  cost_centre: credit-risk-mlops

Two things to notice. First, sample_fraction is a real input. The PR smoke run will override that to 0.05 so the smoke costs pennies. Second, the triggered_by tag is how finance reconciles the Azure bill back to the right cost centre at month end. Tag your jobs, your future self will thank you when accounting comes knocking.

Step 4. Write the GitHub Actions workflow with OIDC login.

This is the heart of it. Save as .github/workflows/retrain.yml.

name: credit-risk-retrain

on:
  pull_request:
    branches: [main]
    paths:
      - 'src/**'
      - 'aml/**'
      - '.github/workflows/retrain.yml'
  push:
    branches: [main]
    paths:
      - 'src/**'
      - 'aml/**'
  workflow_dispatch:
    inputs:
      learning_rate:
        description: LightGBM learning rate
        required: false
        default: '0.04'
      num_leaves:
        description: LightGBM num_leaves
        required: false
        default: '127'

permissions:
  id-token: write
  contents: read

env:
  AZ_RG: rg-ml-credit-prod-uks
  AZ_WS: aml-credit-prod
  AZ_LOC: uksouth

jobs:
  train:
    runs-on: ubuntu-22.04
    environment: ${{ github.event_name == 'push' && 'prod' || 'pr' }}
    steps:
      - uses: actions/checkout@v4

      - name: Azure login via OIDC
        uses: azure/login@v2
        with:
          client-id: ${{ vars.AZURE_CLIENT_ID }}
          tenant-id: ${{ vars.AZURE_TENANT_ID }}
          subscription-id: ${{ vars.AZURE_SUBSCRIPTION_ID }}

      - name: Install az ml extension
        run: az extension add -n ml --version 2.32.0 -y

      - name: Set sample fraction by event
        id: vars
        run: |
          if [ "${{ github.event_name }}" = "pull_request" ]; then
            echo "SAMPLE=0.05" >> $GITHUB_OUTPUT
            echo "LR=0.04"     >> $GITHUB_OUTPUT
            echo "LEAVES=63"   >> $GITHUB_OUTPUT
          else
            echo "SAMPLE=1.0"  >> $GITHUB_OUTPUT
            echo "LR=${{ github.event.inputs.learning_rate || '0.04' }}" >> $GITHUB_OUTPUT
            echo "LEAVES=${{ github.event.inputs.num_leaves || '127' }}" >> $GITHUB_OUTPUT
          fi

      - name: Submit Azure ML training job
        id: submit
        run: |
          RUN_NAME=$(az ml job create \
            -f aml/job-train.yml \
            --resource-group $AZ_RG \
            --workspace-name $AZ_WS \
            --set inputs.sample_fraction=${{ steps.vars.outputs.SAMPLE }} \
            --set inputs.learning_rate=${{ steps.vars.outputs.LR }} \
            --set inputs.num_leaves=${{ steps.vars.outputs.LEAVES }} \
            --stream \
            --query name -o tsv)
          echo "RUN_NAME=$RUN_NAME" >> $GITHUB_OUTPUT

      - name: Capture run metrics
        id: metrics
        run: |
          az ml job show -n ${{ steps.submit.outputs.RUN_NAME }} \
            --resource-group $AZ_RG --workspace-name $AZ_WS \
            --query "properties" -o json > run-props.json
          cat run-props.json

      - name: Upload run props
        uses: actions/upload-artifact@v4
        with:
          name: run-props-${{ steps.submit.outputs.RUN_NAME }}
          path: run-props.json

The permissions: id-token: write block is the one line everyone forgets. Without it the runner cannot mint an OIDC token at all and azure/login@v2 will fail with a Could not fetch access token error that looks unrelated. The environment: per-job line lets you wire GitHub Environments protection rules (required reviewers, wait timer) on the prod path while the pr path runs freely. The --stream flag on az ml job create is the difference between staring at a spinner for 40 minutes and seeing live log output in the Actions log.

Step 5. Pull metrics off the completed run with the SDK and gate registration.

Add a second job in the workflow that only runs on push to main, reads MLflow metrics, and only registers if the new model beats the existing tagged baseline. Drop this script at scripts/gate_register.py.

import argparse
import os
import sys
from azure.ai.ml import MLClient
from azure.ai.ml.entities import Model
from azure.ai.ml.constants import AssetTypes
from azure.identity import DefaultAzureCredential
import mlflow

p = argparse.ArgumentParser()
p.add_argument("--run-name", required=True)
p.add_argument("--rg", required=True)
p.add_argument("--ws", required=True)
p.add_argument("--sub", required=True)
p.add_argument("--model-name", default="credit-risk-lgbm")
p.add_argument("--min-auc-pr", type=float, default=0.78)
p.add_argument("--min-ks", type=float, default=0.42)
args = p.parse_args()

cred = DefaultAzureCredential()
ml = MLClient(cred, args.sub, args.rg, args.ws)

# Point MLflow at the workspace tracking URI
tracking_uri = ml.workspaces.get(args.ws).mlflow_tracking_uri
mlflow.set_tracking_uri(tracking_uri)

run = mlflow.get_run(args.run_name)
metrics = run.data.metrics
auc_pr = metrics.get("auc_pr", 0.0)
ks = metrics.get("ks_stat", 0.0)
auc_roc = metrics.get("auc_roc", 0.0)

print(f"new run: auc_pr={auc_pr:.4f}  ks_stat={ks:.4f}  auc_roc={auc_roc:.4f}")

# Read currently-deployed model's metrics from its tags
baseline = ml.models.get(args.model_name, label="latest")
b_auc_pr = float(baseline.tags.get("auc_pr", "0"))
b_ks = float(baseline.tags.get("ks_stat", "0"))
print(f"baseline v{baseline.version}: auc_pr={b_auc_pr:.4f}  ks_stat={b_ks:.4f}")

# Hard floor + must beat baseline by 0.5pp on auc_pr
if auc_pr < args.min_auc_pr or ks < args.min_ks:
    print("FAIL: new run below absolute floor. Not registering.")
    sys.exit(2)

if auc_pr < (b_auc_pr + 0.005):
    print("FAIL: new run did not beat baseline auc_pr by 0.5pp. Not registering.")
    sys.exit(3)

# Passed the gate. Register from the run's model output.
model_uri = f"azureml://jobs/{args.run_name}/outputs/model_dir"
new_model = Model(
    name=args.model_name,
    path=model_uri,
    type=AssetTypes.MLFLOW_MODEL,
    description=f"Promoted from run {args.run_name} via GH Actions OIDC",
    tags={
        "auc_pr": f"{auc_pr:.4f}",
        "ks_stat": f"{ks:.4f}",
        "auc_roc": f"{auc_roc:.4f}",
        "triggered_by": "github-actions",
        "source_run": args.run_name,
    },
)
registered = ml.models.create_or_update(new_model)
print(f"REGISTERED {registered.name} v{registered.version}")

The gate has three exits. Exit 2 means the run is below the absolute regulatory floor (the FCA-aligned KS floor of 0.42 the risk team agreed on). Exit 3 means the run is fine but does not improve on the baseline meaningfully. Only exit 0 registers a new version. That register call writes the metrics into the model's tags, which becomes the new baseline next week. The whole thing is self-perpetuating.

Step 6: Add the gate job to the workflow and require it before notifying release.

Append this to .github/workflows/retrain.yml under jobs:.

gate-and-register:
    needs: train
    if: github.event_name == 'push' && github.ref == 'refs/heads/main'
    runs-on: ubuntu-22.04
    permissions:
      id-token: write
      contents: read
    steps:
      - uses: actions/checkout@v4

      - uses: azure/setup-python@v5
        with:
          python-version: '3.11'

      - name: Install gate deps
        run: |
          pip install \
            azure-ai-ml==1.21.0 \
            azure-identity==1.19.0 \
            mlflow==2.16.2

      - name: Azure login via OIDC
        uses: azure/login@v2
        with:
          client-id: ${{ vars.AZURE_CLIENT_ID }}
          tenant-id: ${{ vars.AZURE_TENANT_ID }}
          subscription-id: ${{ vars.AZURE_SUBSCRIPTION_ID }}

      - name: Download run props
        uses: actions/download-artifact@v4
        with:
          name: run-props-${{ needs.train.outputs.run_name }}

      - name: Run metric gate
        run: |
          python scripts/gate_register.py \
            --run-name "${{ needs.train.outputs.run_name }}" \
            --rg "$AZ_RG" \
            --ws "$AZ_WS" \
            --sub "${{ vars.AZURE_SUBSCRIPTION_ID }}" \
            --min-auc-pr 0.78 \
            --min-ks 0.42

For this to actually receive needs.train.outputs.run_name you also need to declare the output on the train job:

train:
    runs-on: ubuntu-22.04
    outputs:
      run_name: ${{ steps.submit.outputs.RUN_NAME }}

Notice the if: guard. The gate-and-register job only fires on push to main, never on a PR. PRs run the smoke training but never touch the model registry. That separation is the whole point of branch-based triggers in MLOps.

Step 7. Pass workflow inputs through to the job for emergency manual retrains.

Sometimes the data team wants to sweep learning_rate by hand without opening a PR. The workflow_dispatch block in Step 4 already gives them that. From the GitHub UI, click Actions -> credit-risk-retrain -> Run workflow, fill in learning_rate=0.025 and num_leaves=255, hit Run. The vars.LR step picks the dispatched value up, threads it through to az ml job create --set inputs.learning_rate=0.025, and a new run lands in the workspace tagged triggered_by: github-actions. Same OIDC flow, no secrets.

If you need a real hyperparameter sweep rather than a single override, swap the command job for a sweep job in aml/job-sweep.yml and let Actions submit that instead. The federation, the gate, and the registration logic do not change.

Step 8. Test the PR smoke path end to end.

Open a throwaway branch, change a comment in src/train.py, push, open a PR against main. Within 30 seconds Actions kicks off, azure/login@v2 exchanges the OIDC token, az ml job create submits with sample_fraction=0.05, and you should see the streamed training log in the GitHub Actions log inside about three minutes. The smoke job will not register a model because the gate-and-register job is guarded by if: github.event_name == 'push'. That is correct. PRs prove the pipeline still runs. Merges prove the model is good enough to ship.

Troubleshooting

i> AADSTS70021: No matching federated identity record found for presented assertion. The subject claim on your federated credential does not match the actual OIDC token the runner is presenting. Run a one-off step echo "$ACTIONS_ID_TOKEN_REQUEST_TOKEN" is not the answer (that is the request token, not the assertion). Instead, look at the sub value the action logs when it fails. The most common mismatch is repo:org/repo:environment:prod versus repo:org/repo:ref:refs/heads/main, because you added an environment: line to the job after creating the credential. Either add a second federated credential for the environment subject or drop the environment from the job.

ii> Sweep job stuck in queued state forever. The compute target idle scale-down kicked in mid-sweep. Set min_instances: 1 on cpu-cluster-prod for the duration of the sweep, then put it back to 0 after. az ml compute update -n cpu-cluster-prod -g rg-ml-credit-prod-uks -w aml-credit-prod --min-instances 1 does it. Forgetting to reset min_instances back to 0 is how cost overruns happen; tag a calendar reminder.

iii> Could not fetch access token: getIDToken is not a function. You forgot the permissions: id-token: write block on the job. GitHub disables the OIDC token endpoint per-job by default. Add it at the job level, not the workflow level, otherwise it leaks to jobs that should not be minting tokens.

iv> The gate script fails with KeyError: 'auc_pr' even though MLflow auto-logging is on. LightGBM with mlflow.lightgbm.autolog() logs internal metric names like valid_0-binary_logloss, not your business metrics. You have to log auc_pr and ks_stat explicitly with mlflow.log_metric("auc_pr", value) inside train.py after you compute them on the held-out set. Autolog gets you the loss curve. Autolog does not get you the metrics the model risk committee asked for.

Cost of running this

UK South pricing, rough monthly figures for a weekly retrain plus 4-5 PR smoke runs a week.

• Azure ML workspace control plane: free.
• cpu-cluster-prod (Standard_D8s_v5, min 0 / max 4 nodes, ~6 hours full retrain + ~30 min of smoke per week): around £58 a month.
• ACR (Basic SKU acrmlcreditproduks): £4 a month.
• Blob storage for model artefacts and MLflow run data (~80 GB Hot): £1.60 a month.
• Key Vault kv-credit-prod-uks (handful of secrets, ~3000 ops/month): £0.50 a month.
• Application Insights attached to the workspace (low ingest, ~2 GB/month): £4.40 a month.
• GitHub Actions minutes on ubuntu-22.04 (about 35 minutes per push, 12 minutes per PR, 6 PRs + 4 pushes a week): free under the 2,000-minute GitHub Free tier for public repos, around £6 a month for a Team plan if private.

Total: roughly £70 to £85 a month for the full retrain pipeline, federation, registration, the lot.

Clean up

When you done playing with this:

az group delete \
  --name rg-ml-credit-prod-uks \
  --yes --no-wait

Then in GitHub: Settings -> Secrets and variables -> Actions -> Variables, delete AZURE_CLIENT_ID, AZURE_TENANT_ID, AZURE_SUBSCRIPTION_ID. And drop the federated credentials off the UAMI before you delete the resource group, so nothing dangles:

az identity federated-credential delete \
  --name gh-main-branch \
  --identity-name azureml-runner \
  --resource-group rg-ml-credit-prod-uks --yes

az identity federated-credential delete \
  --name gh-pull-request \
  --identity-name azureml-runner \
  --resource-group rg-ml-credit-prod-uks --yes

So that is the full picture. A weekly retrain that ships through GitHub Actions, authenticates via OIDC federation so the runner never holds an Azure secret, submits a real Azure ML CLI v2 job against a pinned environment, streams logs back to the Actions UI, captures metrics through MLflow, and gates the model registry on both an absolute regulatory floor and a relative beat-the-baseline check. The whole thing reproduces itself week after week. No more SSH into the dusty VM, no more pickle file copied by hand, no more six business days of bad scores before someone notices. Take this, swap in your own dataset, your own metric thresholds, your own UAMI, and you have something a mortgage-lending Fortune 500 actually runs in production. The bones are the same.

azure #azureml #githubactions #cicd #mlops #oidc #seniormlopsengineer