e2am 0.1.0__py3-none-any.whl

e2am/__init__.py

@@ -0,0 +1,45 @@
+"""E2AM — Energy Efficient AI Models.
+
+Automatic energy, carbon, and performance profiling for AI training with
+almost zero code changes.
+
+Example:
+    >>> from e2am import monitor
+    >>> with monitor(project="ResNet50"):
+    ...     train()
+
+    >>> from e2am import Trainer
+    >>> trainer = Trainer(model=model, optimizer=optimizer,
+    ...                   train_loader=train_loader, val_loader=val_loader)
+    >>> trainer.fit()
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+__version__ = "0.1.0"
+
+__all__ = ["Trainer", "__version__", "monitor"]
+
+if TYPE_CHECKING:  # pragma: no cover - import-time only for type checkers
+    from e2am.monitoring import monitor
+    from e2am.trainer import Trainer
+
+
+def __getattr__(name: str) -> Any:
+    """Lazily resolve the public API.
+
+    Keeps ``import e2am`` fast and torch-free: the ``Trainer`` (which needs
+    PyTorch) is only imported when actually accessed, so monitoring-only
+    installs work on machines without PyTorch.
+    """
+    if name == "monitor":
+        from e2am.monitoring import monitor
+
+        return monitor
+    if name == "Trainer":
+        from e2am.trainer import Trainer
+
+        return Trainer
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

e2am/cli/__init__.py

@@ -0,0 +1,7 @@
+"""The ``e2am`` command line interface (typer application)."""
+
+from __future__ import annotations
+
+from e2am.cli.main import app
+
+__all__ = ["app"]

e2am/cli/loader.py

@@ -0,0 +1,60 @@
+"""Loading user model scripts for the CLI.
+
+``e2am train`` and ``e2am benchmark`` operate on a plain Python file that
+exports well-known factory functions — no subclassing, no registration:
+
+* ``get_model() -> nn.Module`` (required)
+* ``get_loaders() -> (train_loader, val_loader | None)`` (train only)
+* ``get_optimizer(model) -> torch.optim.Optimizer`` (optional; Adam 1e-3)
+* ``get_loss() -> callable`` (optional; CrossEntropyLoss)
+"""
+
+from __future__ import annotations
+
+import importlib.util
+import sys
+from pathlib import Path
+from types import ModuleType
+from typing import Any
+
+from e2am.exceptions import ConfigError
+
+
+def load_user_module(path: str | Path) -> ModuleType:
+    """Import a user's Python file as a module.
+
+    Raises:
+        ConfigError: If the file is missing or fails to import.
+    """
+    path = Path(path)
+    if not path.exists():
+        raise ConfigError(f"Script not found: {path}")
+    spec = importlib.util.spec_from_file_location(f"e2am_user_{path.stem}", path)
+    if spec is None or spec.loader is None:
+        raise ConfigError(f"Cannot import {path} as a Python module.")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    try:
+        spec.loader.exec_module(module)
+    except Exception as exc:
+        raise ConfigError(f"Error while importing {path}: {exc}") from exc
+    return module
+
+
+def get_factory(module: ModuleType, name: str, required: bool = True) -> Any:
+    """Fetch a factory function from the user module.
+
+    Raises:
+        ConfigError: If a required factory is missing or not callable.
+    """
+    factory = getattr(module, name, None)
+    if factory is None:
+        if required:
+            raise ConfigError(
+                f"{module.__name__} must define `{name}()` "
+                f"(see the docs for the e2am script convention)."
+            )
+        return None
+    if not callable(factory):
+        raise ConfigError(f"`{name}` in the script is not callable.")
+    return factory

e2am/cli/main.py

@@ -0,0 +1,334 @@
+"""The ``e2am`` command line interface.
+
+Commands mirror the Python API and share its code paths — the CLI is a thin
+orchestration layer, never a second implementation.
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import typer
+from rich.console import Console
+from rich.table import Table
+
+import e2am
+from e2am.config.settings import ExperimentConfig
+from e2am.exceptions import E2AMError
+
+if TYPE_CHECKING:  # pragma: no cover
+    from e2am.monitoring.result import MonitorResult
+    from e2am.trainer.result import TrainingResult
+
+app = typer.Typer(
+    name="e2am",
+    help="E2AM — Energy Efficient AI Models: automatic energy, carbon, and "
+    "performance profiling for AI training.",
+    no_args_is_help=True,
+    pretty_exceptions_show_locals=False,
+)
+console = Console()
+
+
+def _fail(message: str) -> None:
+    console.print(f"[red]error:[/red] {message}")
+    raise typer.Exit(code=1)
+
+
+@app.command()
+def version() -> None:
+    """Print the E2AM version."""
+    console.print(f"e2am {e2am.__version__}")
+
+
+@app.command()
+def hardware(
+    as_json: bool = typer.Option(False, "--json", help="Emit machine-readable JSON."),
+) -> None:
+    """Detect hardware and its energy-measurement capabilities."""
+    from e2am.utils.hardware import detect_hardware
+
+    info = detect_hardware()
+    if as_json:
+        console.print_json(info.model_dump_json())
+        return
+    table = Table(title="Detected hardware", title_style="bold green")
+    table.add_column("Component", style="bold")
+    table.add_column("Details")
+    table.add_row("Host", f"{info.hostname} · {info.os} {info.os_version}")
+    table.add_row("Python", info.python_version)
+    table.add_row(
+        "CPU",
+        f"{info.cpu.model}\n"
+        f"{info.cpu.physical_cores or '?'} cores / {info.cpu.logical_cores or '?'} threads · "
+        f"TDP {info.cpu.tdp_w:.0f} W (estimation basis)",
+    )
+    table.add_row(
+        "RAM", f"{info.ram.total_gb:.1f} GB · ~{info.ram.estimated_power_w:.1f} W estimated"
+    )
+    for gpu in info.gpus:
+        sensor = (
+            "[green]live power sensor[/green]"
+            if gpu.supports_power_reading
+            else f"[yellow]no power sensor — estimating at {gpu.tdp_w:.0f} W × util[/yellow]"
+        )
+        table.add_row(
+            f"GPU {gpu.index}",
+            f"{gpu.name} · {gpu.total_memory_mb or 0:.0f} MB · "
+            f"driver {gpu.driver_version or '?'}\n{sensor}",
+        )
+    if info.torch_version:
+        cuda = f"CUDA {info.cuda_version}" if info.cuda_available else "CPU-only build"
+        table.add_row("PyTorch", f"{info.torch_version} · {cuda}")
+    else:
+        table.add_row("PyTorch", "[yellow]not installed[/yellow] (monitoring still works)")
+    console.print(table)
+
+
+@app.command()
+def train(
+    script: Path = typer.Argument(..., help="Python file defining get_model() and get_loaders()."),
+    config: Path | None = typer.Option(None, "--config", "-c", help="Experiment config YAML."),
+    epochs: int | None = typer.Option(None, help="Override epochs."),
+    device: str | None = typer.Option(None, help="Override device (cuda/cpu)."),
+    run_name: str | None = typer.Option(None, help="Override run name."),
+) -> None:
+    """Train a model from a script with full E2AM telemetry."""
+    from e2am.cli.loader import get_factory, load_user_module
+    from e2am.trainer import Trainer
+
+    try:
+        module = load_user_module(script)
+        get_model = get_factory(module, "get_model")
+        get_loaders = get_factory(module, "get_loaders")
+        get_optimizer = get_factory(module, "get_optimizer", required=False)
+        get_loss = get_factory(module, "get_loss", required=False)
+
+        experiment = ExperimentConfig.from_yaml(config) if config else ExperimentConfig()
+        model = get_model()
+        loaders = get_loaders()
+        train_loader, val_loader = loaders if isinstance(loaders, tuple) else (loaders, None)
+        if get_optimizer is not None:
+            optimizer = get_optimizer(model)
+        else:
+            import torch
+
+            optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+        trainer = Trainer(
+            model=model,
+            optimizer=optimizer,
+            train_loader=train_loader,
+            val_loader=val_loader,
+            loss_fn=get_loss() if get_loss is not None else None,
+            config=experiment,
+            epochs=epochs,
+            device=device,
+            run_name=run_name,
+        )
+        result = trainer.fit()
+    except E2AMError as exc:
+        _fail(str(exc))
+        return
+    console.print(
+        f"[green]Run complete:[/green] {result.run_name} · status {result.status} · "
+        f"{result.total_energy_wh or 0:.4g} Wh · artifacts in "
+        f"{experiment.output.dir / result.run_name}"
+    )
+
+
+@app.command()
+def benchmark(
+    script: Path = typer.Argument(..., help="Python file defining get_model()."),
+    input_size: str = typer.Option(
+        ..., "--input-size", "-i", help="Input shape incl. batch, e.g. 8,3,224,224"
+    ),
+    device: str | None = typer.Option(None, help="Device (default: auto)."),
+    iterations: int = typer.Option(50, help="Timed iterations."),
+    warmup: int = typer.Option(10, help="Warmup iterations."),
+) -> None:
+    """Profile FLOPs, latency, throughput, and energy per inference."""
+    import torch
+
+    from e2am.cli.loader import get_factory, load_user_module
+    from e2am.monitoring.session import MonitorSession
+    from e2am.profiler import benchmark_latency, profile_model
+
+    try:
+        shape = tuple(int(part) for part in input_size.replace(" ", "").split(","))
+    except ValueError:
+        _fail(f"--input-size must be comma-separated integers, got {input_size!r}")
+        return
+    try:
+        module = load_user_module(script)
+        model = get_factory(module, "get_model")()
+        resolved = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
+        model = model.to(resolved)
+
+        profile = profile_model(model, input_size=shape, device=resolved)
+        session = MonitorSession(project="benchmark", run_name=script.stem)
+        session.config.sampling_interval_s = 0.2
+        session.start()
+        latency = benchmark_latency(
+            model, input_size=shape, device=resolved, warmup=warmup, iterations=iterations
+        )
+        monitor_result = session.stop()
+    except E2AMError as exc:
+        _fail(str(exc))
+        return
+
+    energy_j = monitor_result.total_energy_j
+    per_inference_j = energy_j / iterations if iterations else 0.0
+
+    table = Table(
+        title=f"Benchmark · {profile.model_name} @ {list(shape)} on {latency.device}",
+        title_style="bold green",
+    )
+    table.add_column("Metric", style="bold")
+    table.add_column("Value", justify="right")
+    table.add_row("Parameters", f"{profile.params:,}")
+    table.add_row("MACs / forward", f"{profile.macs:,} ({profile.gmacs:.3f} G)")
+    table.add_row("FLOPs / forward", f"{profile.flops:,} ({profile.gflops:.3f} G)")
+    table.add_row("Model size", f"{profile.model_size_mb:.2f} MB")
+    if profile.coverage < 1.0:
+        table.add_row("MAC coverage", f"[yellow]{profile.coverage * 100:.0f} %[/yellow]")
+    table.add_row("Latency mean", f"{latency.mean_ms:.3f} ms")
+    table.add_row("Latency p50 / p95", f"{latency.p50_ms:.3f} / {latency.p95_ms:.3f} ms")
+    table.add_row("Throughput", f"{latency.throughput_samples_per_s:.1f} samples/s")
+    table.add_row("Energy (whole benchmark)", f"{monitor_result.total_energy_wh:.5f} Wh")
+    table.add_row("Energy / inference", f"{per_inference_j:.4f} J")
+    console.print(table)
+
+
+def _load_result(run_dir: Path) -> TrainingResult | MonitorResult:
+    from e2am.monitoring.result import MonitorResult
+    from e2am.trainer.result import TrainingResult
+
+    metrics = run_dir / "metrics.json"
+    if not metrics.exists():
+        _fail(f"No metrics.json in {run_dir}")
+    data = json.loads(metrics.read_text(encoding="utf-8"))
+    if "epochs_requested" in data:
+        return TrainingResult.model_validate(data)
+    return MonitorResult.model_validate(data)
+
+
+@app.command()
+def report(
+    run_dir: Path = typer.Argument(..., help="A run directory containing metrics.json."),
+    pdf: bool = typer.Option(False, "--pdf", help="Also generate report.pdf."),
+) -> None:
+    """(Re)generate plots and reports for a finished run."""
+    from e2am.config.settings import OutputConfig
+    from e2am.reports.generate import generate_run_artifacts
+
+    result = _load_result(run_dir)
+    written = generate_run_artifacts(result, run_dir, OutputConfig(save_pdf=pdf))
+    for path in written:
+        console.print(f"  [green]wrote[/green] {path}")
+    console.print(f"[green]Regenerated {len(written)} artifact(s) for {result.run_name}.[/green]")
+
+
+@app.command()
+def optimize(
+    run_dir: Path = typer.Argument(..., help="A run directory containing metrics.json."),
+) -> None:
+    """Suggest efficiency optimizations based on a run's measurements."""
+    from e2am.optimize import analyze
+    from e2am.trainer.result import TrainingResult
+
+    result = _load_result(run_dir)
+    if not isinstance(result, TrainingResult):
+        _fail(
+            "optimize needs a training run (metrics.json from e2am.Trainer); "
+            "this directory holds a monitoring-only run."
+        )
+        return
+    suggestions = analyze(result)
+    if not suggestions:
+        console.print(
+            "[green]No optimization opportunities detected[/green] — this run "
+            "already uses the measured hardware efficiently."
+        )
+        return
+    table = Table(
+        title=f"Optimization suggestions · {result.project} / {result.run_name}",
+        title_style="bold green",
+        show_lines=True,
+    )
+    table.add_column("Priority", style="bold", width=8)
+    table.add_column("Suggestion")
+    table.add_column("Estimated savings", justify="right")
+    colors = {"high": "red", "medium": "yellow", "low": "cyan"}
+    for s in suggestions:
+        table.add_row(
+            f"[{colors.get(s.priority, 'white')}]{s.priority}[/]",
+            f"[bold]{s.title}[/bold]\n{s.rationale}\n[dim]{s.action}[/dim]",
+            s.savings_label(),
+        )
+    console.print(table)
+    quantified = sum(s.estimated_savings_wh or 0 for s in suggestions)
+    if quantified > 0:
+        console.print(
+            f"[bold]Quantified savings from this run's own data: " f"~{quantified:.4g} Wh[/bold]"
+        )
+
+
+@app.command()
+def compare(
+    run_dirs: list[Path] = typer.Argument(..., help="Two or more run directories."),
+) -> None:
+    """Compare runs side by side on their headline metrics."""
+    if len(run_dirs) < 2:
+        _fail("compare needs at least two run directories.")
+    flats = []
+    for run_dir in run_dirs:
+        result = _load_result(run_dir)
+        flats.append(result.to_flat_dict())
+
+    keys: list[str] = []
+    for flat in flats:
+        for key in flat:
+            if key not in keys:
+                keys.append(key)
+
+    table = Table(title="Run comparison", title_style="bold green")
+    table.add_column("Metric", style="bold")
+    for flat in flats:
+        table.add_column(str(flat.get("run_name", "?")), justify="right")
+    for key in keys:
+        if key == "run_name":
+            continue
+        values = []
+        for flat in flats:
+            value = flat.get(key)
+            values.append("—" if value in (None, "") else str(value))
+        table.add_row(key, *values)
+    console.print(table)
+
+
+@app.command()
+def dashboard(
+    results_dir: Path = typer.Argument(Path("results"), help="Results root directory."),
+    open_browser: bool = typer.Option(
+        True, "--open/--no-open", help="Open the dashboard in the default browser."
+    ),
+) -> None:
+    """Build a local HTML dashboard of all runs."""
+    from e2am.reports.dashboard import generate_dashboard
+
+    try:
+        path = generate_dashboard(results_dir)
+    except E2AMError as exc:
+        _fail(str(exc))
+        return
+    console.print(f"[green]Dashboard written:[/green] {path}")
+    if open_browser:
+        import webbrowser
+
+        webbrowser.open(path.resolve().as_uri())
+
+
+if __name__ == "__main__":
+    app()

