Milalib

Milalib is a library of helpers for Mila researchers to properly use the computing resources at their disposal.

All the listed functionality should work in slurm jobs on the Mila cluster. They may or may not work on DRAC clusters (I can report it doesn't work on Fir) -- you are welcome to try and report.

Monitoring GPU/resource usage

Monitoring resources is a little tricky for a few reasons:

• nvidia-smi's GPU utilization reporting is very loose-grained, leading to overestimation
• More precise profiling metrics such as sm_occupancy are exposed via DCGM
• CPU/RAM measurements require knowing what was allocated to your job, which is hidden deep in /sys/fs/cgroup.

Milalib does all the above for you.

You can run the monitor standalone with uvx. It outputs one measurement per line in JSON format.

# Monitor continuously, sampling every 3 seconds (default is 5 seconds)
uvx milalib monitor -i 3

# Get a single sample and quit
uvx milalib monitor -s

# List all available metrics
uvx milalib monitor -l

# Monitor some metrics specifically:
uvx milalib monitor -m sm_occupancy -m gpu_util

# Pipe to jq to only get the numeric value
uvx milalib monitor -m sm_occupancy | jq .value

# Only output metrics for the first GPU
uvx milalib monitor -d gpu0

# Only output GPU metrics
uvx milalib monitor -d 'gpu.*'

You may also install milalib with pip install milalib or uv add milalib.

Wandb integration

If you use wandb for experiment logging, you can also log the monitoring metrics:

from milalib.wandb import monitor

watcher = monitor(
    interval=5,
    metrics=["gpu_util", "power"],
    devices=["gpu0"],
    log=wandb.log,  # optional, defaults to wandb.log
)
...  # train
watcher.stop()

[!TIP] Run milalib monitor -l to view the list of available metrics.

Or, scoped to a block:

with monitor(interval=5):
    ...  # train

Interface

milalib exposes an async generator for the metrics, streaming objects of type MetricValue. Use it like this:

import asyncio

from milalib.monitor.poll import MetricRequest
from milalib.monitor.stream import stream_metrics


async def main():
    request = MetricRequest(metrics=["sm_occupancy"])
    async for value in stream_metrics(request, interval=0.1):
        print(value.timestamp, value.device, value.name, value.statistic, value.value)


asyncio.run(main())