constraintml 0.1.0__py3-none-any.whl

constraintml/__init__.py

@@ -0,0 +1,5 @@
+from .trainers import ConstraintTrainer, GreenTrainer, SmartTrainer, TrainingReport
+
+__all__ = ["ConstraintTrainer", "GreenTrainer", "SmartTrainer", "TrainingReport"]
+
+__version__ = "0.1.0"

constraintml/cloud/__init__.py

@@ -0,0 +1,29 @@
+"""Extension point for optional, provider-specific cloud adapters (AWS/Azure/GCP).
+
+For v2+ - not implemented in v1 - ConstraintML is local-first and must work fully without
+any cloud adapter. This module only documents the shape a future adapter would take
+so that `CarbonModel.source` and cost accounting have a defined place to plug into.
+No network or authentication code lives here.
+"""
+
+from __future__ import annotations
+
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class CloudAdapter(Protocol):
+    """Shape a future cloud adapter would implement. Not used by v1 -- the
+    RuntimeOptimizationEngine never depends on this Protocol being satisfied.
+    """
+
+    def get_carbon_intensity(self, region: str) -> float:
+        """Return regional grid carbon intensity in kg CO2e/kWh."""
+        ...
+
+    def get_energy_price(self, region: str) -> float:
+        """Return regional electricity price in USD/kWh."""
+        ...
+
+
+__all__ = ["CloudAdapter"]

constraintml/constraints/__init__.py

@@ -0,0 +1,3 @@
+from .spec import ConstraintSpec, OPTIMIZE_FOR_CHOICES
+
+__all__ = ["ConstraintSpec", "OPTIMIZE_FOR_CHOICES"]

constraintml/constraints/parsing.py

@@ -0,0 +1,98 @@
+"""Parsing helpers that turn user-facing constraint strings/numbers into typed values.
+
+Every parser accepts either a bare number (assumed to already be in the target unit)
+or a string like "8h", "<50kWh", "25 kg CO2e". A leading comparator ("<", "<=") is
+informational only -- v1 treats every parsed budget as an upper bound regardless of
+the comparator given. A leading ">"/">=" logs a warning since it is nonsensical for
+a resource budget, but the numeric value is still used as an upper bound.
+"""
+
+from __future__ import annotations
+
+import logging
+import re
+from numbers import Number
+
+from .units import CARBON_UNITS_TO_KG, DURATION_UNITS_TO_SECONDS, ENERGY_UNITS_TO_KWH, normalize_unit
+
+logger = logging.getLogger(__name__)
+
+_VALUE_PATTERN = re.compile(r"^\s*([<>]=?)?\s*([\d.]+)\s*(.*?)\s*$")
+
+
+def _split_value_unit(value: str | Number) -> tuple[float, str]:
+    if isinstance(value, Number):
+        return float(value), ""
+
+    match = _VALUE_PATTERN.match(str(value))
+    if not match:
+        raise ValueError(f"Could not parse constraint value: {value!r}")
+
+    comparator, number, unit = match.groups()
+    if comparator in (">", ">="):
+        logger.warning(
+            "Constraint value %r uses comparator %r, which is nonsensical for an upper-bound "
+            "budget; the numeric value will still be treated as an upper bound.",
+            value,
+            comparator,
+        )
+
+    number = float(number)
+    if number < 0:
+        raise ValueError(f"Constraint value must be non-negative, got: {value!r}")
+
+    return number, unit
+
+
+def parse_duration(value: str | Number) -> float:
+    """Parse a duration like "8h", "30m", "1d", or a bare number of seconds -> seconds."""
+    number, unit = _split_value_unit(value)
+    unit = normalize_unit(unit)
+    if unit == "":
+        return number
+    if unit not in DURATION_UNITS_TO_SECONDS:
+        raise ValueError(f"Unknown duration unit {unit!r} in {value!r}")
+    return number * DURATION_UNITS_TO_SECONDS[unit]
+
+
+def parse_energy(value: str | Number) -> float:
+    """Parse an energy budget like "100kWh", "<50kWh", or a bare number -> kWh."""
+    number, unit = _split_value_unit(value)
+    unit = normalize_unit(unit)
+    if unit not in ENERGY_UNITS_TO_KWH:
+        raise ValueError(f"Unknown energy unit {unit!r} in {value!r}")
+    return number * ENERGY_UNITS_TO_KWH[unit]
+
+
+def parse_carbon(value: str | Number) -> float:
+    """Parse a carbon budget like "25 kg CO2e", "8kgCO2e", or a bare number -> kg CO2e."""
+    number, unit = _split_value_unit(value)
+    unit = normalize_unit(unit)
+    if unit not in CARBON_UNITS_TO_KG:
+        raise ValueError(f"Unknown carbon unit {unit!r} in {value!r}")
+    return number * CARBON_UNITS_TO_KG[unit]
+
+
+def parse_cost(value: str | Number) -> float:
+    """Parse a cost budget like "$50", "50 USD", or a bare number -> USD."""
+    if isinstance(value, Number):
+        number = float(value)
+    else:
+        cleaned = str(value).strip().lstrip("$").replace(",", "")
+        cleaned = re.sub(r"(?i)\s*usd\s*$", "", cleaned)
+        number, _unit = _split_value_unit(cleaned)
+    if number < 0:
+        raise ValueError(f"Cost budget must be non-negative, got: {value!r}")
+    return number
+
+
+def parse_percentage(value: str | Number) -> float:
+    """Parse a percentage like "0.5%" -> 0.005. A bare number < 1 is assumed to already be
+    a fraction; a bare number >= 1 is assumed to be a percentage (e.g. 5 -> 0.05)."""
+    number, unit = _split_value_unit(value)
+    unit = normalize_unit(unit)
+    if unit == "%":
+        return number / 100
+    if unit != "":
+        raise ValueError(f"Unknown percentage unit {unit!r} in {value!r}")
+    return number if number < 1 else number / 100

