pymemora 0.1.0__py3-none-any.whl

memora/__init__.py

@@ -0,0 +1,24 @@
+"""Memora: A memory-aware training optimization framework for PyTorch."""
+
+from memora.config import MemoraConfig, OptimizationRecommendation
+from memora.engine import Memora
+from memora.hardware.gpu_detector import DeviceInfo, GPUDetector, HardwareReport
+from memora.hardware.memory_monitor import MemoryMonitor, MemorySnapshot, MemoryTracker
+from memora.reporting.report import RunReport, compare_runs
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Memora",
+    "MemoraConfig",
+    "OptimizationRecommendation",
+    "GPUDetector",
+    "DeviceInfo",
+    "HardwareReport",
+    "MemoryMonitor",
+    "MemorySnapshot",
+    "MemoryTracker",
+    "RunReport",
+    "compare_runs",
+]
+

memora/cli.py

@@ -0,0 +1,104 @@
+"""Command line interface for Memora using Typer and Rich."""
+
+from __future__ import annotations
+
+import json
+import typer
+from rich.console import Console
+from rich.panel import Panel
+from rich.table import Table
+
+from memora import __version__
+from memora.hardware.gpu_detector import GPUDetector
+from memora.utils.common import format_bytes
+
+app = typer.Typer(help="Memora: PyTorch Memory-Aware Optimization CLI")
+console = Console()
+
+
+@app.command()
+def hardware() -> None:
+    """Detect and display system hardware, CPU RAM, and GPU memory."""
+    report = GPUDetector.get_hardware_report()
+
+    console.print(Panel.fit(
+        f"[bold cyan]Memora Hardware Discovery[/bold cyan]\n"
+        f"OS: {report.platform_system} ({report.platform_release})\n"
+        f"Python: {report.python_version} | PyTorch: {report.torch_version}\n"
+        f"CUDA Available: [{'green' if report.cuda_available else 'yellow'}]{report.cuda_available}[/]"
+        + (f" (v{report.cuda_version})" if report.cuda_version else "")
+        + f"\nSystem RAM: {report.to_dict()['system_ram_available_str']} free / {report.to_dict()['system_ram_total_str']} total",
+        title="System Environment",
+    ))
+
+    table = Table(title="Detected Compute Devices", header_style="bold magenta")
+    table.add_column("Index", style="cyan")
+    table.add_column("Type", style="bold")
+    table.add_column("Device Name", style="white")
+    table.add_column("Total Memory", style="green")
+    table.add_column("Compute Cap.", style="yellow")
+    table.add_column("FP16 / BF16", style="blue")
+
+    for dev in report.devices:
+        half_support = f"FP16: {dev.supports_fp16} | BF16: {dev.supports_bf16}"
+        table.add_row(
+            str(dev.device_index),
+            dev.device_type.upper(),
+            dev.name,
+            format_bytes(dev.total_memory),
+            dev.compute_capability or "N/A",
+            half_support,
+        )
+
+    console.print(table)
+
+
+@app.command()
+def report(file_path: str = typer.Argument(..., help="Path to Memora report JSON file")) -> None:
+    """Read and format a Memora training run report."""
+    try:
+        with open(file_path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+
+        console.print(Panel(
+            f"[bold]Run ID:[/bold] {data.get('run_id')}\n"
+            f"[bold]Timestamp:[/bold] {data.get('timestamp')}\n"
+            f"[bold]Device:[/bold] {data.get('hardware_info', {}).get('default_device', {}).get('name')}",
+            title="[bold green]Report Summary[/bold green]",
+        ))
+
+        opts = data.get("optimization_decisions", {})
+        opt_table = Table(title="Optimization Decisions", header_style="bold cyan")
+        opt_table.add_column("Parameter", style="bold")
+        opt_table.add_column("Value", style="green")
+
+        for k, v in opts.items():
+            if k != "reasons":
+                opt_table.add_row(k, str(v))
+        console.print(opt_table)
+
+        prof = data.get("profiler_summary", {})
+        if prof:
+            prof_table = Table(title="Performance Metrics", header_style="bold yellow")
+            prof_table.add_column("Metric", style="bold")
+            prof_table.add_column("Value", style="cyan")
+            for k, v in prof.items():
+                prof_table.add_row(k, str(v))
+            console.print(prof_table)
+
+    except Exception as e:
+        console.print(f"[bold red]Error loading report:[/bold red] {e}")
+        raise typer.Exit(code=1)
+
+
+@app.command()
+def version() -> None:
+    """Show Memora version."""
+    import torch
+    console.print(f"Memora version: [bold green]{__version__}[/bold green]")
+    console.print(f"PyTorch version: [bold cyan]{torch.__version__}[/bold cyan]")
+
+
+if __name__ == "__main__":
+    app()
+

memora/config.py

@@ -0,0 +1,126 @@
+"""Configuration models and schemas for Memora."""
+
+from __future__ import annotations
+
+from typing import Any, Dict, Literal, Optional, Union
+from pydantic import BaseModel, Field, field_validator
+
+
+PrecisionType = Literal["auto", "fp32", "fp16", "bf16"]
+
+
+class MemoraConfig(BaseModel):
+    """Configuration for Memora optimization and training run."""
+
+    # Device & resource management
+    device: str = Field(
+        default="auto",
+        description="Target compute device: 'auto', 'cuda', 'cuda:0', 'mps', or 'cpu'.",
+    )
+    gpu_memory_limit: Union[float, int, Literal["auto"]] = Field(
+        default="auto",
+        description="VRAM memory limit. If float <= 1.0, fraction of total VRAM. If int > 1, bytes limit. 'auto' uses default safety headroom.",
+    )
+    safety_margin: float = Field(
+        default=0.15,
+        ge=0.01,
+        le=0.50,
+        description="Fraction of memory reserved as a safety buffer against unexpected memory spikes (default: 15%).",
+    )
+
+    # Batch search parameters
+    batch_size: Union[int, Literal["auto"]] = Field(
+        default="auto",
+        description="Micro-batch size for training. 'auto' benchmarks and selects optimal size.",
+    )
+    min_batch_size: int = Field(
+        default=1,
+        ge=1,
+        description="Minimum batch size allowed during search.",
+    )
+    max_batch_size: int = Field(
+        default=512,
+        ge=1,
+        description="Maximum batch size allowed during search.",
+    )
+
+    # Gradient accumulation parameters
+    effective_batch_size: Optional[int] = Field(
+        default=None,
+        ge=1,
+        description="Target effective batch size. When provided with small batch_size, gradient accumulation is automatically configured.",
+    )
+    gradient_accumulation_steps: Union[int, Literal["auto"]] = Field(
+        default="auto",
+        description="Gradient accumulation steps. 'auto' computes steps to match effective_batch_size or 1.",
+    )
+
+    # Precision
+    precision: PrecisionType = Field(
+        default="auto",
+        description="Numerical precision: 'auto', 'fp32', 'fp16', or 'bf16'.",
+    )
+
+    # Activation checkpointing
+    activation_checkpointing: Union[bool, Literal["auto"]] = Field(
+        default="auto",
+        description="Activation checkpointing toggle: 'auto' enables it when VRAM is tight, or explicit boolean.",
+    )
+
+    # Profiling & reporting
+    enable_profiling: bool = Field(
+        default=True,
+        description="Whether to monitor and profile memory and step timings.",
+    )
+    report_output_dir: Optional[str] = Field(
+        default="./memora_reports",
+        description="Directory to persist optimization reports.",
+    )
+    seed: int = Field(
+        default=42,
+        description="Random seed for deterministic runs.",
+    )
+
+    @field_validator("batch_size")
+    @classmethod
+    def validate_batch_size(cls, v: Any) -> Any:
+        if isinstance(v, str):
+            if v.lower() != "auto":
+                raise ValueError("batch_size string value must be 'auto'")
+            return "auto"
+        if isinstance(v, int):
+            if v < 1:
+                raise ValueError("batch_size must be >= 1")
+            return v
+        raise ValueError(f"Invalid batch_size type: {type(v)}")
+
+    @field_validator("gradient_accumulation_steps")
+    @classmethod
+    def validate_grad_accum(cls, v: Any) -> Any:
+        if isinstance(v, str):
+            if v.lower() != "auto":
+                raise ValueError("gradient_accumulation_steps string must be 'auto'")
+            return "auto"
+        if isinstance(v, int):
+            if v < 1:
+                raise ValueError("gradient_accumulation_steps must be >= 1")
+            return v
+        raise ValueError(f"Invalid gradient_accumulation_steps type: {type(v)}")
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert config to dictionary."""
+        return self.model_dump()
+
+
+class OptimizationRecommendation(BaseModel):
+    """Structured recommendation produced by the optimization engine."""
+
+    recommended_batch_size: int
+    recommended_precision: Literal["fp32", "fp16", "bf16"]
+    recommended_grad_accum_steps: int
+    recommended_checkpointing: bool
+    estimated_peak_memory_bytes: int
+    memory_limit_bytes: int
+    reasons: Dict[str, str] = Field(default_factory=dict)
+    probe_results: Dict[str, Any] = Field(default_factory=dict)
+