openg2g.controller¶

class Controller(Generic[DCBackendT, GridBackendT], ABC):
    """Interface for a control component in the G2G framework.

    Controllers receive datacenter and grid state and produce control actions.
    Multiple controllers compose in order within the coordinator.
    """

    _dc_types: tuple[type[DatacenterBackend], ...] = (DatacenterBackend,)
    _grid_types: tuple[type[GridBackend], ...] = (GridBackend,)

    def __init_subclass__(cls, **kwargs: object) -> None:
        super().__init_subclass__(**kwargs)
        dc_types: tuple[type[DatacenterBackend], ...] | None = None
        grid_types: tuple[type[GridBackend], ...] | None = None
        for base in getattr(cls, "__orig_bases__", ()):
            if get_origin(base) is Controller:
                args = get_args(base)
                if len(args) != 2:
                    raise TypeError(
                        f"{cls.__name__} must specialize Controller with two generic args: "
                        "Controller[DatacenterType, GridType]."
                    )
                dc_raw, grid_raw = args
                dc_norm = _normalize_backend_type_arg(dc_raw, required_base=DatacenterBackend)
                grid_norm = _normalize_backend_type_arg(grid_raw, required_base=GridBackend)
                dc_types = tuple(t for t in dc_norm if issubclass(t, DatacenterBackend))
                grid_types = tuple(t for t in grid_norm if issubclass(t, GridBackend))
                break

        if dc_types is None or grid_types is None:
            inherited = [b for b in cls.__bases__ if issubclass(b, Controller)]
            inherited = [b for b in inherited if b is not Controller]
            if inherited:
                parent = inherited[0]
                cls._dc_types = parent.compatible_datacenter_types()
                cls._grid_types = parent.compatible_grid_types()
                return
            raise TypeError(
                f"{cls.__name__} must explicitly specialize Controller generics as "
                "Controller[DatacenterType, GridType]."
            )

        cls._dc_types = dc_types
        cls._grid_types = grid_types

    @final
    @classmethod
    def compatible_datacenter_types(cls) -> tuple[type[DatacenterBackend], ...]:
        return cls._dc_types

    @final
    @classmethod
    def compatible_grid_types(cls) -> tuple[type[GridBackend], ...]:
        return cls._grid_types

    @final
    @classmethod
    def compatibility_signature(cls) -> str:
        dc = " | ".join(t.__name__ for t in cls.compatible_datacenter_types())
        grid = " | ".join(t.__name__ for t in cls.compatible_grid_types())
        return f"Controller[{dc}, {grid}]"

    @property
    @abstractmethod
    def dt_s(self) -> Fraction:
        """Control interval as a Fraction (seconds)."""

    @abstractmethod
    def reset(self) -> None:
        """Reset simulation state to initial conditions.

        Called by the coordinator before each [`start`][..start]. Must
        clear all simulation state: dual variables, counters, cached
        matrices. Configuration (dt_s, fits, step sizes) is not
        affected.

        Abstract so every implementation explicitly enumerates its state.
        A forgotten field is a bug -- not clearing it silently corrupts
        the second run.
        """

    def start(self) -> None:
        """Acquire per-run resources.

        Called after [`reset`][..reset], before the simulation loop.
        No-op by default because most controllers have no resources to
        acquire.
        """

    def stop(self) -> None:
        """Release per-run resources. Simulation state is preserved.

        Called after the simulation loop in LIFO order. No-op by default.
        """

    @abstractmethod
    def step(
        self,
        clock: SimulationClock,
        datacenter: DCBackendT,
        grid: GridBackendT,
        events: EventEmitter,
    ) -> list[DatacenterCommand | GridCommand]:
        """Compute control commands for this step. Return an empty list for no-op."""