constraintml/constraints/spec.py

@@ -0,0 +1,81 @@
+"""ConstraintSpec: the normalized, internal representation of user-supplied budgets."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from numbers import Number
+
+from .parsing import parse_carbon, parse_cost, parse_duration, parse_energy, parse_percentage
+
+OPTIMIZE_FOR_CHOICES = ("accuracy", "energy", "carbon", "cost", "time")
+
+_DICT_KEY_ALIASES = {
+    "energy": "energy_budget",
+    "energy_budget": "energy_budget",
+    "carbon": "carbon_budget",
+    "carbon_budget": "carbon_budget",
+    "deadline": "deadline",
+    "max_accuracy_loss": "max_accuracy_loss",
+    "accuracy_loss": "max_accuracy_loss",
+    "cost": "cost_budget",
+    "cost_budget": "cost_budget",
+}
+
+
+@dataclass(frozen=True)
+class ConstraintSpec:
+    energy_budget_kwh: float | None = None
+    carbon_budget_kgco2e: float | None = None
+    deadline_seconds: float | None = None
+    max_accuracy_loss: float | None = None
+    cost_budget_usd: float | None = None
+    optimize_for: str = "accuracy"
+
+    def __post_init__(self):
+        if self.optimize_for not in OPTIMIZE_FOR_CHOICES:
+            raise ValueError(
+                f"optimize_for must be one of {OPTIMIZE_FOR_CHOICES}, got {self.optimize_for!r}"
+            )
+
+    @classmethod
+    def from_kwargs(
+        cls,
+        *,
+        energy_budget: str | Number | None = None,
+        carbon_budget: str | Number | None = None,
+        deadline: str | Number | None = None,
+        max_accuracy_loss: str | Number | None = None,
+        cost_budget: str | Number | None = None,
+        optimize_for: str = "accuracy",
+    ) -> "ConstraintSpec":
+        return cls(
+            energy_budget_kwh=parse_energy(energy_budget) if energy_budget is not None else None,
+            carbon_budget_kgco2e=parse_carbon(carbon_budget) if carbon_budget is not None else None,
+            deadline_seconds=parse_duration(deadline) if deadline is not None else None,
+            max_accuracy_loss=parse_percentage(max_accuracy_loss) if max_accuracy_loss is not None else None,
+            cost_budget_usd=parse_cost(cost_budget) if cost_budget is not None else None,
+            optimize_for=optimize_for,
+        )
+
+    @classmethod
+    def from_dict(cls, constraints: dict, optimize_for: str = "accuracy") -> "ConstraintSpec":
+        kwargs: dict = {}
+        for key, value in constraints.items():
+            try:
+                field = _DICT_KEY_ALIASES[key]
+            except KeyError:
+                raise ValueError(f"Unknown constraint key: {key!r}") from None
+            kwargs[field] = value
+        return cls.from_kwargs(optimize_for=optimize_for, **kwargs)
+
+    def merge(self, other: "ConstraintSpec") -> "ConstraintSpec":
+        """Return a spec with `other`'s non-None fields overriding this spec's fields.
+
+        Used to let named kwargs (e.g. GreenTrainer(energy_budget=...)) win over a
+        `constraints={...}` dict passed to the same constructor.
+        """
+        merged = {}
+        for field_name in self.__dataclass_fields__:
+            other_value = getattr(other, field_name)
+            merged[field_name] = other_value if other_value is not None else getattr(self, field_name)
+        return ConstraintSpec(**merged)

constraintml/constraints/units.py

@@ -0,0 +1,54 @@
+"""Unit lookup tables used by constraints.parsing."""
+
+DURATION_UNITS_TO_SECONDS = {
+    "s": 1,
+    "sec": 1,
+    "secs": 1,
+    "second": 1,
+    "seconds": 1,
+    "m": 60,
+    "min": 60,
+    "mins": 60,
+    "minute": 60,
+    "minutes": 60,
+    "h": 3600,
+    "hr": 3600,
+    "hrs": 3600,
+    "hour": 3600,
+    "hours": 3600,
+    "d": 86400,
+    "day": 86400,
+    "days": 86400,
+}
+
+# All keys normalized (lowercase, spaces/dashes/underscores stripped) to kWh.
+ENERGY_UNITS_TO_KWH = {
+    "kwh": 1.0,
+    "wh": 1e-3,
+    "mwh": 1e3,
+    "j": 1 / 3_600_000,
+    "joule": 1 / 3_600_000,
+    "joules": 1 / 3_600_000,
+    "kj": 1e3 / 3_600_000,
+    "mj": 1e6 / 3_600_000,
+    "": 1.0,  # bare number defaults to kWh
+}
+
+# All keys normalized (lowercase, spaces/dashes/underscores stripped) to kg CO2e.
+CARBON_UNITS_TO_KG = {
+    "kgco2e": 1.0,
+    "kgco2": 1.0,
+    "kg": 1.0,
+    "gco2e": 1e-3,
+    "gco2": 1e-3,
+    "g": 1e-3,
+    "tco2e": 1e3,
+    "tco2": 1e3,
+    "t": 1e3,
+    "": 1.0,  # bare number defaults to kg CO2e
+}
+
+
+def normalize_unit(unit: str) -> str:
+    """Lowercase and strip spaces/dashes/underscores so 'kg CO2e' == 'kgCO2e' == 'kg-co2e'."""
+    return unit.strip().lower().replace(" ", "").replace("-", "").replace("_", "")

constraintml/optimization/__init__.py

@@ -0,0 +1,17 @@
+from .carbon import CarbonModel, DEFAULT_GRID_CARBON_INTENSITY_KG_PER_KWH
+from .engine import RuntimeOptimizationEngine
+from .planner import ConstraintPlanner
+from .state import EpochSnapshot, RunState
+from .strategies import ExecutionStrategy, PlannerDecision, StrategyAction
+
+__all__ = [
+    "CarbonModel",
+    "DEFAULT_GRID_CARBON_INTENSITY_KG_PER_KWH",
+    "RuntimeOptimizationEngine",
+    "ConstraintPlanner",
+    "RunState",
+    "EpochSnapshot",
+    "ExecutionStrategy",
+    "PlannerDecision",
+    "StrategyAction",
+]

constraintml/optimization/backends/__init__.py

@@ -0,0 +1,4 @@
+from .base import ExecutionBackend
+from .pytorch_backend import PyTorchBackend
+
+__all__ = ["ExecutionBackend", "PyTorchBackend"]

constraintml/optimization/backends/base.py

@@ -0,0 +1,29 @@
+"""Framework-agnostic execution backend interface.
+
+v1 ships only `PyTorchBackend`. JAX/TensorFlow backends would implement this same
+protocol; none exist yet.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from ..strategies import ExecutionStrategy
+
+
+class ExecutionBackend(ABC):
+    @abstractmethod
+    def training_step(self, batch: Any, loss_fn: Any) -> float:
+        """Run one training step (forward + backward + optimizer update as appropriate
+        given accumulation state) and return the scalar loss value."""
+
+    @abstractmethod
+    def apply_strategy(self, strategy: ExecutionStrategy) -> None:
+        """Reconfigure execution (precision, batch size, grad accumulation,
+        checkpointing) to match the given strategy."""
+
+    @abstractmethod
+    def evaluate(self, val_loader: Any) -> float | None:
+        """Run validation and return an accuracy-like metric (higher is better),
+        or None if it cannot be computed."""

