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

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)

mlenvdoctor/utils.py

@@ -5,10 +5,20 @@ import sys
 from pathlib import Path
 from typing import List, Optional, Tuple
 
+import sys
+
 from rich.console import Console
 from rich.progress import Progress, SpinnerColumn, TextColumn
 
-console = Console()
+from .exceptions import DiagnosticError
+from .icons import icon_check, icon_cross, icon_info, icon_warning
+
+# Configure console for Windows compatibility
+if sys.platform == "win32":
+    # Use legacy Windows renderer if needed
+    console = Console(legacy_windows=True, force_terminal=True)
+else:
+    console = Console()
 
 
 def run_command(
@@ -16,8 +26,29 @@ def run_command(
     capture_output: bool = True,
     check: bool = False,
     timeout: Optional[int] = 30,
-) -> subprocess.CompletedProcess:
-    """Run a shell command with error handling."""
+) -> subprocess.CompletedProcess[str]:
+    """
+    Run a shell command with error handling and input validation.
+
+    Args:
+        cmd: Command and arguments as a list
+        capture_output: Whether to capture stdout/stderr
+        check: Whether to raise on non-zero exit code
+        timeout: Command timeout in seconds
+
+    Returns:
+        CompletedProcess with command result
+
+    Raises:
+        DiagnosticError: For command execution errors
+        ConfigurationError: For invalid input
+    """
+    from .validators import sanitize_command, validate_timeout
+
+    # Validate and sanitize inputs
+    cmd = sanitize_command(cmd)
+    timeout = validate_timeout(timeout)
+
     try:
         result = subprocess.run(
             cmd,
@@ -27,29 +58,58 @@ def run_command(
             timeout=timeout,
         )
         return result
-    except subprocess.TimeoutExpired:
-        console.print(f"[red]Command timed out: {' '.join(cmd)}[/red]")
-        raise
-    except FileNotFoundError:
-        console.print(f"[red]Command not found: {cmd[0]}[/red]")
-        raise
+    except subprocess.TimeoutExpired as e:
+        error_msg = f"Command timed out after {timeout}s: {' '.join(cmd)}"
+        console.print(f"[red]{error_msg}[/red]")
+        raise DiagnosticError(
+            error_msg,
+            "Try increasing timeout or check if the command is hanging",
+        ) from e
+    except FileNotFoundError as e:
+        error_msg = f"Command not found: {cmd[0]}"
+        console.print(f"[red]{error_msg}[/red]")
+        raise DiagnosticError(
+            error_msg,
+            f"Install {cmd[0]} or ensure it's in your PATH",
+        ) from e
     except subprocess.CalledProcessError as e:
         if not check:
-            return e  # type: ignore
+            # Return the exception as if it were a result
+            # This maintains backward compatibility but is type-unsafe
+            return subprocess.CompletedProcess(  # type: ignore[return-value]
+                cmd, e.returncode, e.stdout, e.stderr
+            )
         raise
 
 
 def check_command_exists(cmd: str) -> bool:
-    """Check if a command exists in PATH."""
+    """
+    Check if a command exists in PATH.
+
+    Args:
+        cmd: Command name to check
+
+    Returns:
+        True if command exists and is executable, False otherwise
+    """
+    if not isinstance(cmd, str) or not cmd.strip():
+        return False
+
     try:
-        subprocess.run(
-            [cmd, "--version"] if cmd != "nvidia-smi" else [cmd],
+        # Use 'which' on Unix, 'where' on Windows
+        if sys.platform == "win32":
+            check_cmd = ["where", cmd]
+        else:
+            check_cmd = ["which", cmd]
+
+        result = subprocess.run(
+            check_cmd,
             capture_output=True,
             timeout=5,
             check=False,
         )
-        return True
-    except (FileNotFoundError, subprocess.TimeoutExpired):
+        return result.returncode == 0
+    except (FileNotFoundError, subprocess.TimeoutExpired, Exception):
         return False
 
 
@@ -63,22 +123,22 @@ def get_home_config_dir() -> Path:
 
 def print_success(message: str) -> None:
     """Print a success message."""
-    console.print(f"[green]✅ {message}[/green]")
+    console.print(f"[green]{icon_check()} {message}[/green]")
 
 
 def print_error(message: str) -> None:
     """Print an error message."""
-    console.print(f"[red]❌ {message}[/red]")
+    console.print(f"[red]{icon_cross()} {message}[/red]")
 
 
 def print_warning(message: str) -> None:
     """Print a warning message."""
-    console.print(f"[yellow]⚠️  {message}[/yellow]")
+    console.print(f"[yellow]{icon_warning()}  {message}[/yellow]")
 
 
 def print_info(message: str) -> None:
     """Print an info message."""
-    console.print(f"[blue]ℹ️  {message}[/blue]")
+    console.print(f"[blue]{icon_info()}  {message}[/blue]")
 
 
 def with_spinner(message: str):

mlenvdoctor/validators.py

@@ -0,0 +1,217 @@
+"""Input validation and sanitization for ML Environment Doctor."""
+
+import re
+from pathlib import Path
+from typing import Optional
+
+from .exceptions import ConfigurationError
+
+
+def validate_model_name(model_name: str) -> str:
+    """
+    Validate and sanitize model name.
+
+    Args:
+        model_name: Model name to validate
+
+    Returns:
+        Sanitized model name
+
+    Raises:
+        ConfigurationError: If model name is invalid
+    """
+    if not model_name or not isinstance(model_name, str):
+        raise ConfigurationError(
+            "Model name must be a non-empty string",
+            "Use a valid model name like 'tinyllama', 'gpt2', or 'mistral-7b'",
+        )
+
+    # Remove whitespace
+    model_name = model_name.strip()
+
+    # Check for dangerous characters (basic sanitization)
+    if not re.match(r"^[a-zA-Z0-9._-]+$", model_name):
+        raise ConfigurationError(
+            f"Invalid model name: {model_name}",
+            "Model name can only contain letters, numbers, dots, underscores, and hyphens",
+        )
+
+    return model_name.lower()
+
+
+def validate_file_path(file_path: Path, must_exist: bool = False, must_be_file: bool = False) -> Path:
+    """
+    Validate and sanitize file path.
+
+    Args:
+        file_path: Path to validate
+        must_exist: Whether the path must exist
+        must_be_file: Whether the path must be a file
+
+    Returns:
+        Resolved, absolute path
+
+    Raises:
+        ConfigurationError: If path is invalid
+    """
+    if not isinstance(file_path, (Path, str)):
+        raise ConfigurationError(
+            "File path must be a Path object or string",
+            "Use pathlib.Path or a valid string path",
+        )
+
+    path = Path(file_path).resolve()
+
+    # Check for path traversal attempts
+    if ".." in str(path):
+        # Resolve should handle this, but double-check
+        resolved = path.resolve()
+        if ".." in str(resolved):
+            raise ConfigurationError(
+                "Invalid path: contains '..'",
+                "Use absolute paths or relative paths without '..'",
+            )
+
+    if must_exist and not path.exists():
+        raise ConfigurationError(
+            f"Path does not exist: {path}",
+            "Ensure the file or directory exists",
+        )
+
+    if must_be_file and not path.is_file():
+        raise ConfigurationError(
+            f"Path is not a file: {path}",
+            "Provide a valid file path",
+        )
+
+    return path
+
+
+def validate_log_level(level: str) -> str:
+    """
+    Validate logging level.
+
+    Args:
+        level: Log level to validate
+
+    Returns:
+        Validated log level
+
+    Raises:
+        ConfigurationError: If level is invalid
+    """
+    valid_levels = {"DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"}
+    level_upper = level.upper() if isinstance(level, str) else str(level).upper()
+
+    if level_upper not in valid_levels:
+        raise ConfigurationError(
+            f"Invalid log level: {level}",
+            f"Use one of: {', '.join(valid_levels)}",
+        )
+
+    return level_upper
+
+
+def validate_stack_name(stack: str) -> str:
+    """
+    Validate ML stack name.
+
+    Args:
+        stack: Stack name to validate
+
+    Returns:
+        Validated stack name
+
+    Raises:
+        ConfigurationError: If stack is invalid
+    """
+    valid_stacks = {"trl-peft", "minimal"}
+    stack_lower = stack.lower() if isinstance(stack, str) else str(stack).lower()
+
+    if stack_lower not in valid_stacks:
+        raise ConfigurationError(
+            f"Invalid stack: {stack}",
+            f"Use one of: {', '.join(valid_stacks)}",
+        )
+
+    return stack_lower
+
+
+def sanitize_command(cmd: list[str]) -> list[str]:
+    """
+    Sanitize command arguments to prevent injection.
+
+    Args:
+        cmd: Command and arguments list
+
+    Returns:
+        Sanitized command list
+
+    Raises:
+        ConfigurationError: If command contains dangerous patterns
+    """
+    if not isinstance(cmd, list) or not cmd:
+        raise ConfigurationError(
+            "Command must be a non-empty list",
+            "Provide command as a list of strings",
+        )
+
+    sanitized = []
+    for arg in cmd:
+        if not isinstance(arg, str):
+            raise ConfigurationError(
+                "All command arguments must be strings",
+                "Convert all arguments to strings",
+            )
+
+        # Check for command injection patterns
+        dangerous_patterns = [";", "&&", "||", "`", "$(", "<", ">", "|"]
+        for pattern in dangerous_patterns:
+            if pattern in arg:
+                raise ConfigurationError(
+                    f"Dangerous pattern detected in command: {pattern}",
+                    "Do not use shell operators in command arguments",
+                )
+
+        sanitized.append(arg)
+
+    return sanitized
+
+
+def validate_timeout(timeout: Optional[int], min_timeout: int = 1, max_timeout: int = 3600) -> Optional[int]:
+    """
+    Validate timeout value.
+
+    Args:
+        timeout: Timeout in seconds
+        min_timeout: Minimum allowed timeout
+        max_timeout: Maximum allowed timeout
+
+    Returns:
+        Validated timeout
+
+    Raises:
+        ConfigurationError: If timeout is invalid
+    """
+    if timeout is None:
+        return None
+
+    if not isinstance(timeout, int):
+        raise ConfigurationError(
+            f"Timeout must be an integer, got {type(timeout)}",
+            "Provide timeout as an integer number of seconds",
+        )
+
+    if timeout < min_timeout:
+        raise ConfigurationError(
+            f"Timeout too small: {timeout}s (minimum: {min_timeout}s)",
+            f"Increase timeout to at least {min_timeout} seconds",
+        )
+
+    if timeout > max_timeout:
+        raise ConfigurationError(
+            f"Timeout too large: {timeout}s (maximum: {max_timeout}s)",
+            f"Decrease timeout to at most {max_timeout} seconds",
+        )
+
+    return timeout

{mlenvdoctor-0.1.1.dist-info → mlenvdoctor-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mlenvdoctor
-Version: 0.1.1
+Version: 0.1.2
 Summary: Diagnose & fix ML environments for LLM fine-tuning
 Author: ML Environment Doctor Contributors
 License: MIT
@@ -20,6 +20,7 @@ Requires-Python: >=3.8
 Requires-Dist: packaging>=23.0
 Requires-Dist: psutil>=5.9.0
 Requires-Dist: rich>=13.0.0
+Requires-Dist: tomli>=2.0.0; python_version < '3.11'
 Requires-Dist: typer>=0.9.0
 Provides-Extra: dev
 Requires-Dist: black>=23.0.0; extra == 'dev'
@@ -34,7 +35,7 @@ Description-Content-Type: text/markdown
 
 [![Python 3.8+](https://img.shields.io/badge/python-3.8+-blue.svg)](https://www.python.org/downloads/)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](https://opensource.org/licenses/MIT)
-[![PyPI](https://img.shields.io/pypi/v/mlenvdoctor.svg)](https://pypi.org/project/mlenvdoctor/)
+[![PyPI](https://img.shields.io/pypi/v/mlenvdoctor.svg)]([https://pypi.org/project/mlenvdoctor/])
 
 > **Single command fixes 90% of "my torch.cuda.is_available() is False" issues.**
 

mlenvdoctor-0.1.2.dist-info/RECORD

@@ -0,0 +1,21 @@
+mlenvdoctor/__init__.py,sha256=igtJXQ-DiuG4_2BcfW64KPxvhCzSeuXUNUVqDzuLDI8,327
+mlenvdoctor/cli.py,sha256=jvc1MWNciyRGedYCEKSCZeqXfX4wf6xe38Vx44k86xc,7145
+mlenvdoctor/config.py,sha256=c7WUd8XUAo8nz-Ri4TIldc-_p3wiyDNqOsL-MhV9jbI,4770
+mlenvdoctor/constants.py,sha256=ZoYf6-gqmw88CGENyYVur_hCmdcUm0IR5N9pUdOLjJY,1980
+mlenvdoctor/diagnose.py,sha256=7lsAVOPeK2I3nmlzMKmtniyEh1UGv76bJqlqSpjvPL8,20553
+mlenvdoctor/dockerize.py,sha256=AC8HX5sRkSFAM0O0caBnKW4HAdS49MVmMcsplKEDXI4,5562
+mlenvdoctor/exceptions.py,sha256=8wzZE-In0zimXJ3omUA3YmFeklcseO53kqo0SpB58DM,1069
+mlenvdoctor/export.py,sha256=CsuLpbpR2OVs-bud27K8Xv28KCI9veNPIWCsznyFmaw,8671
+mlenvdoctor/fix.py,sha256=fXS4uxBN-FWRFKowmOsdPYI7bnY8jnYxG7UADpJ1hwc,8989
+mlenvdoctor/gpu.py,sha256=sMFgtF4pt-dpOr6IDxvm6f0ChfmCf58E8mApIH6jvAs,6295
+mlenvdoctor/icons.py,sha256=vu35SuBxlZu55rUqldgQfX4UeHeENt082ighZKnzHZY,2289
+mlenvdoctor/logger.py,sha256=OKJQjcdOspARokBcIDyCri14gS_7MlrCt2B3znKc34Y,2377
+mlenvdoctor/parallel.py,sha256=HVJmu8t4k2-XeoQPHlLmhInLOcXjUTiP0VM3THOkAmE,3630
+mlenvdoctor/retry.py,sha256=ZH-KWe6BPNK1sUNFSK2uwdobZ_Z77fxMUfSrSpInt3c,2993
+mlenvdoctor/utils.py,sha256=ehohh-iRLe2qkOMxj5v9yTWONf5gWSdY6CvfrRttTlg,4862
+mlenvdoctor/validators.py,sha256=Kz1FcJM4Cym-S_z5vTocv0cxzKOEgvZIqv8C8c1gSzY,6109
+mlenvdoctor-0.1.2.dist-info/METADATA,sha256=7Rlpv9kjMHQWPBsnZR1gbaUzocWNXNJ3wdGy5UXaNQU,8942
+mlenvdoctor-0.1.2.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+mlenvdoctor-0.1.2.dist-info/entry_points.txt,sha256=Y-WH-ANeiTdECIaqi_EB3ZEf_kACkvsYBHnNhXsCI4k,52
+mlenvdoctor-0.1.2.dist-info/licenses/LICENSE,sha256=rGHdyWGvGWYnEFlthqtB-RtRCTa7WaAOElom5qD-nHw,1114
+mlenvdoctor-0.1.2.dist-info/RECORD,,

mlenvdoctor-0.1.1.dist-info/RECORD

@@ -1,12 +0,0 @@
-mlenvdoctor/__init__.py,sha256=vYK9Wp5kAcHKL7njV76xMfPA24ILdxJvIhoumF6-Sz4,110
-mlenvdoctor/cli.py,sha256=aQ2rpjxsfMDwYAVKsn7cHSar42AZJTrr1tSTl-iM0L4,5488
-mlenvdoctor/diagnose.py,sha256=xa3aqCornGApMJkEWQNIGHwNBRhGA3ud1hBQ6wIVhVQ,17099
-mlenvdoctor/dockerize.py,sha256=AC8HX5sRkSFAM0O0caBnKW4HAdS49MVmMcsplKEDXI4,5562
-mlenvdoctor/fix.py,sha256=P4Qce41LLgjaHugbMFFSg7ldfsSSNDFBz5_T_YA9mig,8945
-mlenvdoctor/gpu.py,sha256=iuiLAW8lZLBpuUL1yapOkx5VYLtY_i1SwK9cE5koZTE,6129
-mlenvdoctor/utils.py,sha256=2gtbiJogEI33IpOLHGEfks6b7Jd1Y7pyfojW9wpYsjU,2893
-mlenvdoctor-0.1.1.dist-info/METADATA,sha256=NQHbKeu7KnZfHWPJcqQEYUXZqIRL7ORCesdMyyqFKU8,8887
-mlenvdoctor-0.1.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-mlenvdoctor-0.1.1.dist-info/entry_points.txt,sha256=Y-WH-ANeiTdECIaqi_EB3ZEf_kACkvsYBHnNhXsCI4k,52
-mlenvdoctor-0.1.1.dist-info/licenses/LICENSE,sha256=rGHdyWGvGWYnEFlthqtB-RtRCTa7WaAOElom5qD-nHw,1114
-mlenvdoctor-0.1.1.dist-info/RECORD,,