@dataclass(frozen=True)
class BatchSizeChange:
    """A batch size change event, optionally with gradual ramp-up.

    Attributes:
        batch_size: Target batch size (max_num_seqs).
        ramp_up_rate: Requests/second ramp-up rate. 0 means immediate.
    """

    batch_size: int
    ramp_up_rate: float = 0.0

    def __post_init__(self) -> None:
        if self.batch_size <= 0:
            raise ValueError(f"batch_size must be positive, got {self.batch_size}.")
        if self.ramp_up_rate < 0:
            raise ValueError(f"ramp_up_rate must be >= 0, got {self.ramp_up_rate}.")

    def at(self, t: float) -> BatchSizeSchedule:
        """Schedule this change at time *t* seconds.

        Returns:
            A single-entry [`BatchSizeSchedule`][...BatchSizeSchedule].
        """
        return BatchSizeSchedule(((t, self),))

class BatchSizeSchedule:
    """Ordered sequence of batch size changes, built with `|` operator.

    Example:

    ```python
    schedule = (
        BatchSizeChange(48).at(40)
        | BatchSizeChange(32).at(60)
        | BatchSizeChange(48, ramp_up_rate=4).at(280)
    )
    ```

    Raises:
        ValueError: If two entries share the same timestamp.
    """

    __slots__ = ("_entries",)

    def __init__(self, entries: tuple[tuple[float, BatchSizeChange], ...]) -> None:
        self._entries = tuple(sorted(entries, key=lambda e: e[0]))
        times = [t for t, _ in self._entries]
        if len(times) != len(set(times)):
            seen: set[float] = set()
            dupes = sorted({t for t in times if t in seen or seen.add(t)})
            raise ValueError(f"BatchSizeSchedule has duplicate timestamps: {dupes}")

    def __or__(self, other: BatchSizeSchedule) -> BatchSizeSchedule:
        return BatchSizeSchedule(self._entries + other._entries)

    def __iter__(self) -> Iterator[tuple[float, BatchSizeChange]]:
        return iter(self._entries)

    def __len__(self) -> int:
        return len(self._entries)

    def __bool__(self) -> bool:
        return bool(self._entries)

    def __repr__(self) -> str:
        parts: list[str] = []
        for t, c in self._entries:
            ramp = f", ramp_up_rate={c.ramp_up_rate}" if c.ramp_up_rate > 0 else ""
            parts.append(f"BatchSizeChange({c.batch_size}{ramp}).at(t={t})")
        return " | ".join(parts)

class BatchSizeScheduleController(Controller[DatacenterBackend, GridBackend]):
    """Applies pre-defined batch size changes at scheduled times.

    Walks each model's schedule and emits
    [`SetBatchSize`][openg2g.datacenter.command.SetBatchSize] commands when the
    simulation clock reaches the scheduled time.

    Args:
        schedules: Per-model batch size schedules, keyed by model label.
        dt_s: How often the controller checks the schedule (seconds).
    """

    def __init__(
        self,
        *,
        schedules: dict[str, BatchSizeSchedule],
        dt_s: Fraction = Fraction(1),
    ) -> None:
        self._dt_s = dt_s
        self._schedules = dict(schedules)
        self._indices: dict[str, int] = {label: 0 for label in schedules}

    def reset(self) -> None:
        self._indices = {label: 0 for label in self._schedules}

    @property
    def dt_s(self) -> Fraction:
        return self._dt_s

    def step(
        self,
        clock: SimulationClock,
        datacenter: DatacenterBackend,
        grid: GridBackend,
        events: EventEmitter,
    ) -> list[DatacenterCommand | GridCommand]:
        t_now = clock.time_s
        batch_changes: dict[str, int] = {}
        ramp_rates: dict[str, float] = {}

        for label, schedule in self._schedules.items():
            entries = list(schedule)
            idx = self._indices[label]

            while idx < len(entries):
                t_ev, change = entries[idx]
                if float(t_ev) <= t_now + 1e-12:
                    batch_changes[label] = change.batch_size
                    if change.ramp_up_rate > 0:
                        ramp_rates[label] = change.ramp_up_rate
                    idx += 1
                else:
                    break

            self._indices[label] = idx

        if batch_changes:
            events.emit(
                "controller.batch_schedule.fired",
                {"batch_size_by_model": batch_changes},
            )
            return [
                SetBatchSize(
                    batch_size_by_model=batch_changes,
                    ramp_up_rate_by_model=ramp_rates,
                )
            ]
        return []

