Skip to main content

Diagnose & fix ML environments for LLM fine-tuning

Project description

ML Environment Doctor

Python 3.8+ License: MIT PyPI CI PyPI - Downloads PyPI - Python Version codecov

Diagnose ML environment issues, export shareable reports, and generate safe starting fixes.

ML Environment Doctor is a Python CLI for inspecting machine learning environments used for local development, CI jobs, and GPU containers. It focuses on practical readiness checks for PyTorch, TensorFlow/Keras, and JAX/Flax workflows, then turns those results into readable reports and automatable outputs.

Repository: https://github.com/Dheena731/Ml-env-doctor

What It Does Today

  • Diagnoses Python runtime compatibility
  • Checks NVIDIA driver visibility with nvidia-smi
  • Validates PyTorch version, CUDA visibility, and a basic CUDA tensor execution path
  • Validates TensorFlow import, GPU enumeration, and a small tensor execution path
  • Validates JAX backend visibility, Flax presence, and a small JAX array execution path
  • Checks common ML training libraries like transformers, peft, trl, datasets, and accelerate
  • Optionally checks GPU memory, disk space, Docker GPU support, and Hugging Face connectivity
  • Exports JSON, CSV, and HTML reports
  • Generates requirements files, Conda environment files, and Dockerfiles
  • Provides a safe fix workflow with planning, dry-run, apply, and verification modes

What It Does Not Claim Yet

  • It does not fully repair every environment problem automatically
  • It does not maintain a large model registry
  • It does not replace framework-native setup guides for CUDA, cuDNN, or TPU drivers
  • Its Docker generation is configurable, but still template-based rather than a full build system

Quick Start

pip install mlenvdoctor

mlenvdoctor diagnose
mlenvdoctor diagnose --full
mlenvdoctor diagnose --json -
mlenvdoctor report
mlenvdoctor fix --plan
mlenvdoctor fix --dry-run
mlenvdoctor fix --apply --yes
mlenvdoctor fix --verify
mlenvdoctor dockerize tinyllama --stack minimal

Main Commands

Diagnose

Detailed evidence and export command.

mlenvdoctor diagnose
mlenvdoctor diagnose --full
mlenvdoctor diagnose --json diagnostics.json
mlenvdoctor diagnose --json -
mlenvdoctor diagnose --csv diagnostics.csv --html diagnostics.html

diagnose returns stable exit codes:

  • 0: no warnings or critical issues
  • 1: warnings present
  • 2: critical issues present

Doctor

Action-oriented triage command. It is intentionally different from diagnose:

  • doctor tells you what failed, the likely cause, the best next fix, and how to verify it
  • diagnose shows the full evidence and supports exports

Compact CI-friendly output:

mlenvdoctor doctor --ci
mlenvdoctor doctor --ci --full
mlenvdoctor doctor --guided

Report

Save a shareable JSON and HTML bundle:

mlenvdoctor report
mlenvdoctor report --quick --output-dir artifacts/mlenvdoctor

Fix

The fix workflow is intentionally explicit:

mlenvdoctor fix --plan
mlenvdoctor fix --dry-run
mlenvdoctor fix --apply
mlenvdoctor fix --apply --yes
mlenvdoctor fix --verify
mlenvdoctor fix --rollback
mlenvdoctor fix --venv --apply --yes
mlenvdoctor fix --conda
mlenvdoctor fix --stack llm-training --dry-run

Current fix behavior:

  • Plans file-generation and environment actions from detected issues
  • Labels each planned action by risk level
  • Supports dry-run mode before changes
  • Can create a virtual environment
  • Can generate either requirements-mlenvdoctor.txt or environment-mlenvdoctor.yml
  • Can install requirements when --apply is used
  • Supports explicit verification via mlenvdoctor fix --verify
  • Re-runs diagnostics for verification after a successful apply

Stack

mlenvdoctor stack llm-training
mlenvdoctor stack llm-training --output requirements-llm-training.txt

Dockerize

