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

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 --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

This is currently a small JSON-lines stub with:

  • diagnose
  • get_fixes
  • doctor_summary

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.4.tar.gz (170.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.4-py3-none-any.whl (49.2 kB view details)

Uploaded Python 3

File details

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

File metadata

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

File hashes

Hashes for mlenvdoctor-0.1.4.tar.gz
Algorithm Hash digest
SHA256 1b08a009042eae466c65294d8526a664568ce87533474a66453ef24be4c2acbb
MD5 35d13b03c91f17975cb31943397b32ba
BLAKE2b-256 3a571f8b0f452303c7c74f1a8f0ab12d74e69d98a52f42a0c2e97e2f2f679e76

See more details on using hashes here.

File details

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

File metadata

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

File hashes

Hashes for mlenvdoctor-0.1.4-py3-none-any.whl
Algorithm Hash digest
SHA256 f5ee66635e55740005a6d4b10228e644e2548440a5a54f06b29e71fff104c6b7
MD5 8b127adc56664b6bdae9e7acf5c6ddfd
BLAKE2b-256 0064095af758d3a73f011b045b329ade12555d2c5af9ff66f095721d1f544a72

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