class LoadShiftConfig(BaseModel):
    """Configuration for cross-site load shifting."""

    enabled: bool = False
    gpus_per_shift: int = 8
    headroom: float = 0.3
    """Fraction of extra server capacity to pre-allocate at each DC
    so incoming replicas have room (e.g. 0.3 = 30% headroom)."""

class LoadShiftController(Controller[LLMBatchSizeControlledDatacenter, OpenDSSGrid]):
    """Shift LLM replicas between datacenters to resolve voltage violations.

    Rules:
    1. Only shift models already running at both source and destination.
    2. Only act when batch sizes are saturated AND violation persists.
    3. For undervoltage: shift load OUT of violated site → highest-voltage site.
       For overvoltage: shift load INTO violated site ← lowest-voltage site.
    4. Shift ``gpus_per_shift`` GPUs worth of replicas per time step.
    5. Repeat until violation resolves or no candidates remain.
    """

    def __init__(
        self,
        *,
        config: LoadShiftConfig,
        dt_s: Fraction,
        datacenters: dict[str, LLMBatchSizeControlledDatacenter],
        site_bus_map: dict[str, str],
        models_by_site: dict[str, list[str]],
        gpus_per_replica_by_model: dict[str, int],
        feasible_batch_sizes_by_model: dict[str, list[int]],
        v_min: float = 0.95,
        v_max: float = 1.05,
    ) -> None:
        self._config = config
        self._dt_s = dt_s
        self._datacenters = datacenters  # site_id -> datacenter
        self._site_bus_map = site_bus_map  # site_id -> bus name
        self._models_by_site = models_by_site  # site_id -> [model_labels]
        self._gpus_per_replica = gpus_per_replica_by_model
        self._feasible_bs = feasible_batch_sizes_by_model
        self._v_min = v_min
        self._v_max = v_max
        self._step_count = 0

    @property
    def dt_s(self) -> Fraction:
        return self._dt_s

    @property
    def site_id(self) -> str | None:
        return None  # cross-site controller

    def step(
        self,
        clock: SimulationClock,
        datacenter: LLMBatchSizeControlledDatacenter,
        grid: OpenDSSGrid,
        events: EventEmitter,
    ) -> list[DatacenterCommand | GridCommand]:
        self._step_count += 1
        if not self._config.enabled:
            return []

        # Build bus -> min voltage mapping from grid
        voltages = grid.voltages_vector()
        v_index = grid.v_index  # list of (bus, phase)
        bus_voltages: dict[str, list[float]] = {}
        for (bus, _phase), v in zip(v_index, voltages, strict=False):
            bus_voltages.setdefault(bus.lower(), []).append(float(v))

        # Per-site min/max voltage
        site_vmin: dict[str, float] = {}
        site_vmax: dict[str, float] = {}
        for site_id, bus in self._site_bus_map.items():
            vs = bus_voltages.get(bus.lower(), [])
            if vs:
                site_vmin[site_id] = min(vs)
                site_vmax[site_id] = max(vs)

        commands: list[DatacenterCommand | GridCommand] = []

        # Check each site for violations
        for site_id in list(self._site_bus_map.keys()):
            vmin = site_vmin.get(site_id, 1.0)
            vmax = site_vmax.get(site_id, 1.0)

            is_undervoltage = vmin < self._v_min
            is_overvoltage = vmax > self._v_max

            if not is_undervoltage and not is_overvoltage:
                continue

            # Check if batch sizes are saturated at this site
            if not self._is_batch_saturated(site_id, datacenter, grid, is_undervoltage):
                continue

            # Find best destination and model to shift
            site_models = set(self._models_by_site.get(site_id, []))

            if is_undervoltage:
                # Shift load OUT: pick destination with highest voltage AND available capacity
                best_dest = None
                best_v = -1.0
                for other_id in self._site_bus_map:
                    if other_id == site_id:
                        continue
                    other_models = set(self._models_by_site.get(other_id, []))
                    shared = site_models & other_models
                    if not shared:
                        continue
                    dest_dc = self._datacenters.get(other_id)
                    if dest_dc is not None and dest_dc.available_gpu_capacity() < self._config.gpus_per_shift:
                        continue  # no room at this destination
                    ov = site_vmin.get(other_id, 0.0)
                    if ov > best_v:
                        best_v = ov
                        best_dest = other_id

                if best_dest is None:
                    continue

                shared_models = site_models & set(self._models_by_site.get(best_dest, []))
                model = self._pick_model(shared_models)
                if model is None:
                    continue

                replicas = max(1, self._config.gpus_per_shift // self._gpus_per_replica[model])
                commands.append(
                    ShiftReplicas(
                        model_label=model,
                        replica_delta=-replicas,
                        target_site_id=site_id,
                    )
                )
                commands.append(
                    ShiftReplicas(
                        model_label=model,
                        replica_delta=+replicas,
                        target_site_id=best_dest,
                    )
                )
                logger.info(
                    "LoadShift: undervoltage at %s (Vmin=%.4f), shift %s ×%d replicas -> %s (Vmin=%.4f, free=%d GPUs)",
                    site_id,
                    vmin,
                    model,
                    replicas,
                    best_dest,
                    best_v,
                    self._datacenters[best_dest].available_gpu_capacity() if best_dest in self._datacenters else -1,
                )

            elif is_overvoltage:
                # Shift load IN: check this site has capacity, pick source with lowest voltage
                site_dc = self._datacenters.get(site_id)
                if site_dc is not None and site_dc.available_gpu_capacity() < self._config.gpus_per_shift:
                    continue  # violated site is full, can't accept more load

                best_src = None
                best_v = 2.0
                for other_id in self._site_bus_map:
                    if other_id == site_id:
                        continue
                    other_models = set(self._models_by_site.get(other_id, []))
                    shared = site_models & other_models
                    if not shared:
                        continue
                    ov = site_vmax.get(other_id, 2.0)
                    if ov < best_v:
                        best_v = ov
                        best_src = other_id

                if best_src is None:
                    continue

                shared_models = site_models & set(self._models_by_site.get(best_src, []))
                model = self._pick_model(shared_models)
                if model is None:
                    continue

                replicas = max(1, self._config.gpus_per_shift // self._gpus_per_replica[model])
                commands.append(
                    ShiftReplicas(
                        model_label=model,
                        replica_delta=-replicas,
                        target_site_id=best_src,
                    )
                )
                commands.append(
                    ShiftReplicas(
                        model_label=model,
                        replica_delta=+replicas,
                        target_site_id=site_id,
                    )
                )
                logger.info(
                    "LoadShift: overvoltage at %s (Vmax=%.4f), shift %s ×%d replicas <- %s (Vmax=%.4f)",
                    site_id,
                    vmax,
                    model,
                    replicas,
                    best_src,
                    best_v,
                )

        return commands

    def _is_batch_saturated(
        self,
        site_id: str,
        datacenter: LLMBatchSizeControlledDatacenter,
        grid: OpenDSSGrid,
        is_undervoltage: bool,
    ) -> bool:
        """Check if all models at site have batch sizes at their limit.

        For undervoltage: saturated = all at minimum batch (can't reduce power further
        via OFO — batch is already at min so load can't be lowered).
        For overvoltage: saturated = all at maximum batch (can't increase power further).
        """
        dc = self._datacenters.get(site_id)
        if dc is None:
            return False
        state = dc.state
        if state is None:
            return False

        site_models = self._models_by_site.get(site_id, [])
        for model_label in site_models:
            current_bs = state.batch_size_by_model.get(model_label)
            feasible = self._feasible_bs.get(model_label, [])
            if not feasible or current_bs is None:
                continue
            if is_undervoltage:
                # Saturated = at minimum batch (OFO can't reduce power further)
                if current_bs > min(feasible):
                    return False  # Still room to reduce
            else:
                # Saturated = at maximum batch (OFO can't increase power further)
                if current_bs < max(feasible):
                    return False  # Still room to increase
        return True

    def _pick_model(self, shared_models: set[str]) -> str | None:
        """Pick the model with the most GPUs per replica (largest power impact)."""
        if not shared_models:
            return None
        return max(shared_models, key=lambda m: self._gpus_per_replica.get(m, 1))

    def reset(self) -> None:
        self._step_count = 0

    def start(self) -> None:
        pass

    def stop(self) -> None:
        pass

class NoopController(Controller[DatacenterBackend, GridBackend]):
    """Controller that always returns an empty action."""

    def __init__(self, dt_s: Fraction = Fraction(1)) -> None:
        self._dt_s = dt_s

    @property
    def dt_s(self) -> Fraction:
        return self._dt_s

    def reset(self) -> None:
        pass

    def step(
        self,
        clock: SimulationClock,
        datacenter: DatacenterBackend,
        grid: GridBackend,
        events: EventEmitter,
    ) -> list[DatacenterCommand | GridCommand]:
        return []

class OFOConfig(BaseModel):
    """Online Feedback Optimization tuning parameters.

    Attributes:
        primal_step_size: Primal descent step size ρ_x (Eq. 8).
        w_throughput: Throughput weight in primal gradient.
        w_switch: Switching cost regularizer weight γ (Eq. 4a).
        voltage_gradient_scale: Scaling factor k_v for voltage dual term
            in the primal gradient.
        v_min: Lower voltage bound (pu).
        v_max: Upper voltage bound (pu).
        voltage_dual_step_size: Voltage dual ascent step size ρ_v (Eqs. 5-6).
        latency_dual_step_size: Latency dual ascent step size ρ_l (Eq. 7).
        sensitivity_update_interval: Steps between H-matrix re-estimation
            (0 = only once at init).
        sensitivity_perturbation_kw: Perturbation magnitude (kW) for
            finite-difference sensitivity estimation.
    """

    model_config = ConfigDict(frozen=True)

    # Primal
    primal_step_size: float = 0.05
    w_throughput: float = 0.1
    w_switch: float = 0.0
    voltage_gradient_scale: float = 1e6

    # Dual
    v_min: float = 0.95
    v_max: float = 1.05
    voltage_dual_step_size: float = 0.5
    latency_dual_step_size: float = 1.0

    # Sensitivity
    sensitivity_update_interval: int = 0
    sensitivity_perturbation_kw: float = 100.0

class LogisticModelStore:
    """Per-model logistic models for power, latency, and throughput.

    Used by
    [`OFOBatchSizeController`][openg2g.controller.ofo.OFOBatchSizeController]
    to compute gradients of the Lagrangian with respect to batch size.

    Attributes:
        COL_MODEL_LABEL: Column name for model label in the CSV.
        COL_METRIC: Column name for metric type in the CSV.
    """

    COL_MODEL_LABEL = "model_label"
    COL_METRIC = "metric"

    def __init__(
        self,
        power: dict[str, LogisticModel],
        latency: dict[str, LogisticModel],
        throughput: dict[str, LogisticModel],
    ) -> None:
        self._power = dict(power)
        self._latency = dict(latency)
        self._throughput = dict(throughput)
        self._by_batch: dict[str, dict[int, list[tuple[float, float, float]]]] | None = None

    def power(self, model: str) -> LogisticModel:
        """Return the power logistic model for a model label."""
        return self._power[model]

    def latency(self, model: str) -> LogisticModel:
        """Return the latency logistic model for a model label."""
        return self._latency[model]

    def throughput(self, model: str) -> LogisticModel:
        """Return the throughput logistic model for a model label."""
        return self._throughput[model]

    @property
    def power_fits(self) -> dict[str, LogisticModel]:
        return dict(self._power)

    @property
    def latency_fits(self) -> dict[str, LogisticModel]:
        return dict(self._latency)

    @property
    def throughput_fits(self) -> dict[str, LogisticModel]:
        return dict(self._throughput)

    @classmethod
    def generate(
        cls,
        models: tuple[InferenceModelSpec, ...],
        data_sources: dict[str, Any],
        *,
        runs: Any = None,
        mlenergy_data_dir: Path | None = None,
    ) -> LogisticModelStore:
        """Generate logistic fits from ML.ENERGY benchmark data.

        Args:
            models: Model specifications.
            data_sources: Per-model `MLEnergySource` instances, keyed by
                `model_label`.
            runs: Pre-loaded `LLMRuns` object. If `None`, loads from
                `mlenergy_data_dir` or the HuggingFace Hub.
            mlenergy_data_dir: Path to compiled mlenergy-data directory.
                Ignored if `runs` is provided.

        Returns:
            A new `LogisticModelStore` with fitted logistic models.
        """
        if runs is None:
            unique_tasks = {src.task for src in data_sources.values()}
            if mlenergy_data_dir:
                runs = LLMRuns.from_directory(str(mlenergy_data_dir), stable_only=False).task(*unique_tasks)
            else:
                runs = LLMRuns.from_hf(stable_only=False).task(*unique_tasks)
        if not runs:
            raise ValueError("No runs found for the specified tasks")

        subsets_by_label: dict[str, Any] = {}
        for ms in models:
            src = data_sources.get(ms.model_label)
            if src is None:
                raise ValueError(f"No data source for model {ms.model_label!r}")
            model_id = ms.model_id
            if not model_id:
                raise ValueError(f"model_id is required for data generation (model={ms.model_label!r})")

            subset = (
                runs.model_id(model_id).gpu_model(src.gpu).num_gpus(ms.gpus_per_replica).max_num_seqs(*src.batch_sizes)
            )
            if not subset:
                raise ValueError(
                    f"Config matched zero runs for logistic fits: model_id={model_id!r}, "
                    f"gpu={src.gpu!r}, num_gpus={ms.gpus_per_replica}, "
                    f"batch_sizes={src.batch_sizes}"
                )
            subsets_by_label[ms.model_label] = subset

        all_by_batch: dict[str, dict[int, list[tuple[float, float, float]]]] = {}
        power: dict[str, LogisticModel] = {}
        latency: dict[str, LogisticModel] = {}
        throughput: dict[str, LogisticModel] = {}
        for model_label, group in subsets_by_label.items():
            exclude = set(data_sources[model_label].fit_exclude_batch_sizes)
            by_batch: dict[int, list[tuple[float, float, float]]] = {}
            for r in group:
                if r.max_num_seqs in exclude:
                    continue
                by_batch.setdefault(r.max_num_seqs, []).append(
                    (r.avg_power_watts, r.mean_itl_ms / 1000.0, r.output_throughput_tokens_per_sec)
                )
            all_by_batch[model_label] = by_batch

            batches = sorted(by_batch.keys())
            if not batches:
                continue

            x = np.log2(np.array(batches, dtype=float).clip(min=1))
            for _metric_name, idx, target in [
                ("power", 0, power),
                ("latency", 1, latency),
                ("throughput", 2, throughput),
            ]:
                y = np.array([float(np.median([t[idx] for t in by_batch[b]])) for b in batches])
                fit = LogisticModel.fit(x, y)
                target[model_label] = fit

        if not power and not latency and not throughput:
            raise ValueError("No logistic fit rows produced")
        store = cls(power=power, latency=latency, throughput=throughput)
        store._by_batch = all_by_batch
        return store

    def save(self, csv_path: Path, *, plot: bool = False) -> None:
        """Save logistic fits to a CSV.

        Args:
            csv_path: Output CSV path.
            plot: If `True`, also write a logistic fits plot to the
                same directory.
        """
        csv_path = Path(csv_path)
        csv_path.parent.mkdir(parents=True, exist_ok=True)
        rows: list[dict[str, Any]] = []
        for metric_name, fits in [("power", self._power), ("latency", self._latency), ("throughput", self._throughput)]:
            for label in sorted(fits):
                model = fits[label]
                rows.append(
                    {
                        self.COL_MODEL_LABEL: label,
                        self.COL_METRIC: metric_name,
                        "L": model.L,
                        "x0": model.x0,
                        "k": model.k,
                        "b0": model.b0,
                    }
                )
        pd.DataFrame(rows).to_csv(csv_path, index=False)

        by_batch = getattr(self, "_by_batch", None)
        if plot and by_batch is not None:
            model_labels = sorted(self._power.keys())
            _plot_logistic_fits(
                by_batch,
                self._power,
                self._latency,
                self._throughput,
                model_labels,
                csv_path.parent,
            )

    @classmethod
    def load(cls, csv_path: Path | str) -> LogisticModelStore:
        """Load power, latency, and throughput fits from a merged CSV.

        Expected columns: `model_label`, `metric`, plus the logistic
        model parameter columns (`L`, `x0`, `k`, `b0`).

        The `metric` column must contain `power`, `latency`, or
        `throughput` (case-insensitive).

        Args:
            csv_path: Path to the logistic fits CSV.
        """
        csv_path = Path(csv_path)
        df = pd.read_csv(csv_path)

        required_cols = [cls.COL_MODEL_LABEL, cls.COL_METRIC]
        missing = [c for c in required_cols if c not in df.columns]
        if missing:
            raise ValueError(f"{csv_path} missing columns: {missing}. Got: {list(df.columns)}")

        power: dict[str, LogisticModel] = {}
        latency: dict[str, LogisticModel] = {}
        throughput: dict[str, LogisticModel] = {}
        targets = {"power": power, "latency": latency, "throughput": throughput}
        for row in df.to_dict(orient="records"):
            metric = str(row[cls.COL_METRIC]).strip().lower()
            if metric in targets:
                targets[metric][str(row[cls.COL_MODEL_LABEL])] = LogisticModel.from_dict(row)

        if not power and not latency and not throughput:
            raise ValueError(f"No logistic model rows loaded from {csv_path}")
        return cls(power=power, latency=latency, throughput=throughput)

    @classmethod
    def ensure(
        cls,
        csv_path: Path,
        models: tuple[InferenceModelSpec, ...] | None = None,
        data_sources: dict[str, Any] | None = None,
        *,
        mlenergy_data_dir: Path | None = None,
        plot: bool = False,
    ) -> LogisticModelStore:
        """Load from `csv_path`, generating first if needed.

        Args:
            csv_path: Path to the logistic fits CSV.
            models: Model specifications. Required when no cached file exists.
            data_sources: Per-model `MLEnergySource` instances, keyed by
                `model_label`. Required when no cached file exists.
            mlenergy_data_dir: Path to compiled mlenergy-data directory.
            plot: If `True`, generate a logistic fits plot on generation.
        """
        csv_path = Path(csv_path)
        if not csv_path.exists():
            if models is None or data_sources is None:
                raise ValueError("models and data_sources required for LogisticModelStore generation (no cached data)")
            logger.info("Generating logistic fits to %s ...", csv_path)
            cls.generate(models, data_sources, mlenergy_data_dir=mlenergy_data_dir).save(csv_path, plot=plot)
        return cls.load(csv_path)