mlenvdoctor dockerize tinyllama
mlenvdoctor dockerize mistral-7b --stack llm-training
mlenvdoctor dockerize --service --base-image nvidia/cuda:12.4.0-runtime-ubuntu22.04
mlenvdoctor dockerize gpt2 --python-version 3.10 -o Dockerfile.gpt2

Current Docker generation supports:

  • model-aware defaults for tinyllama, gpt2, and mistral-7b
  • stack selection
  • base image override
  • Python version selection
  • training and FastAPI service profiles

MCP

mlenvdoctor mcp serve

MCP now exposes a JSON-lines interface over stdin/stdout.

Core integration tools:

  • diagnose
  • get_fixes
  • doctor_summary
  • report_bundle
  • diagnose_environment (v1 alias)
  • get_fix_plan (v1 alias)
  • verify_fix (v1 alias)
  • export_report (v1 alias)
  • list_tools
  • tool_schema
  • health

Request format:

{"tool":"list_tools","arguments":{}}
{"tool":"tool_schema","arguments":{"name":"report_bundle"}}

Compatibility policy:

  • MCP schemas evolve additively.
  • v1 alias names are kept for migration safety.
  • Use tool_schema to discover canonical tool names.

Output Schema

JSON exports include:

  • issue name, status, severity, and fix command
  • check ID and category
  • recommendation, likely cause, and verify steps
  • confidence
  • evidence and metadata blocks
  • runtime context (platform, detected backend, NVIDIA tooling visibility)
  • summary counts and exit code

That makes the tool more useful for CI parsing, dashboards, and future integrations.

Development

git clone https://github.com/Dheena731/Ml-env-doctor.git
cd Ml-env-doctor
python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"
python -m black src tests
python -m ruff check src tests
python -m pytest
python -m mypy src

CI/CD and Release Flow

  • CI runs formatting (black --check), lint (ruff), tests (pytest), and package build checks (python -m build, twine check) on every push/PR.
  • The matrix tests Linux (3.8-3.11), macOS (3.11), and Windows (3.11).
  • A tagged GitHub release triggers the publish job to upload built artifacts to PyPI.
  • If CI fails locally vs. GitHub, run these exact commands first:
python -m black src tests
python -m ruff check src tests
python -m pytest

Test Strategy

The repository uses two test layers:

  • direct Python and Typer CliRunner tests for most command behavior
  • a smaller set of module-level tests for exports, validators, diagnostics, and file generation

This keeps the main test suite independent from whether mlenvdoctor was installed as a shell command in the current environment.

Repository Notes

License

This project is licensed under the MIT License. See LICENSE.

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

mlenvdoctor-0.1.5.tar.gz (176.8 kB view details)

Uploaded Source

Built Distribution

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

mlenvdoctor-0.1.5-py3-none-any.whl (52.6 kB view details)

Uploaded Python 3

File details

Details for the file mlenvdoctor-0.1.5.tar.gz.

File metadata

  • Download URL: mlenvdoctor-0.1.5.tar.gz
  • Upload date:
  • Size: 176.8 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.13

File hashes

Hashes for mlenvdoctor-0.1.5.tar.gz
Algorithm Hash digest
SHA256 c19a5eebf915826f727d8e0eee6f22995907f947aa6d6cafb478d2d0593a9e37
MD5 271a1275f2181e6871ddac81e1b9096e
BLAKE2b-256 fb154b0db22baf06f1e03a46139cd11330b13e40da39aa58a87effb0ed3f21cc

See more details on using hashes here.

File details

Details for the file mlenvdoctor-0.1.5-py3-none-any.whl.

File metadata

  • Download URL: mlenvdoctor-0.1.5-py3-none-any.whl
  • Upload date:
  • Size: 52.6 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/6.2.0 CPython/3.13.13

File hashes

Hashes for mlenvdoctor-0.1.5-py3-none-any.whl
Algorithm Hash digest
SHA256 e83bc89f7d6631ac818464eedae1089792f9e33191639fa9009489b0e345a146
MD5 ccb4c7633c1d7093582a5450faf29a14
BLAKE2b-256 70b1fc69e8795b5dcc68e66405ef8b3d0735686230c393ccbb8cd400f7f009eb

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