Skip to content

Experiment Tracking (MLflow)

PAIS runs MLflow 3.x with PostgreSQL backend and VAST artifact storage. All experiments are persistent and accessible to your entire research group.

Access

UI: https://mlflow.pais.auckland.ac.nz — log in with Tuakiri SSO

Python: 

import mlflow
import os

mlflow.set_tracking_uri("https://mlflow.pais.auckland.ac.nz")

Key features on PAIS

Feature Notes
Experiment tracking Log metrics, parameters, tags per run
Artifact storage Models and files stored on VAST (persistent)
Model registry Register, version, and stage models
Team access All group members see the same experiments
OIDC auth Single sign-on via Tuakiri

Auto-logging with PyTorch Lightning

import mlflow
from pytorch_lightning.loggers import MLFlowLogger

mlflow.set_tracking_uri("https://mlflow.pais.auckland.ac.nz")

logger = MLFlowLogger(
    experiment_name="my_experiment",
    tracking_uri="https://mlflow.pais.auckland.ac.nz",
    run_name="run-1",
)

trainer = pl.Trainer(logger=logger, max_epochs=10)

Auto-logging with HuggingFace Transformers

import mlflow
mlflow.set_tracking_uri("https://mlflow.pais.auckland.ac.nz")
mlflow.transformers.autolog()

# All training_args metrics are logged automatically
trainer = Trainer(training_args, model=model, ...)
trainer.train()

See also