mlenvdoctor 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

mlenvdoctor/__init__.py

@@ -1,4 +1,18 @@
 """ML Environment Doctor - Diagnose and fix ML environments for LLM fine-tuning."""
 
-__version__ = "0.1.0"
+__version__ = "0.1.2"
 
+__all__ = [
+    "__version__",
+    "diagnose",
+    "fix",
+    "dockerize",
+    "export",
+    "exceptions",
+    "logger",
+    "config",
+    "validators",
+    "retry",
+    "parallel",
+    "constants",
+]

mlenvdoctor/cli.py

@@ -1,5 +1,6 @@
 """CLI entrypoint for ML Environment Doctor."""
 
+from pathlib import Path
 from typing import Optional
 
 import typer
@@ -7,13 +8,16 @@ import typer
 from . import __version__
 from .diagnose import diagnose_env, print_diagnostic_table
 from .dockerize import generate_dockerfile, generate_service_template
+from .export import export_csv, export_html, export_json
 from .fix import auto_fix
-from .gpu import benchmark_gpu_ops, smoke_test_lora, test_model
+from .gpu import benchmark_gpu_ops, smoke_test_lora, test_model as gpu_test_model
+from .icons import icon_check, icon_cross, icon_search, icon_test, icon_whale, icon_wrench
+from .logger import get_default_log_file, setup_logger
 from .utils import console
 
 app = typer.Typer(
     name="mlenvdoctor",
-    help="🔍 ML Environment Doctor - Diagnose & fix ML environments for LLM fine-tuning",
+    help=f"{icon_search()} ML Environment Doctor - Diagnose & fix ML environments for LLM fine-tuning",
     add_completion=False,
 )
 
