Diagnose & fix ML environments for LLM fine-tuning
Project description
ML Environment Doctor
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, andaccelerate - 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 issues1: warnings present2: critical issues present
Doctor
Action-oriented triage command. It is intentionally different from diagnose:
doctortells you what failed, the likely cause, the best next fix, and how to verify itdiagnoseshows 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.txtorenvironment-mlenvdoctor.yml - Can install requirements when
--applyis 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, andmistral-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:
diagnoseget_fixesdoctor_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
CliRunnertests 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
- CONTRIBUTING.md covers local setup and contribution flow
- IMPROVEMENTS.md tracks active improvement themes
- IMPROVEMENTS_ROADMAP.md outlines future milestones
- docs/PROJECT_OVERVIEW.md explains the current project architecture and product shape
- docs/KILLER_PRODUCT_PLAN.md is the strategic roadmap for turning the project into a category-leading product
- docs/PHASED_EXECUTION_PLAN.md breaks the strategy into execution phases and current build focus
- docker/README.md documents Docker-specific workflows
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
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
1b08a009042eae466c65294d8526a664568ce87533474a66453ef24be4c2acbb
|
|
| MD5 |
35d13b03c91f17975cb31943397b32ba
|
|
| BLAKE2b-256 |
3a571f8b0f452303c7c74f1a8f0ab12d74e69d98a52f42a0c2e97e2f2f679e76
|
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
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
f5ee66635e55740005a6d4b10228e644e2548440a5a54f06b29e71fff104c6b7
|
|
| MD5 |
8b127adc56664b6bdae9e7acf5c6ddfd
|
|
| BLAKE2b-256 |
0064095af758d3a73f011b045b329ade12555d2c5af9ff66f095721d1f544a72
|