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

mlenvdoctor/__init__.py

@@ -1,3 +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 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,
 )
 
@@ -37,9 +41,21 @@ def main(
         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()
@@ -47,9 +63,18 @@ def diagnose(
     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.
@@ -57,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]")
@@ -78,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.
@@ -87,7 +123,7 @@ 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]")
 
 
@@ -101,8 +137,8 @@ def dockerize(
         "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.
@@ -115,45 +151,45 @@ 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(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")
+    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]"
+            f"[bold red]{icon_cross()} Smoke test failed. Run 'mlenvdoctor diagnose' for details.[/bold red]"
         )
         raise typer.Exit(1)
 

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"

mlenvdoctor/diagnose.py

@@ -11,6 +11,7 @@ except ImportError:
 
 from rich.table import Table
 
+from .icons import icon_check, icon_cross, icon_info, icon_search, icon_warning
 from .utils import check_command_exists, console, format_size, get_home_config_dir, run_command
 
 
@@ -33,12 +34,13 @@ class DiagnosticIssue:
 
     def to_row(self) -> Tuple[str, str, str, str]:
         """Convert to table row."""
-        status_icon = {
-            "PASS": "✅",
-            "FAIL": "❌",
-            "WARN": "⚠️",
-            "INFO": "ℹ️",
-        }.get(self.status.split()[0], "❓")
+        status_icon_map = {
+            "PASS": icon_check(),
+            "FAIL": icon_cross(),
+            "WARN": icon_warning(),
+            "INFO": icon_info(),
+        }
+        status_icon = status_icon_map.get(self.status.split()[0], "?")
         return (
             self.name,
             f"{status_icon} {self.status}",
@@ -411,11 +413,19 @@ def check_docker_gpu() -> List[DiagnosticIssue]:
 
 def check_internet_connectivity() -> List[DiagnosticIssue]:
     """Check internet connectivity for HF Hub."""
+    from .retry import retry_network
+
     issues = []
-    try:
+
+    @retry_network
+    def _check_connectivity() -> bool:
         import urllib.request
 
         urllib.request.urlopen("https://huggingface.co", timeout=5)
+        return True
+
+    try:
+        _check_connectivity()
         issues.append(
             DiagnosticIssue(
                 name="Internet Connectivity",
@@ -424,36 +434,117 @@ def check_internet_connectivity() -> List[DiagnosticIssue]:
                 fix="",
             )
         )
-    except Exception:
+    except Exception as e:
         issues.append(
             DiagnosticIssue(
                 name="Internet Connectivity",
                 status="WARN - Cannot reach HF Hub",
                 severity="warning",
                 fix="Check internet connection and firewall settings",
+                details=str(e),
             )
         )
 
     return issues
 
 
-def diagnose_env(full: bool = False) -> List[DiagnosticIssue]:
-    """Run all diagnostic checks."""
-    all_issues: List[DiagnosticIssue] = []
+def diagnose_env(full: bool = False, parallel: bool = True) -> List[DiagnosticIssue]:
+    """
+    Run all diagnostic checks.
+
+    Args:
+        full: Whether to run full diagnostics including extended checks
+        parallel: Whether to run independent checks in parallel
 
-    console.print("[bold blue]🔍 Running ML Environment Diagnostics...[/bold blue]\n")
+    Returns:
+        List of diagnostic issues found
+    """
+    from .parallel import run_parallel_with_results
 
-    # Core checks (always run)
-    all_issues.extend(check_cuda_driver())
-    all_issues.extend(check_pytorch_cuda())
-    all_issues.extend(check_ml_libraries())
+    all_issues: List[DiagnosticIssue] = []
 
-    # Extended checks (if --full)
+    console.print(f"[bold blue]{icon_search()} Running ML Environment Diagnostics...[/bold blue]\n")
+
+    # Core checks (always run) - these can run in parallel
+    core_checks = [
+        check_cuda_driver,
+        check_pytorch_cuda,
+        check_ml_libraries,
+    ]
+
+    if parallel:
+        # Run core checks in parallel
+        results = run_parallel_with_results(
+            lambda check_func: check_func(),
+            core_checks,
+            max_workers=3,
+            timeout=60.0,
+        )
+        for check_func, result in results:
+            if isinstance(result, Exception):
+                # Log error but continue with other checks
+                from .logger import logger
+                logger.error(f"Check {check_func.__name__} failed: {result}")
+                # Add a diagnostic issue for the failure
+                all_issues.append(
+                    DiagnosticIssue(
+                        name=check_func.__name__.replace("check_", "").replace("_", " ").title(),
+                        status="FAIL - Check error",
+                        severity="critical",
+                        fix="Run diagnostics again or check logs",
+                        details=str(result),
+                    )
+                )
+            else:
+                all_issues.extend(result)
+    else:
+        # Sequential execution (fallback)
+        for check_func in core_checks:
+            try:
+                all_issues.extend(check_func())
+            except Exception as e:
+                from .logger import logger
+                logger.error(f"Check {check_func.__name__} failed: {e}")
+                all_issues.append(
+                    DiagnosticIssue(
+                        name=check_func.__name__.replace("check_", "").replace("_", " ").title(),
+                        status="FAIL - Check error",
+                        severity="critical",
+                        fix="Run diagnostics again or check logs",
+                        details=str(e),
+                    )
+                )
+
+    # Extended checks (if --full) - can also run in parallel
     if full:
-        all_issues.extend(check_gpu_memory())
-        all_issues.extend(check_disk_space())
-        all_issues.extend(check_docker_gpu())
-        all_issues.extend(check_internet_connectivity())
+        extended_checks = [
+            check_gpu_memory,
+            check_disk_space,
+            check_docker_gpu,
+            check_internet_connectivity,
+        ]
+
+        if parallel:
+            results = run_parallel_with_results(
+                lambda check_func: check_func(),
+                extended_checks,
+                max_workers=4,
+                timeout=120.0,
+            )
+            for check_func, result in results:
+                if isinstance(result, Exception):
+                    from .logger import logger
+                    logger.warning(f"Extended check {check_func.__name__} failed: {result}")
+                    # Extended checks are less critical, so we log but don't fail
+                else:
+                    all_issues.extend(result)
+        else:
+            for check_func in extended_checks:
+                try:
+                    all_issues.extend(check_func())
+                except Exception as e:
+                    from .logger import logger
+                    logger.warning(f"Extended check {check_func.__name__} failed: {e}")
 
     return all_issues
 
@@ -484,18 +575,18 @@ def print_diagnostic_table(issues: List[DiagnosticIssue]) -> None:
     pass_count = sum(1 for i in issues if "PASS" in i.status)
 
     console.print()
-    console.print(f"[green]✅ Passed: {pass_count}[/green]")
+    console.print(f"[green]{icon_check()} Passed: {pass_count}[/green]")
     if warning_count > 0:
-        console.print(f"[yellow]⚠️  Warnings: {warning_count}[/yellow]")
+        console.print(f"[yellow]{icon_warning()}  Warnings: {warning_count}[/yellow]")
     if critical_count > 0:
-        console.print(f"[red]❌ Critical Issues: {critical_count}[/red]")
+        console.print(f"[red]{icon_cross()} Critical Issues: {critical_count}[/red]")
 
     if critical_count == 0 and warning_count == 0:
         console.print(
             "\n[bold green]🎉 Your ML environment looks ready for fine-tuning![/bold green]"
         )
     elif critical_count > 0:
-        console.print("\n[bold red]⚠️  Please fix critical issues before proceeding.[/bold red]")
+        console.print(f"\n[bold red]{icon_warning()}  Please fix critical issues before proceeding.[/bold red]")
     else:
         console.print(
             "\n[bold yellow]💡 Consider addressing warnings for optimal performance.[/bold yellow]"

mlenvdoctor/exceptions.py

@@ -0,0 +1,51 @@
+"""Custom exceptions for ML Environment Doctor."""
+
+
+class MLEnvDoctorError(Exception):
+    """Base exception for ML Environment Doctor."""
+
+    def __init__(self, message: str, suggestion: str = ""):
+        super().__init__(message)
+        self.message = message
+        self.suggestion = suggestion
+
+    def __str__(self) -> str:
+        if self.suggestion:
+            return f"{self.message}\n💡 Suggestion: {self.suggestion}"
+        return self.message
+
+
+class DiagnosticError(MLEnvDoctorError):
+    """Error during diagnostic checks."""
+
+    pass
+
+
+class FixError(MLEnvDoctorError):
+    """Error during auto-fix operations."""
+
+    pass
+
+
+class DockerError(MLEnvDoctorError):
+    """Error during Docker operations."""
+
+    pass
+
+
+class GPUError(MLEnvDoctorError):
+    """Error related to GPU operations."""
+
+    pass
+
+
+class ConfigurationError(MLEnvDoctorError):
+    """Error in configuration."""
+
+    pass
+
+
+class InstallationError(MLEnvDoctorError):
+    """Error during package installation."""
+
+    pass