@@ -21,26 +25,56 @@ app = typer.Typer(
 def version_callback(value: bool):
     """Print version and exit."""
     if value:
-        console.print(f"[bold blue]ML Environment Doctor[/bold blue] version [cyan]{__version__}[/cyan]")
+        console.print(
+            f"[bold blue]ML Environment Doctor[/bold blue] version [cyan]{__version__}[/cyan]"
+        )
         raise typer.Exit()
 
 
 @app.callback()
 def main(
     version: Optional[bool] = typer.Option(
-        None, "--version", "-v", callback=version_callback, is_eager=True, help="Show version and exit"
+        None,
+        "--version",
+        "-v",
+        callback=version_callback,
+        is_eager=True,
+        help="Show version and exit",
+    ),
+    log_file: Optional[Path] = typer.Option(
+        None,
+        "--log-file",
+        help="Path to log file (default: ~/.mlenvdoctor/logs/mlenvdoctor.log)",
+    ),
+    log_level: str = typer.Option(
+        "INFO",
+        "--log-level",
+        help="Logging level: DEBUG, INFO, WARNING, ERROR, CRITICAL",
     ),
 ):
     """ML Environment Doctor - Diagnose & fix ML environments for LLM fine-tuning."""
-    pass
+    # Set up logging
+    log_path = log_file or get_default_log_file()
+    setup_logger(log_file=log_path, level=log_level)
 
 
 @app.command()
 def diagnose(
-    full: bool = typer.Option(False, "--full", "-f", help="Run full diagnostics including GPU benchmarks"),
+    full: bool = typer.Option(
+        False, "--full", "-f", help="Run full diagnostics including GPU benchmarks"
+    ),
+    json_output: Optional[Path] = typer.Option(
+        None, "--json", help="Export results to JSON file"
+    ),
+    csv_output: Optional[Path] = typer.Option(
+        None, "--csv", help="Export results to CSV file"
+    ),
+    html_output: Optional[Path] = typer.Option(
+        None, "--html", help="Export results to HTML file"
+    ),
 ):
-    """
-    🔍 Diagnose your ML environment.
+    f"""
+    {icon_search()} Diagnose your ML environment.
 
     Quick scan: Checks CUDA, PyTorch, and required ML libraries.
     Full scan (--full): Also checks GPU memory, disk space, Docker GPU support, and connectivity.
@@ -48,6 +82,17 @@ def diagnose(
     issues = diagnose_env(full=full)
     print_diagnostic_table(issues)
 
+    # Export to formats if requested
+    if json_output:
+        export_json(issues, json_output)
+        console.print(f"[green]{icon_check()} Exported to {json_output}[/green]")
+    if csv_output:
+        export_csv(issues, csv_output)
+        console.print(f"[green]{icon_check()} Exported to {csv_output}[/green]")
+    if html_output:
+        export_html(issues, html_output)
+        console.print(f"[green]{icon_check()} Exported to {html_output}[/green]")
+
     if full:
         console.print()
         console.print("[bold blue]Running GPU benchmark...[/bold blue]")
@@ -69,8 +114,8 @@ def fix(
     venv: bool = typer.Option(False, "--venv", "-v", help="Create virtual environment"),
     stack: str = typer.Option("trl-peft", "--stack", "-s", help="ML stack: trl-peft or minimal"),
 ):
-    """
-    🔧 Auto-fix environment issues and generate requirements.
+    f"""
+    {icon_wrench()} Auto-fix environment issues and generate requirements.
 
     Generates requirements.txt or conda environment file based on detected issues.
     Optionally creates a virtual environment and installs dependencies.
@@ -78,18 +123,22 @@ def fix(
     success = auto_fix(use_conda=conda, create_venv=venv, stack=stack)
     if success:
         console.print()
-        console.print("[bold green]✅ Auto-fix completed![/bold green]")
+        console.print(f"[bold green]{icon_check()} Auto-fix completed![/bold green]")
         console.print("[yellow]💡 Run 'mlenvdoctor diagnose' to verify fixes[/yellow]")
 
 
 @app.command()
 def dockerize(
     model: Optional[str] = typer.Argument(None, help="Model name (mistral-7b, tinyllama, gpt2)"),
-    service: bool = typer.Option(False, "--service", "-s", help="Generate FastAPI service template"),
-    output: str = typer.Option("Dockerfile.mlenvdoctor", "--output", "-o", help="Output Dockerfile name"),
+    service: bool = typer.Option(
+        False, "--service", "-s", help="Generate FastAPI service template"
+    ),
+    output: str = typer.Option(
+        "Dockerfile.mlenvdoctor", "--output", "-o", help="Output Dockerfile name"
+    ),
 ):
-    """
-    🐳 Generate Dockerfile for ML fine-tuning.
+    f"""
+    {icon_whale()} Generate Dockerfile for ML fine-tuning.
 
     Creates a production-ready Dockerfile with CUDA support.
     Optionally generates a FastAPI service template.
@@ -102,44 +151,46 @@ def dockerize(
         generate_dockerfile(model_name=model, service=service, output_file=output)
 
     console.print()
-    console.print("[bold green]✅ Dockerfile generated![/bold green]")
+    console.print(f"[bold green]{icon_check()} Dockerfile generated![/bold green]")
 
 
-@app.command()
-def test_model(
+@app.command(name="test-model")
+def test_model_cmd(
     model: str = typer.Argument("tinyllama", help="Model to test (tinyllama, gpt2, mistral-7b)"),
 ):
-    """
-    🧪 Run smoke test with a real LLM model.
+    f"""
+    {icon_test()} Run smoke test with a real LLM model.
 
     Tests model loading and forward pass to verify fine-tuning readiness.
     """
-    console.print(f"[bold blue]🧪 Testing model: {model}[/bold blue]\n")
-    success = test_model(model_name=model)
+    console.print(f"[bold blue]{icon_test()} Testing model: {model}[/bold blue]\n")
+    success = gpu_test_model(model_name=model)
     if success:
         console.print()
-        console.print("[bold green]✅ Model test passed! Ready for fine-tuning.[/bold green]")
+        console.print(f"[bold green]{icon_check()} Model test passed! Ready for fine-tuning.[/bold green]")
     else:
         console.print()
-        console.print("[bold red]❌ Model test failed. Check diagnostics.[/bold red]")
+        console.print(f"[bold red]{icon_cross()} Model test failed. Check diagnostics.[/bold red]")
         raise typer.Exit(1)
 
 
 @app.command()
 def smoke_test():
-    """
-    🧪 Run LoRA fine-tuning smoke test.
+    f"""
+    {icon_test()} Run LoRA fine-tuning smoke test.
 
     Performs a minimal LoRA fine-tuning test to verify environment setup.
     """
-    console.print("[bold blue]🧪 Running LoRA smoke test...[/bold blue]\n")
+    console.print(f"[bold blue]{icon_test()} Running LoRA smoke test...[/bold blue]\n")
     success = smoke_test_lora()
     if success:
         console.print()
-        console.print("[bold green]✅ Smoke test passed! Environment is ready.[/bold green]")
+        console.print(f"[bold green]{icon_check()} Smoke test passed! Environment is ready.[/bold green]")
     else:
         console.print()
-        console.print("[bold red]❌ Smoke test failed. Run 'mlenvdoctor diagnose' for details.[/bold red]")
+        console.print(
+            f"[bold red]{icon_cross()} Smoke test failed. Run 'mlenvdoctor diagnose' for details.[/bold red]"
+        )
         raise typer.Exit(1)
 
 
@@ -150,4 +201,3 @@ def main_cli():
 
 if __name__ == "__main__":
     main_cli()
-

mlenvdoctor/config.py

@@ -0,0 +1,169 @@
+"""Configuration management for ML Environment Doctor."""
+
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+# Try tomllib (Python 3.11+)
+try:
+    import tomllib
+except ImportError:
+    tomllib = None  # type: ignore
+
+# Fallback to tomli for older Python versions
+try:
+    import tomli
+except ImportError:
+    tomli = None
+
+from .exceptions import ConfigurationError
+from .utils import get_home_config_dir
+
+
+def load_config(config_path: Optional[Path] = None) -> Dict[str, Any]:
+    """
+    Load configuration from TOML file.
+
+    Args:
+        config_path: Path to config file. If None, searches for:
+            1. mlenvdoctor.toml in current directory
+            2. .mlenvdoctorrc in current directory
+            3. ~/.mlenvdoctor/config.toml
+
+    Returns:
+        Configuration dictionary
+
+    Raises:
+        ConfigurationError: If config file is invalid
+    """
+    default_config: Dict[str, Any] = {
+        "diagnostics": {
+            "full_scan": False,
+            "skip_checks": [],
+        },
+        "fix": {
+            "default_stack": "trl-peft",
+            "auto_install": False,
+        },
+        "docker": {
+            "default_base_image": "nvidia/cuda:12.4.0-devel-ubuntu22.04",
+        },
+        "logging": {
+            "level": "INFO",
+            "file": None,
+        },
+    }
+
+    if config_path is None:
+        # Search for config files
+        search_paths = [
+            Path("mlenvdoctor.toml"),
+            Path(".mlenvdoctorrc"),
+            get_home_config_dir() / "config.toml",
+        ]
+
+        for path in search_paths:
+            if path.exists():
+                config_path = path
+                break
+    else:
+        if not config_path.exists():
+            raise ConfigurationError(
+                f"Config file not found: {config_path}",
+                "Create the file or use default configuration",
+            )
+
+    if config_path is None or not config_path.exists():
+        return default_config
+
+    try:
+        # Try tomllib (Python 3.11+)
+        if tomllib is not None:
+            with config_path.open("rb") as f:
+                user_config = tomllib.load(f)
+        elif tomli is not None:
+            # Fallback to tomli for older Python
+            with config_path.open("rb") as f:
+                user_config = tomli.load(f)
+        else:
+            raise ConfigurationError(
+                "TOML parsing not available. Install tomli: pip install tomli",
+                "Or upgrade to Python 3.11+",
+            )
+
+        # Merge with defaults
+        merged_config = default_config.copy()
+        for section, values in user_config.items():
+            if section in merged_config and isinstance(merged_config[section], dict):
+                merged_config[section].update(values)
+            else:
+                merged_config[section] = values
+
+        return merged_config
+
+    except Exception as e:
+        raise ConfigurationError(
+            f"Error parsing config file {config_path}: {e}",
+            "Check TOML syntax and file permissions",
+        ) from e
+
+
+def get_config_value(config: Dict[str, Any], *keys: str, default: Any = None) -> Any:
+    """
+    Get nested config value safely.
+
+    Args:
+        config: Configuration dictionary
+        *keys: Nested keys to traverse
+        default: Default value if key not found
+
+    Returns:
+        Config value or default
+    """
+    value = config
+    for key in keys:
+        if isinstance(value, dict):
+            value = value.get(key)
+            if value is None:
+                return default
+        else:
+            return default
+    return value if value is not None else default
+
+
+def create_default_config(output_path: Path) -> Path:
+    """
+    Create a default configuration file.
+
+    Args:
+        output_path: Path where to create config file
+
+    Returns:
+        Path to created config file
+    """
+    default_content = """# ML Environment Doctor Configuration
+
+[diagnostics]
+# Run full scan by default
+full_scan = false
+# Skip specific checks (e.g., ["docker_gpu", "internet"])
+skip_checks = []
+
+[fix]
+# Default ML stack: "trl-peft" or "minimal"
+default_stack = "trl-peft"
+# Automatically install dependencies without prompting
+auto_install = false
+
+[docker]
+# Default base image for Dockerfiles
+default_base_image = "nvidia/cuda:12.4.0-devel-ubuntu22.04"
+
+[logging]
+# Logging level: DEBUG, INFO, WARNING, ERROR, CRITICAL
+level = "INFO"
+# Log file path (None for default: ~/.mlenvdoctor/logs/mlenvdoctor.log)
+file = null
+"""
+
+    output_path.write_text(default_content, encoding="utf-8")
+    return output_path

mlenvdoctor/constants.py

@@ -0,0 +1,63 @@
+"""Constants used throughout ML Environment Doctor."""
+
+from typing import Final
+
+# Version compatibility
+MIN_PYTHON_VERSION: Final[tuple[int, int]] = (3, 8)
+MIN_PYTORCH_VERSION: Final[str] = "2.4.0"
+
+# CUDA versions
+SUPPORTED_CUDA_VERSIONS: Final[list[str]] = ["12.1", "12.4"]
+DEFAULT_CUDA_VERSION: Final[str] = "12.4"
+
+# ML Library versions
+MIN_TRANSFORMERS_VERSION: Final[str] = "4.44.0"
+MIN_PEFT_VERSION: Final[str] = "0.12.0"
+MIN_TRL_VERSION: Final[str] = "0.9.0"
+MIN_DATASETS_VERSION: Final[str] = "2.20.0"
+MIN_ACCELERATE_VERSION: Final[str] = "1.0.0"
+
+# Memory requirements (GB)
+MIN_GPU_MEMORY_GB: Final[int] = 8
+RECOMMENDED_GPU_MEMORY_GB: Final[int] = 16
+MIN_DISK_SPACE_GB: Final[int] = 50
+
+# Timeouts (seconds)
+DEFAULT_COMMAND_TIMEOUT: Final[int] = 30
+DEFAULT_NETWORK_TIMEOUT: Final[int] = 10
+DEFAULT_INSTALL_TIMEOUT: Final[int] = 600
+
+# File paths
+DEFAULT_CONFIG_FILE: Final[str] = "mlenvdoctor.toml"
+DEFAULT_REQUIREMENTS_FILE: Final[str] = "requirements-mlenvdoctor.txt"
+DEFAULT_DOCKERFILE: Final[str] = "Dockerfile.mlenvdoctor"
+
+# Model names
+SUPPORTED_MODELS: Final[dict[str, str]] = {
+    "tinyllama": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
+    "gpt2": "gpt2",
+    "mistral-7b": "mistralai/Mistral-7B-v0.1",
+}
+
+# ML Stacks
+ML_STACKS: Final[list[str]] = ["trl-peft", "minimal"]
+
+# Diagnostic check names
+CHECK_CUDA_DRIVER: Final[str] = "cuda_driver"
+CHECK_PYTORCH_CUDA: Final[str] = "pytorch_cuda"
+CHECK_ML_LIBRARIES: Final[str] = "ml_libraries"
+CHECK_GPU_MEMORY: Final[str] = "gpu_memory"
+CHECK_DISK_SPACE: Final[str] = "disk_space"
+CHECK_DOCKER_GPU: Final[str] = "docker_gpu"
+CHECK_INTERNET: Final[str] = "internet"
+
+# Severity levels
+SEVERITY_CRITICAL: Final[str] = "critical"
+SEVERITY_WARNING: Final[str] = "warning"
+SEVERITY_INFO: Final[str] = "info"
+
+# Status values
+STATUS_PASS: Final[str] = "PASS"
+STATUS_FAIL: Final[str] = "FAIL"
+STATUS_WARN: Final[str] = "WARN"
+STATUS_INFO: Final[str] = "INFO"