openg2g.datacenter¶

@dataclass(frozen=True)
class DatacenterState:
    """State emitted by a datacenter backend each timestep.

    Contains only universally applicable fields. LLM-inference-specific
    fields (batch sizes, replicas, latency) live on child classes like
    [`LLMDatacenterState`][..LLMDatacenterState].

    Attributes:
        time_s: Simulation time in seconds.
        power_w: Three-phase power in watts.
    """

    time_s: float
    power_w: ThreePhase

@dataclass(frozen=True)
class LLMDatacenterState(DatacenterState):
    """State from a datacenter serving LLM workloads.

    Extends [`DatacenterState`][..DatacenterState] with per-model batch
    size, replica count, and observed inter-token latency fields used
    by LLM controllers.

    Attributes:
        batch_size_by_model: Current batch size per model label.
        active_replicas_by_model: Number of active replicas per model.
        observed_itl_s_by_model: Observed average inter-token latency
            (seconds) per model. `NaN` if unavailable.
    """

    batch_size_by_model: dict[str, int] = field(default_factory=dict)
    active_replicas_by_model: dict[str, int] = field(default_factory=dict)
    observed_itl_s_by_model: dict[str, float] = field(default_factory=dict)

class DatacenterBackend(Generic[DCStateT], ABC):
    """Interface for datacenter power simulation backends."""

    _INIT_SENTINEL = object()

    def __init__(self) -> None:
        self._state: DCStateT | None = None
        self._history: list[DCStateT] = []
        self._dc_base_init = DatacenterBackend._INIT_SENTINEL

    def _check_base_init(self) -> None:
        if getattr(self, "_dc_base_init", None) is not DatacenterBackend._INIT_SENTINEL:
            raise TypeError(f"{type(self).__name__}.__init__ must call super().__init__() ")

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

    @final
    @property
    def state(self) -> DCStateT:
        """Latest emitted state.

        Raises:
            RuntimeError: If accessed before the first `step()` call.
        """
        self._check_base_init()
        if self._state is None:
            raise RuntimeError(f"{type(self).__name__}.state accessed before first step().")
        return self._state

    @final
    def history(self, n: int | None = None) -> list[DCStateT]:
        """Return emitted state history (all, or latest `n`)."""
        self._check_base_init()
        if n is None:
            return list(self._history)
        if n <= 0:
            return []
        return list(self._history[-int(n) :])

    @final
    def do_step(self, clock: SimulationClock, events: EventEmitter) -> DCStateT:
        """Call `step`, record the state, and return it.

        Called by the coordinator. Subclasses should not override this.
        """
        self._check_base_init()
        state = self.step(clock, events)
        self._state = state
        self._history.append(state)
        return state

    @abstractmethod
    def step(self, clock: SimulationClock, events: EventEmitter) -> DCStateT:
        """Advance one native timestep. Return state for this step."""

    @abstractmethod
    def apply_control(self, command: DatacenterCommand, events: EventEmitter) -> None:
        """Apply one command. Takes effect on next step() call."""

    @final
    def do_reset(self) -> None:
        """Clear history and call `reset`.

        Called by the coordinator. Subclasses should not override this.
        """
        self._check_base_init()
        self._state = None
        self._history.clear()
        self.reset()

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

        Called by the coordinator (via `do_reset`) before each
        [`start`][..start]. Must clear all simulation state: counters,
        RNG seeds, cached values. Configuration (dt_s, models,
        templates) is not affected. History is cleared automatically
        by `do_reset`.

        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 (threads, solver circuits).

        Called after [`reset`][..reset], before the simulation loop.
        Override for backends that need resource acquisition (e.g.,
        [`OpenDSSGrid`][openg2g.grid.opendss.OpenDSSGrid] compiles its
        DSS circuit here). No-op by default because most offline
        components 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. Override for
        backends that acquired resources in [`start`][..start]. No-op
        by default.
        """

class LLMBatchSizeControlledDatacenter(DatacenterBackend[DCStateT]):
    """Datacenter that serves LLM workloads and supports batch-size control.

    Marker layer between [`DatacenterBackend`][..DatacenterBackend] and
    concrete implementations. Controllers that issue
    [`SetBatchSize`][openg2g.datacenter.command.SetBatchSize] commands or read
    `active_replicas_by_model` / `observed_itl_s_by_model`
    from state should bind their generic to this class.
    """

    @property
    def phase_share_by_model(self) -> dict[str, np.ndarray]:
        """Per-model phase share vectors `[frac_A, frac_B, frac_C]`.

        Returns an empty dict by default. Consumers treat missing keys
        as uniform `[1/3, 1/3, 1/3]`. Override in subclasses that know
        actual server-to-phase placement.
        """
        return {}

class DatacenterCommand:
    """Base for commands targeting the datacenter backend.

    Subclass this for each concrete datacenter command kind.
    The coordinator routes commands to backends based on this type hierarchy.
    """

    def __init__(self) -> None:
        if type(self) is DatacenterCommand:
            raise TypeError("DatacenterCommand cannot be instantiated directly; subclass it.")

@dataclass(frozen=True)
class SetBatchSize(DatacenterCommand):
    """Set batch sizes for one or more models.

    Attributes:
        batch_size_by_model: Mapping of model label to target batch size.
        ramp_up_rate_by_model: Per-model requests/second ramp-up rate.
            Models not present get immediate changes (rate 0).
        target_site_id: Site this command targets. The coordinator uses
            this to route the command to the correct datacenter.
    """

    batch_size_by_model: dict[str, int]
    ramp_up_rate_by_model: dict[str, float] = field(default_factory=dict)
    target_site_id: str | None = None

@dataclass(frozen=True)
class ShiftReplicas(DatacenterCommand):
    """Shift replicas for a model at this datacenter.

    Positive ``replica_delta`` adds replicas (receiving site);
    negative removes them (sending site).

    Attributes:
        model_label: Which model to shift.
        replica_delta: Number of replicas to add (>0) or remove (<0).
        target_site_id: Site this command targets.  The coordinator uses
            this to route the command to the correct datacenter.
    """

    model_label: str
    replica_delta: int
    target_site_id: str | None = None

class InferenceModelSpec(BaseModel):
    """Specification for one LLM model served in the datacenter.

    This is a pure model-identity object describing *what* is served, not
    *how many* or *at what batch size*.  Deployment-specific parameters
    (replica count, initial batch size) are specified via
    :class:`ModelDeployment`.

    Attributes:
        model_label: Human-readable model identifier (e.g. `"Llama-3.1-70B"`).
        model_id: HuggingFace model ID (e.g. `"meta-llama/Llama-3.1-70B-Instruct"`).
            Used for benchmark data lookups and online API model fields.
        gpus_per_replica: GPUs allocated to each replica (determines model
            parallelism and per-replica power draw).
        itl_deadline_s: Per-model inter-token latency deadline for the OFO
            latency dual (seconds).
        feasible_batch_sizes: Allowed batch sizes. Used by the OFO
            controller for discretizing continuous batch-size updates
            and by the online datacenter for load-generator sizing.
    """

    model_config = ConfigDict(frozen=True)

    model_label: str
    model_id: str = ""
    gpus_per_replica: int
    itl_deadline_s: float
    feasible_batch_sizes: tuple[int, ...]

    @model_validator(mode="after")
    def _validate(self) -> InferenceModelSpec:
        if self.gpus_per_replica < 1:
            raise ValueError(f"gpus_per_replica must be >= 1, got {self.gpus_per_replica}.")
        if self.itl_deadline_s <= 0:
            raise ValueError(f"itl_deadline_s must be > 0, got {self.itl_deadline_s}.")
        if not self.feasible_batch_sizes:
            raise ValueError("feasible_batch_sizes must not be empty.")
        return self

@dataclass(frozen=True)
class ModelDeployment:
    """One model's deployment at a datacenter site.

    Pairs an :class:`InferenceModelSpec` (model identity) with
    deployment-specific parameters.

    Attributes:
        spec: The model specification.
        num_replicas: Number of replicas deployed at this site.
        initial_batch_size: Starting batch size for this deployment.
            Must be in ``spec.feasible_batch_sizes``.
    """

    spec: InferenceModelSpec
    num_replicas: int
    initial_batch_size: int

    def __post_init__(self) -> None:
        if self.num_replicas < 0:
            raise ValueError(f"num_replicas must be >= 0, got {self.num_replicas}.")
        if self.initial_batch_size <= 0:
            raise ValueError(f"initial_batch_size must be > 0, got {self.initial_batch_size}.")
        if self.initial_batch_size not in self.spec.feasible_batch_sizes:
            raise ValueError(
                f"initial_batch_size ({self.initial_batch_size}) must be in "
                f"feasible_batch_sizes ({self.spec.feasible_batch_sizes})."
            )

class TrainingRun:
    """Training workload parameters.

    The trace is eagerly rescaled so its peak matches `target_peak_W_per_gpu`.
    Use `eval_power` to evaluate total training power at a given simulation time.

    Combine with [`at`][.at] and `|` to build a [`TrainingSchedule`][..TrainingSchedule]:

    ```python
    schedule = (
        TrainingRun(n_gpus=2400, trace=trace_a).at(t_start=1000, t_end=2000)
        | TrainingRun(n_gpus=1200, trace=trace_b).at(t_start=2500, t_end=3500)
    )
    ```

    Attributes:
        n_gpus: Number of GPUs running the training workload.
        trace: Single-GPU [`TrainingTrace`][openg2g.datacenter.workloads.training.TrainingTrace].
        target_peak_W_per_gpu: The trace is rescaled so its peak equals this value.
    """

    __slots__ = ("_period", "_rescaled_power", "_trace_time", "n_gpus", "target_peak_W_per_gpu", "trace")

    def __init__(self, *, n_gpus: int, trace: TrainingTrace, target_peak_W_per_gpu: float = 400.0) -> None:
        if n_gpus <= 0:
            raise ValueError(f"TrainingRun n_gpus must be > 0, got {n_gpus}.")
        self.n_gpus = n_gpus
        self.trace = trace
        self.target_peak_W_per_gpu = target_peak_W_per_gpu

        t = np.asarray(trace.t_s, float)
        p = np.asarray(trace.power_w, float)
        t = t - t[0]
        period = float(t[-1] - t[0])
        if period <= 0:
            raise ValueError("Training trace time span must be positive.")
        peak = float(np.max(p))
        if peak <= 0:
            raise ValueError("Training trace has non-positive peak; cannot scale.")
        self._rescaled_power = p * (target_peak_W_per_gpu / peak)
        self._trace_time = t
        self._period = period

    def eval_power(self, t: float, t_start: float, t_end: float) -> float:
        """Evaluate total training power at simulation time `t`.

        Returns zero if `t` is outside `[t_start, t_end]`.

        Args:
            t: Global simulation time (seconds).
            t_start: Time when training becomes active (seconds).
            t_end: Time when training stops (seconds).

        Returns:
            Total training power (W) across all `n_gpus` GPUs.
        """
        if t < t_start or t > t_end:
            return 0.0
        t_local = t - t_start
        t_mod = t_local % self._period
        p_1gpu = float(np.interp(t_mod, self._trace_time, self._rescaled_power))
        return p_1gpu * self.n_gpus

    def at(self, t_start: float, t_end: float) -> TrainingSchedule:
        """Schedule this training run over `[t_start, t_end]`.

        Args:
            t_start: Global simulation time when training becomes active (seconds).
            t_end: Global simulation time when training stops (seconds).

        Returns:
            A single-entry [`TrainingSchedule`][...TrainingSchedule].
        """
        if t_end < t_start:
            raise ValueError(f"t_end ({t_end}) must be >= t_start ({t_start}).")
        return TrainingSchedule(((self, float(t_start), float(t_end)),))

class TrainingSchedule:
    """Ordered collection of [`TrainingRun`][..TrainingRun] objects scheduled
    over time windows.

    Each entry is a `(TrainingRun, t_start, t_end)` tuple. Entries are
    sorted by `t_start`.

    Built with [`TrainingRun.at`][..TrainingRun.at] and `|`.

    Example:

    ```python
    schedule = (
        TrainingRun(n_gpus=2400, trace=trace_a).at(t_start=1000, t_end=2000)
        | TrainingRun(n_gpus=1200, trace=trace_b).at(t_start=2500, t_end=3500)
    )
    ```
    """

    __slots__ = ("_entries",)

    def __init__(self, entries: tuple[tuple[TrainingRun, float, float], ...] = ()) -> None:
        self._entries = tuple(sorted(entries, key=lambda e: e[1]))

    def __or__(self, other: TrainingSchedule) -> TrainingSchedule:
        return TrainingSchedule((*self._entries, *other._entries))

    def __iter__(self) -> Iterator[tuple[TrainingRun, float, float]]:
        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 = [f"TrainingRun(n_gpus={r.n_gpus}).at(t_start={s}, t_end={e})" for r, s, e in self._entries]
        return " | ".join(parts)

@dataclass(frozen=True)
class InferenceRamp:
    """Inference server ramp parameters.

    Transitions the active replica count for a specific model to `target`.
    Combine with [`at`][.at] and `|` to build an
    [`InferenceRampSchedule`][..InferenceRampSchedule]:

    ```python
    ramps = (
        InferenceRamp(target=144, model="Llama-3.1-8B").at(t_start=2500, t_end=3000)
        | InferenceRamp(target=864, model="Llama-3.1-8B").at(t_start=3200, t_end=3400)
    )
    ```

    Attributes:
        target: Target number of active replicas after the ramp completes.
        model: Model label this ramp applies to.
    """

    target: int
    model: str

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

    def at(self, t_start: float, t_end: float) -> InferenceRampSchedule:
        """Schedule this ramp over `[t_start, t_end]`.

        Args:
            t_start: Global simulation time when the ramp begins (seconds).
            t_end: Global simulation time when the ramp ends (seconds).

        Returns:
            A single-entry [`InferenceRampSchedule`][...InferenceRampSchedule].
        """
        if t_end < t_start:
            raise ValueError(f"t_end ({t_end}) must be >= t_start ({t_start}).")
        return InferenceRampSchedule(((self, float(t_start), float(t_end)),))

class InferenceRampSchedule:
    """Ordered collection of [`InferenceRamp`][..InferenceRamp] events for
    a single model.

    Each entry is an `(InferenceRamp, t_start, t_end)` tuple. Entries are
    sorted by `t_start`.

    Built with [`InferenceRamp.at`][..InferenceRamp.at] and `|`.

    Semantics: before the first ramp, the active count equals
    ``initial_count``.  During each `[t_start, t_end]` window the count
    linearly interpolates from the previous level to ``target``.  Between
    ramps, the count holds at the last target.

    An empty schedule means ``initial_count`` replicas are active at all
    times.

    Example:

    ```python
    ramps = (
        InferenceRamp(target=144, model="Llama-3.1-8B").at(t_start=2500, t_end=3000)
        | InferenceRamp(target=720, model="Llama-3.1-8B").at(t_start=3200, t_end=3400)
    )
    ```
    """

    __slots__ = ("_entries", "_initial_count")

    def __init__(
        self,
        entries: tuple[tuple[InferenceRamp, float, float], ...] = (),
        *,
        initial_count: int = 0,
    ) -> None:
        self._entries = tuple(sorted(entries, key=lambda e: e[1]))
        self._initial_count = initial_count

    @property
    def initial_count(self) -> int:
        """Replica count before any ramp event."""
        return self._initial_count

    def __or__(self, other: InferenceRampSchedule) -> InferenceRampSchedule:
        # Preserve initial_count from left operand (the one built first).
        return InferenceRampSchedule(
            (*self._entries, *other._entries),
            initial_count=self._initial_count,
        )

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

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

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

    def for_model(self, model_label: str, *, initial_count: int | None = None) -> InferenceRampSchedule:
        """Return a schedule containing only entries for *model_label*.

        Args:
            model_label: Model to filter for.
            initial_count: Override the initial replica count for this
                per-model schedule.  If ``None``, inherits from ``self``.
        """
        filtered = tuple(e for e in self._entries if e[0].model == model_label)
        ic = initial_count if initial_count is not None else self._initial_count
        return InferenceRampSchedule(filtered, initial_count=ic)

    def max_count(self) -> int:
        """Return the maximum target across all entries, or ``initial_count`` if empty."""
        if not self._entries:
            return self._initial_count
        return max(self._initial_count, *(r.target for r, _, _ in self._entries))

    def __repr__(self) -> str:
        parts = []
        for r, s, e in self._entries:
            parts.append(f"InferenceRamp(target={r.target}, model={r.model!r}).at(t_start={s}, t_end={e})")
        prefix = f"InferenceRampSchedule(initial_count={self._initial_count}): "
        return prefix + (" | ".join(parts) if parts else "(empty)")

    def count_at(self, t: float | np.ndarray) -> float | np.ndarray:
        """Evaluate the active replica count at time(s) *t*.

        Piecewise-linear interpolation between ramp events.
        Before the first ramp, returns ``initial_count``.

        Args:
            t: Scalar or array of global simulation times (seconds).

        Returns:
            Active replica count(s), same shape as *t*.
        """
        if isinstance(t, np.ndarray):
            return self._count_array(t)
        return float(self._count_scalar(float(t)))

    def _count_scalar(self, t: float) -> float:
        level = float(self._initial_count)
        for ramp, t_start, t_end in self._entries:
            if t < t_start:
                return level
            if t <= t_end:
                if t_end == t_start:
                    return float(ramp.target)
                alpha = (t - t_start) / (t_end - t_start)
                return level + (float(ramp.target) - level) * alpha
            level = float(ramp.target)
        return level

    def _count_array(self, t: np.ndarray) -> np.ndarray:
        vfunc = np.vectorize(self._count_scalar, otypes=[float])
        return vfunc(t)

class DatacenterConfig(BaseModel):
    """Physical datacenter facility configuration.

    Attributes:
        gpus_per_server: Number of GPUs per physical server rack.
        base_kw_per_phase: Constant base load per phase (kW).
        power_factor: Power factor of the datacenter loads (lagging).
    """

    model_config = ConfigDict(frozen=True)

    gpus_per_server: int = 8
    base_kw_per_phase: float = 0.0
    power_factor: float = 0.95

    @model_validator(mode="after")
    def _validate(self) -> DatacenterConfig:
        if self.gpus_per_server < 1:
            raise ValueError(f"gpus_per_server must be >= 1, got {self.gpus_per_server}.")
        if not (0.0 < self.power_factor <= 1.0):
            raise ValueError(f"power_factor must be in (0, 1], got {self.power_factor}.")
        return self

class PowerAugmentationConfig(BaseModel):
    """Power augmentation settings for virtual server scaling.

    Controls per-server amplitude jitter and additive noise applied during
    power augmentation.

    Attributes:
        amplitude_scale_range: `(low, high)` range for per-server amplitude
            scaling. Each virtual server draws a uniform multiplier from this range.
        noise_fraction: Gaussian noise standard deviation as a fraction of
            per-server power.
    """

    model_config = ConfigDict(frozen=True)

    amplitude_scale_range: tuple[float, float] = (1.0, 1.0)
    noise_fraction: float = 0.0

class ActivationPolicy(ABC):
    """Per-model activation policy that answers "which servers are active?"

    Subclass to implement custom activation logic. The datacenter creates
    one policy per model and passes it to
    [`InferencePowerAugmenter`][openg2g.datacenter.workloads.inference.InferencePowerAugmenter].
    """

    @abstractmethod
    def active_mask(self, t: float) -> np.ndarray:
        """Boolean mask of active servers at time *t*.

        Returns:
            Array of shape `(num_servers,)` with `True` for active servers.
        """

    def active_indices(self, t: float) -> np.ndarray:
        """Indices of active servers at time *t*.

        The default implementation returns indices in ascending order
        via `np.where(`[`active_mask`][..active_mask]`(t))`. Subclasses
        may override to return
        indices in a specific order (e.g., priority order) to control
        floating-point summation order in the datacenter.

        Returns:
            1-D int array of active server indices.
        """
        return np.where(self.active_mask(t))[0]

class RampActivationPolicy(ActivationPolicy):
    """Activate servers by fixed random priority, following an
    [`InferenceRampSchedule`][openg2g.datacenter.config.InferenceRampSchedule].

    At time *t*, the top-*k* servers (by random priority) are active,
    where *k* is derived from the schedule's absolute replica count and
    the model's GPU requirements.

    This is the default policy used by
    [`OfflineDatacenter`][openg2g.datacenter.offline.OfflineDatacenter].

    Args:
        schedule: Per-model ramp schedule (absolute replica counts).
        num_servers: Total allocated servers (may exceed baseline when
            ramp targets exceed the initial replica count).
        rng: RNG for randomizing priority ordering. Consumed once at
            construction time.
        gpus_per_replica: GPUs required per model replica.
        gpus_per_server: GPUs per physical server.
    """

    __slots__ = ("_gpus_per_replica", "_gpus_per_server", "_n", "_priority", "_replica_offset", "_schedule")

    def __init__(
        self,
        schedule: InferenceRampSchedule,
        num_servers: int,
        rng: np.random.Generator,
        *,
        gpus_per_replica: int,
        gpus_per_server: int,
    ) -> None:
        self._schedule = schedule
        self._n = num_servers
        self._gpus_per_replica = gpus_per_replica
        self._gpus_per_server = gpus_per_server
        self._replica_offset: int = 0
        priority = np.arange(num_servers, dtype=int)
        rng.shuffle(priority)
        self._priority = priority

    def _servers_for_count(self, replica_count: float) -> int:
        """Convert a replica count to a server count."""
        import math

        gpus_needed = max(0.0, replica_count + self._replica_offset) * self._gpus_per_replica
        return max(0, min(self._n, math.ceil(gpus_needed / self._gpus_per_server)))

    def active_mask(self, t: float) -> np.ndarray:
        count = self._schedule.count_at(t)
        k = self._servers_for_count(float(count))
        mask = np.zeros(self._n, dtype=bool)
        mask[self._priority[:k]] = True
        return mask

    def active_indices(self, t: float) -> np.ndarray:
        """Return active server indices in priority order."""
        count = self._schedule.count_at(t)
        k = self._servers_for_count(float(count))
        return self._priority[:k].copy()