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

mlenvdoctor/fix.py

@@ -1,14 +1,11 @@
 """Auto-fix and requirements generation for ML Environment Doctor."""
 
-import subprocess
 import sys
 from pathlib import Path
-from typing import List, Optional
+from typing import Optional
 
-from rich.console import Console
-from rich.progress import Progress, SpinnerColumn, TextColumn
-
-from .diagnose import DiagnosticIssue, diagnose_env
+from .diagnose import diagnose_env
+from .icons import icon_wrench
 from .utils import (
     check_command_exists,
     console,
@@ -39,7 +36,9 @@ ML_STACKS = {
 }
 
 
-def generate_requirements_txt(stack: str = "trl-peft", output_file: str = "requirements-mlenvdoctor.txt") -> Path:
+def generate_requirements_txt(
+    stack: str = "trl-peft", output_file: str = "requirements-mlenvdoctor.txt"
+) -> Path:
     """Generate requirements.txt file."""
     if stack not in ML_STACKS:
         print_error(f"Unknown stack: {stack}. Available: {list(ML_STACKS.keys())}")
@@ -61,7 +60,9 @@ def generate_requirements_txt(stack: str = "trl-peft", output_file: str = "requi
             content = "# Standard PyTorch (CPU or CUDA)\n\n"
     except ImportError:
         content = "# PyTorch installation\n"
-        content += "# For CUDA: pip install torch --index-url https://download.pytorch.org/whl/cu124\n"
+        content += (
+            "# For CUDA: pip install torch --index-url https://download.pytorch.org/whl/cu124\n"
+        )
         content += "# For CPU: pip install torch\n\n"
 
     content += "\n".join(requirements)
@@ -72,7 +73,9 @@ def generate_requirements_txt(stack: str = "trl-peft", output_file: str = "requi
     return output_path
 
 
-def generate_conda_env(stack: str = "trl-peft", output_file: str = "environment-mlenvdoctor.yml") -> Path:
+def generate_conda_env(
+    stack: str = "trl-peft", output_file: str = "environment-mlenvdoctor.yml"
+) -> Path:
     """Generate conda environment file."""
     if stack not in ML_STACKS:
         print_error(f"Unknown stack: {stack}. Available: {list(ML_STACKS.keys())}")
@@ -197,7 +200,11 @@ def create_virtualenv(env_name: str = ".venv") -> Optional[Path]:
 
         venv.create(env_path, with_pip=True)
         print_success(f"Virtual environment created: {env_name}")
-        print_info(f"Activate with: {'.venv\\Scripts\\activate' if sys.platform == 'win32' else 'source .venv/bin/activate'}")
+        if sys.platform == "win32":
+            activate_cmd = r".venv\Scripts\activate"
+        else:
+            activate_cmd = "source .venv/bin/activate"
+        print_info(f"Activate with: {activate_cmd}")
         return env_path
     except Exception as e:
         print_error(f"Failed to create virtual environment: {e}")
@@ -206,7 +213,7 @@ def create_virtualenv(env_name: str = ".venv") -> Optional[Path]:
 
 def auto_fix(use_conda: bool = False, create_venv: bool = False, stack: str = "trl-peft") -> bool:
     """Auto-fix environment issues based on diagnostics."""
-    console.print("[bold blue]🔧 Running Auto-Fix...[/bold blue]\n")
+    console.print(f"[bold blue]{icon_wrench()} Running Auto-Fix...[/bold blue]\n")
 
     # Run diagnostics
     issues = diagnose_env(full=False)
@@ -243,7 +250,6 @@ def auto_fix(use_conda: bool = False, create_venv: bool = False, stack: str = "t
         if install.lower() in ["y", "yes"]:
             return install_requirements(str(req_file), use_conda=use_conda)
         else:
-            print_info(f"Requirements file generated. Install manually with:")
+            print_info("Requirements file generated. Install manually with:")
             console.print(f"[cyan]  pip install -r {req_file}[/cyan]")
             return True
-

mlenvdoctor/gpu.py

@@ -1,19 +1,24 @@
 """GPU benchmarks and smoke tests for ML Environment Doctor."""
 
 import time
-from typing import Dict, List, Optional
+from typing import Dict
 
 try:
     import torch
 except ImportError:
     torch = None  # type: ignore
 
+import sys
+
 from rich.console import Console
-from rich.progress import Progress, SpinnerColumn, TextColumn
 
 from .utils import print_error, print_info, print_success
 
-console = Console()
+# Configure console for Windows compatibility
+if sys.platform == "win32":
+    console = Console(legacy_windows=True, force_terminal=True)
+else:
+    console = Console()
 
 
 def benchmark_gpu_ops() -> Dict[str, float]:
@@ -78,7 +83,8 @@ def smoke_test_lora() -> bool:
                 if tokenizer.pad_token is None:
                     tokenizer.pad_token = tokenizer.eos_token
                 model = AutoModelForCausalLM.from_pretrained(
-                    model_name, torch_dtype=torch.float16 if device.type == "cuda" else torch.float32
+                    model_name,
+                    torch_dtype=torch.float16 if device.type == "cuda" else torch.float32,
                 ).to(device)
 
             # Configure LoRA
@@ -101,8 +107,7 @@ def smoke_test_lora() -> bool:
             # Forward pass
             with console.status("[bold green]Running forward pass..."):
                 with torch.no_grad():
-                    outputs = model(**inputs)
-                    loss = outputs.loss if hasattr(outputs, "loss") else None
+                    _ = model(**inputs)
 
             print_success("LoRA smoke test passed!")
             return True
@@ -146,7 +151,9 @@ def test_model(model_name: str = "tinyllama") -> bool:
         # Estimate memory requirements (rough)
         if "7b" in actual_model_name.lower() or "7B" in actual_model_name:
             if free_gb < 16:
-                print_error(f"Insufficient GPU memory: {free_gb:.1f}GB free, need ~16GB for 7B model")
+                print_error(
+                    f"Insufficient GPU memory: {free_gb:.1f}GB free, need ~16GB for 7B model"
+                )
                 return False
 
         with console.status(f"[bold green]Loading {actual_model_name}..."):
@@ -166,7 +173,7 @@ def test_model(model_name: str = "tinyllama") -> bool:
                 inputs = tokenizer(dummy_text, return_tensors="pt").to(device)
 
                 with torch.no_grad():
-                    outputs = model(**inputs)
+                    _ = model(**inputs)
 
                 print_success(f"Model {actual_model_name} loaded and tested successfully!")
                 return True
@@ -181,4 +188,3 @@ def test_model(model_name: str = "tinyllama") -> bool:
     except Exception as e:
         print_error(f"Model test error: {e}")
         return False
-

mlenvdoctor/icons.py

@@ -0,0 +1,100 @@
+"""Safe emoji/icon handling for cross-platform compatibility."""
+
+import sys
+from typing import Literal
+
+# Check if we can safely use emojis
+_USE_EMOJIS = True
+
+try:
+    # Try to write an emoji to see if it works
+    if sys.platform == "win32":
+        import io
+
+        # Test if console supports UTF-8
+        test_output = io.StringIO()
+        try:
+            test_output.write("🔍")
+            test_output.getvalue()
+        except (UnicodeEncodeError, UnicodeError):
+            _USE_EMOJIS = False
+except Exception:
+    _USE_EMOJIS = False
+
+
+def get_icon(icon_name: Literal["search", "check", "cross", "warning", "info", "wrench", "whale", "test"]) -> str:
+    """
+    Get a safe icon/emoji for the current platform.
+
+    Args:
+        icon_name: Name of the icon to get
+
+    Returns:
+        Emoji if supported, ASCII alternative otherwise
+    """
+    if _USE_EMOJIS:
+        icons = {
+            "search": "🔍",
+            "check": "✅",
+            "cross": "❌",
+            "warning": "⚠️",
+            "info": "ℹ️",
+            "wrench": "🔧",
+            "whale": "🐳",
+            "test": "🧪",
+        }
+    else:
+        # ASCII alternatives
+        icons = {
+            "search": "[*]",
+            "check": "[OK]",
+            "cross": "[X]",
+            "warning": "[!]",
+            "info": "[i]",
+            "wrench": "[FIX]",
+            "whale": "[DOCKER]",
+            "test": "[TEST]",
+        }
+
+    return icons.get(icon_name, "")
+
+
+# Convenience functions
+def icon_search() -> str:
+    """Get search icon."""
+    return get_icon("search")
+
+
+def icon_check() -> str:
+    """Get check icon."""
+    return get_icon("check")
+
+
+def icon_cross() -> str:
+    """Get cross/error icon."""
+    return get_icon("cross")
+
+
+def icon_warning() -> str:
+    """Get warning icon."""
+    return get_icon("warning")
+
+
+def icon_info() -> str:
+    """Get info icon."""
+    return get_icon("info")
+
+
+def icon_wrench() -> str:
+    """Get wrench/fix icon."""
+    return get_icon("wrench")
+
+
+def icon_whale() -> str:
+    """Get whale/docker icon."""
+    return get_icon("whale")
+
+
+def icon_test() -> str:
+    """Get test icon."""
+    return get_icon("test")

mlenvdoctor/logger.py

@@ -0,0 +1,81 @@
+"""Logging configuration for ML Environment Doctor."""
+
+import logging
+import sys
+from pathlib import Path
+from typing import Optional
+
+from rich.console import Console
+from rich.logging import RichHandler
+
+from .utils import get_home_config_dir
+
+console = Console()
+
+
+def setup_logger(
+    name: str = "mlenvdoctor",
+    level: str = "INFO",
+    log_file: Optional[Path] = None,
+    enable_rich: bool = True,
+) -> logging.Logger:
+    """
+    Set up logger with Rich console handler and optional file handler.
+
+    Args:
+        name: Logger name
+        level: Logging level (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+        log_file: Optional path to log file
+        enable_rich: Use Rich handler for console output
+
+    Returns:
+        Configured logger instance
+    """
+    logger = logging.getLogger(name)
+    logger.setLevel(getattr(logging, level.upper()))
+
+    # Remove existing handlers to avoid duplicates
+    logger.handlers.clear()
+
+    # Console handler with Rich formatting
+    if enable_rich:
+        console_handler = RichHandler(
+            console=console,
+            show_time=True,
+            show_path=False,
+            rich_tracebacks=True,
+            tracebacks_show_locals=False,
+        )
+    else:
+        console_handler = logging.StreamHandler(sys.stdout)
+        console_handler.setFormatter(
+            logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
+        )
+
+    console_handler.setLevel(getattr(logging, level.upper()))
+    logger.addHandler(console_handler)
+
+    # File handler if log file specified
+    if log_file:
+        log_file.parent.mkdir(parents=True, exist_ok=True)
+        file_handler = logging.FileHandler(log_file, encoding="utf-8")
+        file_handler.setLevel(logging.DEBUG)  # Always log everything to file
+        file_handler.setFormatter(
+            logging.Formatter(
+                "%(asctime)s - %(name)s - %(levelname)s - %(funcName)s:%(lineno)d - %(message)s"
+            )
+        )
+        logger.addHandler(file_handler)
+
+    return logger
+
+
+def get_default_log_file() -> Path:
+    """Get default log file path."""
+    log_dir = get_home_config_dir() / "logs"
+    log_dir.mkdir(parents=True, exist_ok=True)
+    return log_dir / "mlenvdoctor.log"
+
+
+# Default logger instance
+logger = setup_logger()

mlenvdoctor/parallel.py

@@ -0,0 +1,115 @@
+"""Parallel execution utilities for independent operations."""
+
+import concurrent.futures
+from typing import Callable, Iterable, List, TypeVar
+
+from .logger import logger
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+def run_parallel(
+    func: Callable[[T], R],
+    items: Iterable[T],
+    max_workers: int = 4,
+    timeout: float | None = None,
+) -> List[R]:
+    """
+    Run a function in parallel on multiple items.
+
+    Args:
+        func: Function to execute
+        items: Iterable of items to process
+        items_list: List of items to process
+        max_workers: Maximum number of parallel workers
+        timeout: Maximum time to wait for all tasks (None = no timeout)
+
+    Returns:
+        List of results in the same order as input items
+
+    Example:
+        def check_library(name: str) -> bool:
+            return importlib.util.find_spec(name) is not None
+
+        results = run_parallel(check_library, ["torch", "transformers", "peft"])
+    """
+    items_list = list(items)
+
+    if not items_list:
+        return []
+
+    # Use ThreadPoolExecutor for I/O-bound operations
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        # Submit all tasks
+        future_to_item = {executor.submit(func, item): item for item in items_list}
+
+        results: List[R] = []
+        completed = 0
+
+        # Process completed tasks
+        for future in concurrent.futures.as_completed(future_to_item, timeout=timeout):
+            item = future_to_item[future]
+            try:
+                result = future.result()
+                results.append(result)
+                completed += 1
+            except Exception as e:
+                logger.error(f"Error processing {item}: {e}")
+                # Re-raise to maintain error behavior
+                raise
+
+        if completed != len(items_list):
+            raise RuntimeError(f"Only {completed}/{len(items_list)} tasks completed")
+
+        return results
+
+
+def run_parallel_with_results(
+    func: Callable[[T], R],
+    items: Iterable[T],
+    max_workers: int = 4,
+    timeout: float | None = None,
+) -> List[tuple[T, R | Exception]]:
+    """
+    Run a function in parallel and return results with original items.
+
+    Unlike run_parallel, this catches exceptions and returns them as results.
+
+    Args:
+        func: Function to execute
+        items: Iterable of items to process
+        max_workers: Maximum number of parallel workers
+        timeout: Maximum time to wait for all tasks
+
+    Returns:
+        List of (item, result_or_exception) tuples
+
+    Example:
+        def check_library(name: str) -> bool:
+            if name == "bad":
+                raise ValueError("Bad library")
+            return True
+
+        results = run_parallel_with_results(check_library, ["torch", "bad", "peft"])
+        # Returns: [("torch", True), ("bad", ValueError(...)), ("peft", True)]
+    """
+    items_list = list(items)
+
+    if not items_list:
+        return []
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
+        future_to_item = {executor.submit(func, item): item for item in items_list}
+
+        results: List[tuple[T, R | Exception]] = []
+
+        for future in concurrent.futures.as_completed(future_to_item, timeout=timeout):
+            item = future_to_item[future]
+            try:
+                result = future.result()
+                results.append((item, result))
+            except Exception as e:
+                results.append((item, e))
+
+        return results

mlenvdoctor/retry.py

@@ -0,0 +1,92 @@
+"""Retry logic for transient failures."""
+
+import functools
+import time
+from typing import Any, Callable, Optional, TypeVar
+
+from .exceptions import DiagnosticError
+from .logger import logger
+
+T = TypeVar("T")
+
+
+def retry(
+    max_attempts: int = 3,
+    delay: float = 1.0,
+    backoff: float = 2.0,
+    exceptions: tuple[type[Exception], ...] = (Exception,),
+    on_retry: Optional[Callable[[Exception, int], None]] = None,
+) -> Callable[[Callable[..., T]], Callable[..., T]]:
+    """
+    Decorator to retry a function on failure.
+
+    Args:
+        max_attempts: Maximum number of retry attempts
+        delay: Initial delay between retries in seconds
+        backoff: Multiplier for delay after each retry
+        exceptions: Tuple of exceptions to catch and retry
+        on_retry: Optional callback called on each retry
+
+    Returns:
+        Decorated function with retry logic
+
+    Example:
+        @retry(max_attempts=3, delay=1.0, exceptions=(ConnectionError,))
+        def fetch_data():
+            ...
+    """
+
+    def decorator(func: Callable[..., T]) -> Callable[..., T]:
+        @functools.wraps(func)
+        def wrapper(*args: Any, **kwargs: Any) -> T:
+            current_delay = delay
+            last_exception: Optional[Exception] = None
+
+            for attempt in range(1, max_attempts + 1):
+                try:
+                    return func(*args, **kwargs)
+                except exceptions as e:
+                    last_exception = e
+                    if attempt < max_attempts:
+                        logger.warning(
+                            f"{func.__name__} failed (attempt {attempt}/{max_attempts}): {e}. "
+                            f"Retrying in {current_delay:.1f}s..."
+                        )
+                        if on_retry:
+                            on_retry(e, attempt)
+                        time.sleep(current_delay)
+                        current_delay *= backoff
+                    else:
+                        logger.error(f"{func.__name__} failed after {max_attempts} attempts: {e}")
+
+            # All attempts failed
+            if last_exception:
+                raise DiagnosticError(
+                    f"{func.__name__} failed after {max_attempts} attempts",
+                    f"Last error: {last_exception}",
+                ) from last_exception
+
+            # Should never reach here, but satisfy type checker
+            raise RuntimeError("Retry logic failed unexpectedly")
+
+        return wrapper
+
+    return decorator
+
+
+def retry_network(func: Callable[..., T]) -> Callable[..., T]:
+    """
+    Decorator specifically for network operations.
+
+    Retries on network-related exceptions with exponential backoff.
+    """
+    return retry(
+        max_attempts=3,
+        delay=1.0,
+        backoff=2.0,
+        exceptions=(
+            ConnectionError,
+            TimeoutError,
+            OSError,
+        ),
+    )(func)