constraintml/optimization/backends/pytorch_backend.py

@@ -0,0 +1,116 @@
+"""PyTorch execution backend: the only real ExecutionBackend implementation for v1."""
+
+from __future__ import annotations
+
+import logging
+
+import torch
+from torch.utils.data import DataLoader, RandomSampler
+
+from ..strategies import ExecutionStrategy
+from .base import ExecutionBackend
+
+logger = logging.getLogger(__name__)
+
+
+class PyTorchBackend(ExecutionBackend):
+    def __init__(self, model, optimizer, train_loader, loss_fn, device=None):
+        self.model = model
+        self.optimizer = optimizer
+        self.train_loader = train_loader
+        self.loss_fn = loss_fn
+        self.device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+
+        initial_batch_size = getattr(train_loader, "batch_size", None) or 0
+        self.strategy = ExecutionStrategy(batch_size=initial_batch_size)
+        self._initial_batch_size = initial_batch_size
+        self._micro_step = 0
+        self._checkpoint_warned = False
+
+        self.optimizer.zero_grad()
+
+    def apply_strategy(self, strategy: ExecutionStrategy) -> None:
+        if strategy.batch_size and strategy.batch_size != self.strategy.batch_size:
+            self._resize_batch(strategy.batch_size)
+        self.strategy = strategy
+
+    def _resize_batch(self, batch_size: int) -> None:
+        dataset = getattr(self.train_loader, "dataset", None)
+        if dataset is None:
+            logger.warning(
+                "Cannot resize batch size to %d: train_loader has no `.dataset` "
+                "(only map-style DataLoaders support runtime batch-size changes in v1).",
+                batch_size,
+            )
+            return
+        shuffled = isinstance(getattr(self.train_loader, "sampler", None), RandomSampler)
+        self.train_loader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=shuffled,
+            num_workers=getattr(self.train_loader, "num_workers", 0),
+        )
+
+    def training_step(self, batch, loss_fn=None) -> float:
+        loss_fn = loss_fn or self.loss_fn
+        inputs, targets = batch
+        inputs = inputs.to(self.device)
+        targets = targets.to(self.device)
+
+        autocast_enabled = self.strategy.precision == "fp16"
+        # CPU has no native fp16 compute kernels -- torch.autocast(dtype=float16) on CPU
+        # falls back to emulation and measured ~44x SLOWER than fp32 for a mid-sized MLP
+        # on this machine, actively defeating the time/energy budget instead of helping it.
+        # bfloat16 is the dtype CPU autocast is actually designed around; fp16 remains
+        # correct for CUDA, where it's backed by real tensor-core throughput.
+        autocast_dtype = torch.bfloat16 if self.device.type == "cpu" else torch.float16
+        with torch.autocast(device_type=self.device.type, enabled=autocast_enabled, dtype=autocast_dtype):
+            outputs = self._forward(inputs)
+            loss = loss_fn(outputs, targets)
+
+        (loss / self.strategy.grad_accumulation_steps).backward()
+        self._micro_step += 1
+
+        if self._micro_step % self.strategy.grad_accumulation_steps == 0:
+            self.optimizer.step()
+            self.optimizer.zero_grad()
+
+        return loss.item()
+
+    def _forward(self, inputs):
+        if self.strategy.activation_checkpointing:
+            if isinstance(self.model, torch.nn.Sequential) and len(self.model) > 1:
+                return torch.utils.checkpoint.checkpoint_sequential(
+                    self.model, len(self.model), inputs, use_reentrant=False
+                )
+            if not self._checkpoint_warned:
+                logger.warning(
+                    "Activation checkpointing requested but the model is not an "
+                    "nn.Sequential, so its forward pass cannot be automatically "
+                    "wrapped in v1. Continuing without checkpointing."
+                )
+                self._checkpoint_warned = True
+        return self.model(inputs)
+
+    def evaluate(self, val_loader) -> float | None:
+        if val_loader is None or self.loss_fn is None:
+            return None
+
+        self.model.eval()
+        total_loss = 0.0
+        total_batches = 0
+        with torch.no_grad():
+            for inputs, targets in val_loader:
+                inputs = inputs.to(self.device)
+                targets = targets.to(self.device)
+                outputs = self.model(inputs)
+                total_loss += self.loss_fn(outputs, targets).item()
+                total_batches += 1
+        self.model.train()
+
+        if total_batches == 0:
+            return None
+        # "Accuracy-like" metric where higher is better; without a task-specific
+        # metric available, negative average validation loss becomes the generic proxy.
+        return -(total_loss / total_batches)

