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,
        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,
        *,
        datacenter: DatacenterBackend,
        schedules: dict[str, BatchSizeSchedule],
        dt_s: Fraction = Fraction(1),
    ) -> None:
        self._datacenter = datacenter
        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,
        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,
                    target=self._datacenter,
                )
            ]
        return []

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

    Attributes:
        enabled: If False, the controller is a no-op. Useful for ablations.
        gpus_per_shift: Number of GPUs worth of replicas to shift in a single
            controller tick.
    """

    enabled: bool = False
    gpus_per_shift: int = 8

class LoadShiftController(Controller[OfflineDatacenter, 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: list[OfflineDatacenter],
        grid: OpenDSSGrid,
        models_by_dc: dict[OfflineDatacenter, 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 = list(datacenters)
        self._grid = grid
        self._dc_bus_map = {dc: grid.dc_bus(dc) for dc in datacenters}
        self._models_by_dc = dict(models_by_dc)
        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

    def step(
        self,
        clock: SimulationClock,
        events: EventEmitter,
    ) -> list[DatacenterCommand | GridCommand]:
        grid = self._grid
        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
        bus_voltages: dict[str, list[float]] = {}
        for (bus, _phase), v in zip(v_index, voltages, strict=True):
            bus_voltages.setdefault(bus.lower(), []).append(float(v))

        # Per-DC min/max voltage
        dc_vmin: dict[OfflineDatacenter, float] = {}
        dc_vmax: dict[OfflineDatacenter, float] = {}
        for dc, bus in self._dc_bus_map.items():
            vs = bus_voltages.get(bus.lower(), [])
            if vs:
                dc_vmin[dc] = min(vs)
                dc_vmax[dc] = max(vs)

        commands: list[DatacenterCommand | GridCommand] = []

        for dc in self._datacenters:
            vmin = dc_vmin.get(dc, 1.0)
            vmax = dc_vmax.get(dc, 1.0)

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

            if not is_undervoltage and not is_overvoltage:
                continue

            if not self._is_batch_saturated(dc, is_undervoltage):
                continue

            dc_models = set(self._models_by_dc.get(dc, []))

            if is_undervoltage:
                best_dest = None
                best_v = -1.0
                for other_dc in self._datacenters:
                    if other_dc is dc:
                        continue
                    other_models = set(self._models_by_dc.get(other_dc, []))
                    shared = dc_models & other_models
                    if not shared:
                        continue
                    if other_dc.available_gpu_capacity() < self._config.gpus_per_shift:
                        continue
                    ov = dc_vmin.get(other_dc, 0.0)
                    if ov > best_v:
                        best_v = ov
                        best_dest = other_dc

                if best_dest is None:
                    continue

                shared_models = dc_models & set(self._models_by_dc.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])
                src_active = dc.state.active_replicas_by_model[model]
                if src_active < replicas:
                    continue
                commands.append(ShiftReplicas(model_label=model, replica_delta=-replicas, target=dc))
                commands.append(ShiftReplicas(model_label=model, replica_delta=+replicas, target=best_dest))
                logger.info(
                    "LoadShift: undervoltage at %s (Vmin=%.4f), shift %s x%d replicas -> %s (Vmin=%.4f, free=%d GPUs)",
                    dc.name,
                    vmin,
                    model,
                    replicas,
                    best_dest.name,
                    best_v,
                    best_dest.available_gpu_capacity(),
                )

            elif is_overvoltage:
                if dc.available_gpu_capacity() < self._config.gpus_per_shift:
                    continue

                best_src = None
                best_v = 2.0
                for other_dc in self._datacenters:
                    if other_dc is dc:
                        continue
                    other_models = set(self._models_by_dc.get(other_dc, []))
                    shared = dc_models & other_models
                    if not shared:
                        continue
                    ov = dc_vmax.get(other_dc, 2.0)
                    if ov < best_v:
                        best_v = ov
                        best_src = other_dc

                if best_src is None:
                    continue

                shared_models = dc_models & set(self._models_by_dc.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])
                src_active = best_src.state.active_replicas_by_model[model]
                if src_active < replicas:
                    continue
                commands.append(ShiftReplicas(model_label=model, replica_delta=-replicas, target=best_src))
                commands.append(ShiftReplicas(model_label=model, replica_delta=+replicas, target=dc))
                logger.info(
                    "LoadShift: overvoltage at %s (Vmax=%.4f), shift %s x%d replicas <- %s (Vmax=%.4f)",
                    dc.name,
                    vmax,
                    model,
                    replicas,
                    best_src.name,
                    best_v,
                )

        return commands

    def _is_batch_saturated(
        self,
        dc: OfflineDatacenter,
        is_undervoltage: bool,
    ) -> bool:
        """Check if all models at DC have batch sizes at their limit."""
        state = dc.state
        dc_models = self._models_by_dc.get(dc, [])
        for model_label in dc_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:
                if current_bs > min(feasible):
                    return False
            else:
                if current_bs < max(feasible):
                    return False
        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,
        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, ...],
        *,
        runs: Any = None,
        mlenergy_data_dir: Path | None = None,
    ) -> LogisticModelStore:
        """Generate logistic fits from ML.ENERGY benchmark data.

        Args:
            models: Model specifications; every field (`task`,
                `gpu_model`, `gpus_per_replica`, `batch_sizes`,
                `fit_exclude_batch_sizes`) is used for the benchmark query
                and fit selection.
            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 = {ms.task for ms in models}
            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] = {}
        exclude_by_label: dict[str, set[int]] = {}
        for ms in models:
            if not ms.model_id:
                raise ValueError(f"model_id is required for data generation (model={ms.model_label!r})")

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

        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 = exclude_by_label[model_label]
            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
                _warn_if_fit_suspicious(model_label, _metric_name, fit, batches, y)

        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,
        base_dir: Path,
        specs: tuple[InferenceModelSpec, ...],
        *,
        plot: bool = False,
    ) -> None:
        """Save per-spec logistic fits to `base_dir/<spec.cache_hash()>/logistic_fit.json`.

        Args:
            base_dir: Root of the per-spec cache (typically `data/specs/`).
            specs: Specs whose fits to save.
            plot: If `True`, write a diagnostic fits plot per spec's
                directory when `_by_batch` observations are available.
        """
        import json as _json

        base_dir = Path(base_dir)
        by_batch = getattr(self, "_by_batch", None)
        for spec in specs:
            label = spec.model_label
            if label not in self._power and label not in self._latency and label not in self._throughput:
                continue
            spec_dir = base_dir / spec.cache_hash()
            spec_dir.mkdir(parents=True, exist_ok=True)
            payload: dict[str, Any] = {
                "schema": "logistic_v1",
                "model_label": label,
            }
            for metric_name, store in (
                ("power", self._power),
                ("latency", self._latency),
                ("throughput", self._throughput),
            ):
                model = store.get(label)
                if model is not None:
                    payload[metric_name] = {"L": model.L, "x0": model.x0, "k": model.k, "b0": model.b0}
            (spec_dir / "logistic_fit.json").write_text(_json.dumps(payload, indent=2, sort_keys=True))

            if plot and by_batch is not None and label in by_batch:
                _plot_logistic_fits(
                    {label: by_batch[label]},
                    self._power,
                    self._latency,
                    self._throughput,
                    [label],
                    spec_dir,
                )

    @classmethod
    def load(
        cls,
        base_dir: Path | str,
        specs: tuple[InferenceModelSpec, ...],
    ) -> LogisticModelStore:
        """Load per-spec logistic fits from `base_dir/<spec.cache_hash()>/logistic_fit.json`.

        Args:
            base_dir: Root of the per-spec cache.
            specs: Specs whose fits to load.
        """
        import json as _json

        base_dir = Path(base_dir)
        power: dict[str, LogisticModel] = {}
        latency: dict[str, LogisticModel] = {}
        throughput: dict[str, LogisticModel] = {}
        for spec in specs:
            spec_dir = base_dir / spec.cache_hash()
            path = spec_dir / "logistic_fit.json"
            if not path.exists():
                raise FileNotFoundError(f"Logistic fit not found at {path} (spec={spec.model_label!r})")
            payload = _json.loads(path.read_text())
            label = spec.model_label
            for metric_name, target in (
                ("power", power),
                ("latency", latency),
                ("throughput", throughput),
            ):
                params = payload.get(metric_name)
                if params is not None:
                    target[label] = LogisticModel.from_dict(params)
        if not power and not latency and not throughput:
            raise ValueError(f"No logistic model entries loaded from {base_dir}")
        return cls(power=power, latency=latency, throughput=throughput)

    @classmethod
    def ensure(
        cls,
        base_dir: Path,
        models: tuple[InferenceModelSpec, ...],
        *,
        mlenergy_data_dir: Path | None = None,
        plot: bool = False,
    ) -> LogisticModelStore:
        """Load per-spec logistic fits under `base_dir`, generating missing ones.

        Any spec whose `base_dir/<hash>/logistic_fit.json` is absent
        triggers a targeted regeneration for just that spec; the cached
        ones are read as-is.

        Args:
            base_dir: Root of the per-spec cache.
            models: Model specifications.
            mlenergy_data_dir: Path to compiled mlenergy-data directory.
            plot: If `True`, generate a logistic fits plot per newly
                generated spec directory.
        """
        base_dir = Path(base_dir)
        base_dir.mkdir(parents=True, exist_ok=True)

        missing = tuple(ms for ms in models if not (base_dir / ms.cache_hash() / "logistic_fit.json").exists())
        if missing:
            missing_labels = [ms.model_label for ms in missing]
            logger.info(
                "Generating logistic fits for %d/%d specs (missing=%s) under %s ...",
                len(missing),
                len(models),
                missing_labels,
                base_dir,
            )
            cls.generate(missing, mlenergy_data_dir=mlenergy_data_dir).save(base_dir, missing, plot=plot)
        return cls.load(base_dir, models)