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

mlenvdoctor/diagnose.py

@@ -2,8 +2,7 @@
 
 import importlib
 import re
-import subprocess
-from typing import Dict, List, Optional, Tuple
+from typing import List, Optional, Tuple
 
 try:
     import torch
@@ -12,16 +11,8 @@ except ImportError:
 
 from rich.table import Table
 
-from .utils import (
-    check_command_exists,
-    console,
-    format_size,
-    get_home_config_dir,
-    print_error,
-    print_info,
-    print_warning,
-    run_command,
-)
+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
 
 
 class DiagnosticIssue:
@@ -43,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}",
@@ -164,7 +156,7 @@ def check_pytorch_cuda() -> List[DiagnosticIssue]:
                         name="PyTorch Version",
                         status="WARN - Old version",
                         severity="warning",
-                        fix=f"Upgrade: pip install torch>=2.4.0 --index-url https://download.pytorch.org/whl/cu124",
+                        fix="Upgrade: pip install torch>=2.4.0 --index-url https://download.pytorch.org/whl/cu124",
                         details=f"Current: {torch_version}, Recommended: >=2.4.0",
                     )
                 )
@@ -237,7 +229,7 @@ def check_ml_libraries() -> List[DiagnosticIssue]:
                     issues.append(
                         DiagnosticIssue(
                             name=f"{lib_name}",
-                            status=f"PASS - Installed",
+                            status="PASS - Installed",
                             severity="info",
                             fix="",
                             details=f"Version: {version}",
@@ -275,7 +267,6 @@ def check_gpu_memory() -> List[DiagnosticIssue]:
         free_mem, total_mem = torch.cuda.mem_get_info(0)
         free_gb = free_mem / (1024**3)
         total_gb = total_mem / (1024**3)
-        used_gb = total_gb - free_gb
 
         if free_gb < 8:
             issues.append(
@@ -377,7 +368,18 @@ def check_docker_gpu() -> List[DiagnosticIssue]:
         return issues
 
     try:
-        result = run_command(["docker", "run", "--rm", "--gpus", "all", "nvidia/cuda:12.4.0-base-ubuntu22.04", "nvidia-smi"], timeout=30)
+        result = run_command(
+            [
+                "docker",
+                "run",
+                "--rm",
+                "--gpus",
+                "all",
+                "nvidia/cuda:12.4.0-base-ubuntu22.04",
+                "nvidia-smi",
+            ],
+            timeout=30,
+        )
         if result.returncode == 0:
             issues.append(
                 DiagnosticIssue(
@@ -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,43 +434,128 @@ 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.
 
-    console.print("[bold blue]🔍 Running ML Environment Diagnostics...[/bold blue]\n")
+    Args:
+        full: Whether to run full diagnostics including extended checks
+        parallel: Whether to run independent checks in parallel
 
-    # Core checks (always run)
-    all_issues.extend(check_cuda_driver())
-    all_issues.extend(check_pytorch_cuda())
-    all_issues.extend(check_ml_libraries())
+    Returns:
+        List of diagnostic issues found
+    """
+    from .parallel import run_parallel_with_results
+
+    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
 
 
 def print_diagnostic_table(issues: List[DiagnosticIssue]) -> None:
     """Print diagnostic results as a Rich table."""
-    table = Table(title="ML Environment Doctor - Diagnostic Results", show_header=True, header_style="bold magenta")
+    table = Table(
+        title="ML Environment Doctor - Diagnostic Results",
+        show_header=True,
+        header_style="bold magenta",
+    )
     table.add_column("Issue", style="cyan", no_wrap=False)
     table.add_column("Status", style="bold")
     table.add_column("Severity", style="yellow")
@@ -474,20 +569,25 @@ def print_diagnostic_table(issues: List[DiagnosticIssue]) -> None:
 
     # Summary
     critical_count = sum(1 for i in issues if i.severity == "critical" and "FAIL" in i.status)
-    warning_count = sum(1 for i in issues if i.severity == "warning" and ("WARN" in i.status or "FAIL" in i.status))
+    warning_count = sum(
+        1 for i in issues if i.severity == "warning" and ("WARN" in i.status or "FAIL" in i.status)
+    )
     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]")
+        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]")
-
+        console.print(
+            "\n[bold yellow]💡 Consider addressing warnings for optimal performance.[/bold yellow]"
+        )

mlenvdoctor/dockerize.py

@@ -3,9 +3,7 @@
 from pathlib import Path
 from typing import Optional
 
-from rich.console import Console
-
-from .utils import console, print_error, print_info, print_success
+from .utils import console, print_info, print_success
 
 # Model-specific templates
 MODEL_TEMPLATES = {
@@ -149,9 +147,9 @@ CMD ["python", "train.py"]
     console.print("[bold]Build and run:[/bold]")
     console.print(f"[cyan]  docker build -f {output_file} -t mlenvdoctor .[/cyan]")
     if service:
-        console.print(f"[cyan]  docker run --gpus all -p 8000:8000 mlenvdoctor[/cyan]")
+        console.print("[cyan]  docker run --gpus all -p 8000:8000 mlenvdoctor[/cyan]")
     else:
-        console.print(f"[cyan]  docker run --gpus all -v $(pwd)/data:/app/data mlenvdoctor[/cyan]")
+        console.print("[cyan]  docker run --gpus all -v $(pwd)/data:/app/data mlenvdoctor[/cyan]")
 
     return output_path
 
@@ -201,4 +199,3 @@ if __name__ == "__main__":
     output_path.write_text(service_content, encoding="utf-8")
     print_success(f"Generated service template: {output_file}")
     return output_path
-

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

mlenvdoctor/export.py

@@ -0,0 +1,290 @@
+"""Export functionality for diagnostic results."""
+
+import json
+from datetime import datetime
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+from .diagnose import DiagnosticIssue
+
+
+def issue_to_dict(issue: DiagnosticIssue) -> Dict[str, Any]:
+    """Convert DiagnosticIssue to dictionary."""
+    return {
+        "name": issue.name,
+        "status": issue.status,
+        "severity": issue.severity,
+        "fix": issue.fix,
+        "details": issue.details,
+    }
+
+
+def export_json(
+    issues: List[DiagnosticIssue],
+    output_file: Optional[Path] = None,
+    include_metadata: bool = True,
+) -> Path:
+    """
+    Export diagnostic results to JSON.
+
+    Args:
+        issues: List of diagnostic issues
+        output_file: Output file path (default: diagnostic-results.json)
+        include_metadata: Include metadata (timestamp, version, etc.)
+
+    Returns:
+        Path to exported file
+    """
+    if output_file is None:
+        output_file = Path("diagnostic-results.json")
+
+    # Convert issues to dictionaries
+    issues_data = [issue_to_dict(issue) for issue in issues]
+
+    # Calculate summary
+    critical_count = sum(
+        1 for i in issues if i.severity == "critical" and "FAIL" in i.status
+    )
+    warning_count = sum(
+        1 for i in issues if i.severity == "warning" and ("WARN" in i.status or "FAIL" in i.status)
+    )
+    pass_count = sum(1 for i in issues if "PASS" in i.status)
+
+    # Build export data
+    export_data: Dict[str, Any] = {
+        "issues": issues_data,
+        "summary": {
+            "total": len(issues),
+            "passed": pass_count,
+            "warnings": warning_count,
+            "critical": critical_count,
+        },
+    }
+
+    # Add metadata if requested
+    if include_metadata:
+        from . import __version__
+
+        export_data["metadata"] = {
+            "version": __version__,
+            "timestamp": datetime.now().isoformat(),
+            "tool": "mlenvdoctor",
+        }
+
+    # Write to file
+    output_file.write_text(json.dumps(export_data, indent=2, ensure_ascii=False), encoding="utf-8")
+
+    return output_file
+
+
+def export_csv(issues: List[DiagnosticIssue], output_file: Optional[Path] = None) -> Path:
+    """
+    Export diagnostic results to CSV.
+
+    Args:
+        issues: List of diagnostic issues
+        output_file: Output file path (default: diagnostic-results.csv)
+
+    Returns:
+        Path to exported file
+    """
+    import csv
+
+    if output_file is None:
+        output_file = Path("diagnostic-results.csv")
+
+    with output_file.open("w", newline="", encoding="utf-8") as f:
+        writer = csv.writer(f)
+        writer.writerow(["Issue", "Status", "Severity", "Fix", "Details"])
+
+        for issue in issues:
+            writer.writerow(
+                [
+                    issue.name,
+                    issue.status,
+                    issue.severity,
+                    issue.fix,
+                    issue.details or "",
+                ]
+            )
+
+    return output_file
+
+
+def export_html(issues: List[DiagnosticIssue], output_file: Optional[Path] = None) -> Path:
+    """
+    Export diagnostic results to HTML report.
+
+    Args:
+        issues: List of diagnostic issues
+        output_file: Output file path (default: diagnostic-results.html)
+
+    Returns:
+        Path to exported file
+    """
+    if output_file is None:
+        output_file = Path("diagnostic-results.html")
+
+    # Calculate summary
+    critical_count = sum(
+        1 for i in issues if i.severity == "critical" and "FAIL" in i.status
+    )
+    warning_count = sum(
+        1 for i in issues if i.severity == "warning" and ("WARN" in i.status or "FAIL" in i.status)
+    )
+    pass_count = sum(1 for i in issues if "PASS" in i.status)
+
+    from . import __version__
+
+    html_content = f"""<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>ML Environment Doctor - Diagnostic Report</title>
+    <style>
+        body {{
+            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, Cantarell, sans-serif;
+            max-width: 1200px;
+            margin: 0 auto;
+            padding: 20px;
+            background-color: #f5f5f5;
+        }}
+        .header {{
+            background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+            color: white;
+            padding: 30px;
+            border-radius: 10px;
+            margin-bottom: 20px;
+        }}
+        .summary {{
+            display: grid;
+            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+            gap: 15px;
+            margin-bottom: 30px;
+        }}
+        .summary-card {{
+            background: white;
+            padding: 20px;
+            border-radius: 8px;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        }}
+        .summary-card h3 {{
+            margin: 0 0 10px 0;
+            font-size: 14px;
+            color: #666;
+        }}
+        .summary-card .value {{
+            font-size: 32px;
+            font-weight: bold;
+        }}
+        .passed {{ color: #10b981; }}
+        .warning {{ color: #f59e0b; }}
+        .critical {{ color: #ef4444; }}
+        table {{
+            width: 100%;
+            background: white;
+            border-collapse: collapse;
+            border-radius: 8px;
+            overflow: hidden;
+            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+        }}
+        th {{
+            background-color: #667eea;
+            color: white;
+            padding: 15px;
+            text-align: left;
+            font-weight: 600;
+        }}
+        td {{
+            padding: 12px 15px;
+            border-bottom: 1px solid #e5e7eb;
+        }}
+        tr:hover {{
+            background-color: #f9fafb;
+        }}
+        .status-pass {{ color: #10b981; font-weight: 600; }}
+        .status-fail {{ color: #ef4444; font-weight: 600; }}
+        .status-warn {{ color: #f59e0b; font-weight: 600; }}
+        .severity-critical {{ background-color: #fee2e2; }}
+        .severity-warning {{ background-color: #fef3c7; }}
+        .severity-info {{ background-color: #dbeafe; }}
+        .footer {{
+            margin-top: 30px;
+            text-align: center;
+            color: #666;
+            font-size: 12px;
+        }}
+    </style>
+</head>
+<body>
+    <div class="header">
+        <h1>ML Environment Doctor</h1>
+        <p>Diagnostic Report - Generated on {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}</p>
+        <p style="font-size: 14px; opacity: 0.9;">Version {__version__}</p>
+    </div>
+
+    <div class="summary">
+        <div class="summary-card">
+            <h3>Total Checks</h3>
+            <div class="value">{len(issues)}</div>
+        </div>
+        <div class="summary-card">
+            <h3>Passed</h3>
+            <div class="value passed">{pass_count}</div>
+        </div>
+        <div class="summary-card">
+            <h3>Warnings</h3>
+            <div class="value warning">{warning_count}</div>
+        </div>
+        <div class="summary-card">
+            <h3>Critical Issues</h3>
+            <div class="value critical">{critical_count}</div>
+        </div>
+    </div>
+
+    <table>
+        <thead>
+            <tr>
+                <th>Issue</th>
+                <th>Status</th>
+                <th>Severity</th>
+                <th>Fix</th>
+                <th>Details</th>
+            </tr>
+        </thead>
+        <tbody>
+"""
+
+    for issue in issues:
+        status_class = "status-pass"
+        if "FAIL" in issue.status:
+            status_class = "status-fail"
+        elif "WARN" in issue.status:
+            status_class = "status-warn"
+
+        severity_class = f"severity-{issue.severity}"
+
+        html_content += f"""
+            <tr class="{severity_class}">
+                <td><strong>{issue.name}</strong></td>
+                <td class="{status_class}">{issue.status}</td>
+                <td>{issue.severity.upper()}</td>
+                <td>{issue.fix or '-'}</td>
+                <td>{issue.details or '-'}</td>
+            </tr>
+"""
+
+    html_content += """
+        </tbody>
+    </table>
+
+    <div class="footer">
+        <p>Generated by ML Environment Doctor | <a href="https://github.com/dheena731/ml_env_doctor">GitHub</a></p>
+    </div>
+</body>
+</html>
+"""
+
+    output_file.write_text(html_content, encoding="utf-8")
+    return output_file