constraintml/optimization/carbon.py

@@ -0,0 +1,20 @@
+"""Carbon accounting: converts energy consumption into estimated carbon emissions."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# Rough global-average grid carbon intensity. This is a coarse approximation --
+# override with a locally accurate figure via `carbon_intensity_kg_per_kwh=` on any
+# trainer. A future `constraintml.cloud` adapter is the intended extension point for
+# sourcing a regionally accurate figure automatically (see CarbonModel.source).
+DEFAULT_GRID_CARBON_INTENSITY_KG_PER_KWH = 0.4
+
+
+@dataclass
+class CarbonModel:
+    intensity_kg_per_kwh: float = DEFAULT_GRID_CARBON_INTENSITY_KG_PER_KWH
+    source: str = "static_default"
+
+    def to_carbon(self, energy_kwh: float) -> float:
+        return energy_kwh * self.intensity_kg_per_kwh

constraintml/optimization/engine.py

@@ -0,0 +1,76 @@
+"""RuntimeOptimizationEngine: orchestrates telemetry, carbon accounting, and the
+planner around each training step, applying any resulting strategy to the backend.
+"""
+
+from __future__ import annotations
+
+import time
+
+from .backends.base import ExecutionBackend
+from .carbon import CarbonModel
+from .planner import ConstraintPlanner
+from .state import RunState
+from .strategies import StrategyAction
+from .telemetry.base import TelemetryCollector, TelemetrySample
+
+
+class RuntimeOptimizationEngine:
+    def __init__(
+        self,
+        planner: ConstraintPlanner,
+        telemetry: TelemetryCollector,
+        backend: ExecutionBackend,
+        carbon_model: CarbonModel | None = None,
+        cost_per_kwh_usd: float | None = None,
+        state: RunState | None = None,
+    ):
+        self.planner = planner
+        self.telemetry = telemetry
+        self.backend = backend
+        self.carbon_model = carbon_model or CarbonModel()
+        self.cost_per_kwh_usd = cost_per_kwh_usd
+        self.state = state or RunState()
+
+        self._step_start_time: float | None = None
+        self._step_start_sample: TelemetrySample | None = None
+
+    def begin_step(self) -> None:
+        self._step_start_time = time.perf_counter()
+        self._step_start_sample = self.telemetry.sample()
+
+    def end_step(self, loss: float | None = None) -> None:
+        if self._step_start_time is None or self._step_start_sample is None:
+            raise RuntimeError("end_step() called without a matching begin_step().")
+
+        end_sample = self.telemetry.sample()
+        elapsed = time.perf_counter() - self._step_start_time
+
+        energy_kwh = self.telemetry.integrate_energy(self._step_start_sample, end_sample, elapsed)
+        carbon_kgco2e = self.carbon_model.to_carbon(energy_kwh)
+
+        self.state.cumulative_energy_kwh += energy_kwh
+        self.state.cumulative_carbon_kgco2e += carbon_kgco2e
+        if self.cost_per_kwh_usd is not None:
+            self.state.cumulative_cost_usd += energy_kwh * self.cost_per_kwh_usd
+        self.state.elapsed_seconds += elapsed
+        self.state.step_count += 1
+
+        self._step_start_time = None
+        self._step_start_sample = None
+
+        if self.planner.should_evaluate(self.state):
+            decision = self.planner.decide(self.state)
+            self._apply_decision(decision)
+
+    def _apply_decision(self, decision) -> None:
+        if decision.action is StrategyAction.EARLY_STOP:
+            self.state.stopped_early = True
+            self.state.stop_reason = decision.reason
+            return
+        if decision.action is StrategyAction.NO_CHANGE:
+            return
+        self.backend.apply_strategy(decision.resulting_strategy)
+        self.state.record_strategy(decision.resulting_strategy)
+
+    def should_stop(self) -> bool:
+        return self.state.stopped_early