Skip to main content

An MLflow plugin that automatically logs system-level metrics (CPU, memory, disk, GPU, OS) as run tags for improved experiment observability.

Project description

mlflow-sysmetrics

CI/CD Pipeline Coverage PyPI version Python versions License

๐Ÿง  A lightweight MLflow Run Context Provider that automatically logs system-level metrics (CPU, memory, disk, GPU, OS) as run tags.

โœ… Cross-platform ยท ๐Ÿ”Œ Plugin-ready ยท ๐Ÿงช Tested ยท โš™๏ธ Minimal dependencies


๐Ÿ“ฆ What it does

This plugin automatically adds system environment metadata to each MLflow run. It enables lightweight observability for experiment tracking โ€” useful in both local development and remote execution contexts.

โœ… Captured Tags

Tag Key Description
sys.cpu CPU model or architecture
sys.cpu_cores Logical CPU core count
sys.memory_gb Total system memory (GB)
sys.disk_free_gb Free disk space in current directory (GB)
sys.platform OS and kernel version
sys.gpu GPU name via nvidia-smi (Linux), system_profiler (macOS), or PowerShell (Windows)
sysmetrics.error Captures any exception during tagging

๐Ÿ’ก GPU detection is OS-aware:

  • macOS: via system_profiler
  • Linux: via nvidia-smi
  • Windows: via PowerShell (Get-CimInstance Win32_VideoController)

๐Ÿš€ Installation

poetry add mlflow-sysmetrics

Requires: Python โ‰ฅ 3.9 ยท mlflow โ‰ฅ 2.0 ยท psutil (automatically included)


๐Ÿ› ๏ธ Usage

Set the environment variable to activate the plugin:

export MLFLOW_RUN_CONTEXT_PROVIDER=sysmetrics

Then run any MLflow experiment:

import mlflow

with mlflow.start_run():
    mlflow.log_param("foo", "bar")
    # Plugin will automatically add sys.* tags

๐Ÿงช Testing

Run both unit and integration tests:

poetry run pytest -m unit
poetry run pytest -m integration

To manually verify plugin behavior:

export MLFLOW_RUN_CONTEXT_PROVIDER=sysmetrics
poetry run python scripts/debug_run.py

๐Ÿ“ท Example: Debug Script Output

You can verify system metrics manually with the debug script. Below is a sample output:

Example terminal output of sysmetrics plugin


๐Ÿ” Project Structure

mlflow-sysmetrics/
โ”œโ”€โ”€ src/mlflow_sysmetrics/
โ”‚   โ”œโ”€โ”€ __init__.py
โ”‚   โ”œโ”€โ”€ system_context.py         # Plugin implementation
โ”‚   โ””โ”€โ”€ utils/
โ”‚       โ”œโ”€โ”€ __init__.py
โ”‚       โ”œโ”€โ”€ constants.py          # Constants
โ”‚       โ”œโ”€โ”€ mac.py                # macOS-specific logic
โ”‚       โ””โ”€โ”€ windows.py            # Windows-specific GPU logic
โ”œโ”€โ”€ tests/
โ”‚   โ”œโ”€โ”€ unit/
โ”‚   โ”‚   โ”œโ”€โ”€ test_sysmetrics_context.py  # Core plugin logic
โ”‚   โ”‚   โ”œโ”€โ”€ test_mac_utils.py           # macOS GPU tests
โ”‚   โ”‚   โ””โ”€โ”€ test_windows_utils.py       # Windows GPU tests
โ”‚   โ””โ”€โ”€ integration/              # MLflow integration tests
โ”œโ”€โ”€ scripts/
โ”‚   โ””โ”€โ”€ debug_run.py              # Manual testing script
โ”œโ”€โ”€ assets/                       # Image and media assets
โ”‚   โ””โ”€โ”€ debug_run.png             # Screenshot of debug script
โ”œโ”€โ”€ pyproject.toml
โ”œโ”€โ”€ README.md
โ”œโ”€โ”€ LICENSE                       # Apache 2.0 License
โ””โ”€โ”€ CHANGELOG.md                  # Semantic changelog (Keep a Changelog format)

๐Ÿ“ฉ Plugin Registration

This plugin is exposed to MLflow via entry points:

[tool.poetry.plugins."mlflow.run_context_provider"]
sysmetrics = "mlflow_sysmetrics:SysMetricsRunContextProvider"

๐Ÿค Contributing

Pull requests, bug reports, and suggestions are welcome!

  1. Fork the repo
  2. Create a virtual environment: poetry install
  3. Write or update tests
  4. Run tests with poetry run pytest
  5. Submit your PR ๐Ÿš€

๐Ÿ“„ License

Apache License 2.0. See LICENSE for full terms.


๐Ÿ’ฌ Questions?

Feel free to open an issue or reach out via GitHub Discussions.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

mlflow_sysmetrics-0.1.3.tar.gz (9.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

mlflow_sysmetrics-0.1.3-py3-none-any.whl (11.8 kB view details)

Uploaded Python 3

File details

Details for the file mlflow_sysmetrics-0.1.3.tar.gz.

File metadata

  • Download URL: mlflow_sysmetrics-0.1.3.tar.gz
  • Upload date:
  • Size: 9.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/2.1.3 CPython/3.12.10 Linux/6.11.0-1014-azure

File hashes

Hashes for mlflow_sysmetrics-0.1.3.tar.gz
Algorithm Hash digest
SHA256 ea12a7410f38319b6ab0c9c83a0f6bf155d61074209e639aa20ef5ce8822fb06
MD5 d7b6bc454c6ff05ef2d17b106238822b
BLAKE2b-256 2a4d4c0673234a5c45674de8a88c991c5c7b99c5610fc79d42b8b9ddbd8d4303

See more details on using hashes here.

File details

Details for the file mlflow_sysmetrics-0.1.3-py3-none-any.whl.

File metadata

  • Download URL: mlflow_sysmetrics-0.1.3-py3-none-any.whl
  • Upload date:
  • Size: 11.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: poetry/2.1.3 CPython/3.12.10 Linux/6.11.0-1014-azure

File hashes

Hashes for mlflow_sysmetrics-0.1.3-py3-none-any.whl
Algorithm Hash digest
SHA256 20bbf86939a04bbc17e704f68cba6bedcaf19858a00597f3cd29306c5b9e1d15
MD5 59405bc1f0f4f6617a427c22be875d0e
BLAKE2b-256 3ed36b9723914145011d3932cf6755e725113adc0851649ec15894e6d3798c3e

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page