e2am/config/__init__.py

@@ -0,0 +1,21 @@
+"""Typed, YAML-loadable configuration for E2AM experiments."""
+
+from __future__ import annotations
+
+from e2am.config.settings import (
+    CarbonConfig,
+    ExperimentConfig,
+    MonitorConfig,
+    OutputConfig,
+    TrainerConfig,
+    load_config,
+)
+
+__all__ = [
+    "CarbonConfig",
+    "ExperimentConfig",
+    "MonitorConfig",
+    "OutputConfig",
+    "TrainerConfig",
+    "load_config",
+]

e2am/config/settings.py

@@ -0,0 +1,163 @@
+"""Experiment configuration models.
+
+Every knob in E2AM lives in a pydantic model here, giving three things at
+once: validation with helpful errors, YAML round-tripping for reproducible
+runs (``config.yaml`` is written into every results directory), and a single
+documented source of truth for defaults.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime
+from pathlib import Path
+
+import yaml
+from pydantic import BaseModel, Field, field_validator
+
+from e2am.exceptions import ConfigError
+
+#: Global average grid carbon intensity in gCO2eq/kWh (IEA 2023 estimate).
+WORLD_AVG_CARBON_INTENSITY = 475.0
+
+
+class CarbonConfig(BaseModel):
+    """How carbon emissions are computed from energy."""
+
+    country_iso_code: str | None = Field(
+        default=None,
+        description="ISO 3166 country code (e.g. 'IND', 'USA') for region-aware intensity.",
+    )
+    carbon_intensity_g_per_kwh: float = Field(
+        default=WORLD_AVG_CARBON_INTENSITY,
+        gt=0,
+        description="Grid carbon intensity in gCO2eq/kWh. Overridden by "
+        "region lookup when a country code is set and lookup data is available.",
+    )
+
+
+class MonitorConfig(BaseModel):
+    """Background monitoring behavior."""
+
+    sampling_interval_s: float = Field(
+        default=1.0,
+        ge=0.05,
+        le=60.0,
+        description="Seconds between hardware samples.",
+    )
+    gpu_indices: list[int] | None = Field(
+        default=None,
+        description="GPUs to monitor (None = all detected GPUs).",
+    )
+    cpu_tdp_w: float | None = Field(
+        default=None,
+        gt=0,
+        description="CPU TDP override in watts for energy estimation.",
+    )
+    carbon: CarbonConfig = Field(default_factory=CarbonConfig)
+
+
+class OutputConfig(BaseModel):
+    """What gets written to disk after a run."""
+
+    dir: Path = Field(default=Path("results"), description="Root results directory.")
+    save_plots: bool = True
+    save_html: bool = True
+    save_pdf: bool = False
+    save_json: bool = True
+    save_csv: bool = True
+
+    @field_validator("dir", mode="before")
+    @classmethod
+    def _coerce_path(cls, value: str | Path) -> Path:
+        return Path(value)
+
+
+class TrainerConfig(BaseModel):
+    """Training loop behavior for :class:`e2am.Trainer`."""
+
+    epochs: int = Field(default=10, ge=1)
+    device: str | None = Field(
+        default=None,
+        description="Target device ('cuda', 'cuda:1', 'cpu'). None = auto-detect.",
+    )
+    mixed_precision: bool = Field(
+        default=False,
+        description="Enable automatic mixed precision (torch.autocast + GradScaler).",
+    )
+    gradient_accumulation_steps: int = Field(default=1, ge=1)
+    max_grad_norm: float | None = Field(
+        default=None,
+        gt=0,
+        description="Clip gradients to this norm when set.",
+    )
+    log_every_n_steps: int = Field(default=50, ge=1)
+
+
+class ExperimentConfig(BaseModel):
+    """Top-level configuration for one experiment run."""
+
+    project: str = Field(default="e2am", min_length=1)
+    run_name: str | None = Field(
+        default=None,
+        description="Run identifier. None = auto ('<project>-<timestamp>').",
+    )
+    seed: int | None = None
+    tags: list[str] = Field(default_factory=list)
+    notes: str = ""
+    monitor: MonitorConfig = Field(default_factory=MonitorConfig)
+    output: OutputConfig = Field(default_factory=OutputConfig)
+    trainer: TrainerConfig = Field(default_factory=TrainerConfig)
+
+    def resolved_run_name(self) -> str:
+        """Return the explicit run name or generate a timestamped one."""
+        if self.run_name:
+            return self.run_name
+        stamp = datetime.now().strftime("%Y%m%d-%H%M%S")
+        return f"{self.project}-{stamp}"
+
+    def run_dir(self) -> Path:
+        """Directory where this run's artifacts are written."""
+        return self.output.dir / self.resolved_run_name()
+
+    # ------------------------------------------------------------------
+    # YAML round-tripping
+    # ------------------------------------------------------------------
+
+    @classmethod
+    def from_yaml(cls, path: str | Path) -> ExperimentConfig:
+        """Load a configuration from a YAML file.
+
+        Args:
+            path: Path to a YAML file produced by :meth:`to_yaml` or written
+                by hand.
+
+        Raises:
+            ConfigError: If the file is missing, unparsable, or invalid.
+        """
+        path = Path(path)
+        if not path.exists():
+            raise ConfigError(f"Config file not found: {path}")
+        try:
+            data = yaml.safe_load(path.read_text(encoding="utf-8")) or {}
+        except yaml.YAMLError as exc:
+            raise ConfigError(f"Invalid YAML in {path}: {exc}") from exc
+        try:
+            return cls.model_validate(data)
+        except Exception as exc:
+            raise ConfigError(f"Invalid configuration in {path}: {exc}") from exc
+
+    def to_yaml(self, path: str | Path) -> Path:
+        """Write this configuration to a YAML file and return the path."""
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        data = self.model_dump(mode="json")
+        path.write_text(
+            yaml.safe_dump(data, sort_keys=False, default_flow_style=False),
+            encoding="utf-8",
+        )
+        return path
+
+
+def load_config(path: str | Path) -> ExperimentConfig:
+    """Convenience wrapper for :meth:`ExperimentConfig.from_yaml`."""
+    return ExperimentConfig.from_yaml(path)

e2am/exceptions.py

@@ -0,0 +1,45 @@
+"""Exception hierarchy for E2AM.
+
+All E2AM exceptions inherit from :class:`E2AMError`, so callers can catch a
+single base class. Modules raise the most specific subclass available.
+"""
+
+from __future__ import annotations
+
+__all__ = [
+    "ConfigError",
+    "E2AMError",
+    "HardwareError",
+    "MonitorError",
+    "ProfilerError",
+    "ReportError",
+    "TrainerError",
+]
+
+
+class E2AMError(Exception):
+    """Base class for all E2AM errors."""
+
+
+class ConfigError(E2AMError):
+    """Raised when configuration is invalid or cannot be loaded."""
+
+
+class HardwareError(E2AMError):
+    """Raised when hardware detection or access fails unexpectedly."""
+
+
+class MonitorError(E2AMError):
+    """Raised when the monitoring session cannot start or record data."""
+
+
+class ProfilerError(E2AMError):
+    """Raised when model profiling (FLOPs, latency, memory) fails."""
+
+
+class TrainerError(E2AMError):
+    """Raised when the training loop is misconfigured or fails."""
+
+
+class ReportError(E2AMError):
+    """Raised when report generation fails."""