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 `LLMDatacenterState`.
    """

    time_s: float
    power_w: ThreePhase

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

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

    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."""

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

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

        Raises:
            RuntimeError: If accessed before the first `step()` call.
        """

    @abstractmethod
    def history(self, n: int | None = None) -> Sequence[DCStateT]:
        """Return emitted state history (all, or latest `n`)."""

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

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

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

        Called by the coordinator before each `start()`. Must clear all
        simulation state: history, counters, RNG seeds, cached values.
        Configuration (dt_s, models, templates) 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 (threads, solver circuits).

        Called after `reset()`, before the simulation loop. Override for
        backends that need resource acquisition (e.g., `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()`. No-op by default.
        """

    def bind_event_emitter(self, emitter: EventEmitter) -> None:
        """Attach a clock-bound emitter for backend-originated events."""

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

    Marker layer between `DatacenterBackend` and concrete implementations.
    Controllers that issue `set_batch_size` 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 {}

@dataclass(frozen=True)
class TrainingRun:
    """A single training workload window.

    Attributes:
        t_start: Global simulation time when training becomes active (seconds).
        t_end: Global simulation time when training stops (seconds).
        n_gpus: Number of GPUs running the training workload.
        trace: Single-GPU training power trace.
        target_peak_W_per_gpu: The trace is rescaled so its peak equals this value.
    """

    t_start: float
    t_end: float
    n_gpus: int
    trace: TrainingTrace
    target_peak_W_per_gpu: float = 400.0

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

    def __or__(self, other: TrainingRun | TrainingSchedule) -> TrainingSchedule:
        if isinstance(other, TrainingRun):
            return TrainingSchedule(entries=(self, other))
        return TrainingSchedule(entries=(self, *other))

class TrainingSchedule:
    """Ordered collection of training windows, built with `|`.

    Example:

        schedule = (
            TrainingRun(t_start=500, t_end=1500, n_gpus=2400, trace=trace_a)
            | TrainingRun(t_start=2000, t_end=3000, n_gpus=1200, trace=trace_b)
        )
    """

    __slots__ = ("_entries",)

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

    def __or__(self, other: TrainingRun | TrainingSchedule) -> TrainingSchedule:
        if isinstance(other, TrainingRun):
            return TrainingSchedule(entries=(*self._entries, other))
        return TrainingSchedule(entries=(*self._entries, *other._entries))

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

@dataclass(frozen=True)
class ServerRamp:
    """A single server ramp event.

    Transitions the active-server fraction to `target` linearly over
    `[t_start, t_end]`.

    Attributes:
        t_start: Global simulation time when the ramp begins (seconds).
        t_end: Global simulation time when the ramp ends (seconds).
        target: Target active-server fraction after the ramp (0.0--1.0).
    """

    t_start: float
    t_end: float
    target: float

    def __post_init__(self) -> None:
        if self.t_end < self.t_start:
            raise ValueError(f"ServerRamp t_end ({self.t_end}) must be >= t_start ({self.t_start}).")
        if not (0.0 <= self.target <= 1.0):
            raise ValueError(f"ServerRamp target must be in [0.0, 1.0], got {self.target}.")

    def __or__(self, other: ServerRamp | ServerRampSchedule) -> ServerRampSchedule:
        if isinstance(other, ServerRamp):
            return ServerRampSchedule(entries=(self, other))
        return ServerRampSchedule(entries=(self, *other))

class ServerRampSchedule:
    """Ordered collection of server ramp events, built with `|`.

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

    An empty schedule means all servers are active (fraction = 1.0) at all times.

    Example:

        ramps = (
            ServerRamp(t_start=2500, t_end=3000, target=0.2)
            | ServerRamp(t_start=3200, t_end=3400, target=1.0)
        )
    """

    __slots__ = ("_entries",)

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

    def __or__(self, other: ServerRamp | ServerRampSchedule) -> ServerRampSchedule:
        if isinstance(other, ServerRamp):
            return ServerRampSchedule(entries=(*self._entries, other))
        return ServerRampSchedule(entries=(*self._entries, *other._entries))

    def __iter__(self) -> Iterator[ServerRamp]:
        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"ServerRamp(t_start={r.t_start}, t_end={r.t_end}, target={r.target})" for r in self._entries]
        return " | ".join(parts)

    def fraction_at(self, t: float | np.ndarray) -> float | np.ndarray:
        """Evaluate the active-server fraction at time(s) *t*.

        Piecewise-linear interpolation between ramp events.
        Before the first ramp, fraction = 1.0.

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

        Returns:
            Active-server fraction(s), same shape as *t*.
        """
        if isinstance(t, np.ndarray):
            return self._fraction_array(t)
        return float(self._fraction_scalar(float(t)))

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

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

@dataclass(frozen=True)
class DatacenterConfig:
    """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).
    """

    gpus_per_server: int = 8
    base_kw_per_phase: float = 0.0

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

class WorkloadConfig:
    """What runs in the datacenter: inference, training, and ramp events.

    Accepts flexible input types and normalizes them internally:
    - A single `TrainingRun` is wrapped in a `TrainingSchedule`.
    - A single `ServerRamp` is wrapped in a `ServerRampSchedule`.
    - `None` yields an empty schedule.

    Properties always return schedule types, eliminating `isinstance`
    checks at consumption sites.

    Args:
        inference: LLM inference workload specification.
        training: Training workload window(s). `None` disables training overlay.
        server_ramps: Server ramp event(s). `None` keeps all servers active.
    """

    def __init__(
        self,
        inference: LLMInferenceWorkload,
        training: TrainingRun | TrainingSchedule | None = None,
        server_ramps: ServerRamp | ServerRampSchedule | None = None,
    ) -> None:
        self._inference = inference

        if training is None:
            self._training = TrainingSchedule(entries=())
        elif isinstance(training, TrainingRun):
            self._training = TrainingSchedule(entries=(training,))
        else:
            self._training = training

        if server_ramps is None:
            self._server_ramps = ServerRampSchedule(entries=())
        elif isinstance(server_ramps, ServerRamp):
            self._server_ramps = ServerRampSchedule(entries=(server_ramps,))
        else:
            self._server_ramps = server_ramps

    @property
    def inference(self) -> LLMInferenceWorkload:
        return self._inference

    @property
    def training(self) -> TrainingSchedule:
        return self._training

    @property
    def server_ramps(self) -> ServerRampSchedule:
        return self._server_ramps

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

    Created by `ActivationStrategy.for_model` and bound to a specific
    model's server pool.
    """

    @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(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 ActivationStrategy(ABC):
    """Factory that creates per-model `ActivationPolicy` instances.

    A strategy is instantiated once and passed to the datacenter. When
    the datacenter builds each model's server layout, it calls
    `for_model` to create a model-specific `ActivationPolicy`.

    Subclass to implement custom activation strategies. The `phase_list`
    argument in `for_model` enables phase-aware load balancing.
    """

    @abstractmethod
    def for_model(
        self,
        *,
        num_servers: int,
        phase_list: np.ndarray,
        rng: np.random.Generator,
    ) -> ActivationPolicy:
        """Create a policy for one model's server pool.

        Args:
            num_servers: Number of physical servers for this model.
            phase_list: Phase assignment per server (0=A, 1=B, 2=C), shape
                `(num_servers,)`.
            rng: RNG for randomized decisions (priority ordering, etc.).
                Implementations must consume RNG calls deterministically
                so that downstream layout generation is reproducible.

        Returns:
            Policy that answers `active_mask(t)` queries.
        """

class RampActivationStrategy(ActivationStrategy):
    """Activate servers by fixed random priority, following a `ServerRampSchedule`.

    At time *t*, the top-*k* servers (by random priority) are active, where
    `k = round(schedule.fraction_at(t) * num_servers)`.

    This is the default strategy used by `OfflineDatacenter`.

    Args:
        schedule: Temporal ramp schedule mapping time to active-server fraction.
    """

    def __init__(self, schedule: ServerRampSchedule) -> None:
        self._schedule = schedule

    def for_model(
        self,
        *,
        num_servers: int,
        phase_list: np.ndarray,
        rng: np.random.Generator,
    ) -> ActivationPolicy:
        priority = np.arange(num_servers, dtype=int)
        rng.shuffle(priority)
        return _RampActivationPolicy(self._schedule, num_servers, priority)

@dataclass
class ServerLayout:
    """Per-model server layout describing how GPUs are organized.

    Attributes:
        num_servers: Number of physical servers for this model.
        total_gpus: Total GPU count across all servers.
        gpus_per_replica: GPUs per model replica.
        gpus_per_server_list: GPU count per server (last may be partial).
        phase_list: Phase assignment per server (0=A, 1=B, 2=C).
        activation_policy: Determines which servers are active at time *t*.
        stagger_offsets: Per-server offsets for desynchronization. In offline
            mode these are integer indices into a power template; in online
            mode they can be float time offsets into a rolling buffer.
        amplitude_scales: Per-server power multiplier for inter-server variation.
        noise_fraction: Gaussian noise standard deviation as a fraction of
            per-server power.
    """

    num_servers: int
    total_gpus: int
    gpus_per_replica: int
    gpus_per_server_list: np.ndarray
    phase_list: np.ndarray
    activation_policy: ActivationPolicy
    stagger_offsets: np.ndarray
    amplitude_scales: np.ndarray
    noise_fraction: float

    @classmethod
    def build(
        cls,
        model_spec: LLMInferenceModelSpec,
        *,
        gpus_per_server: int,
        stagger_range: int | float,
        activation_strategy: ActivationStrategy,
        amplitude_scale_range: tuple[float, float],
        noise_fraction: float,
        rng: np.random.Generator,
    ) -> ServerLayout:
        """Build a server layout for one model.

        This is a pure function of its inputs (plus RNG state). The caller
        is responsible for providing a consistently-seeded RNG so that
        layout generation is reproducible.

        Args:
            model_spec: Model specification (replicas, GPUs per replica, etc.).
            gpus_per_server: Number of GPUs per physical server rack.
            stagger_range: Upper bound for per-server stagger offsets. If `int`,
                offsets are drawn from `rng.integers(0, stagger_range)` (for
                offline template indexing). If `float`, offsets are drawn from
                `rng.uniform(0, stagger_range)` (for online time-based staggering).
            activation_strategy: Strategy for determining active servers.
            amplitude_scale_range: `(low, high)` range for per-server amplitude
                scaling. Each server draws a uniform multiplier from this range.
            noise_fraction: Gaussian noise standard deviation as a fraction
                of per-server power.
            rng: Random number generator (consumed for phase assignment,
                activation policy, stagger offsets, and amplitude scales).

        Returns:
            Frozen `ServerLayout` for the model.
        """
        num_replicas = int(model_spec.num_replicas)
        gpus_per_replica = int(model_spec.gpus_per_replica)
        total_gpus = num_replicas * gpus_per_replica
        num_servers = int(math.ceil(total_gpus / gpus_per_server))

        gpus_per_server_list = np.full(num_servers, gpus_per_server, dtype=int)
        tail = total_gpus - (num_servers - 1) * gpus_per_server
        gpus_per_server_list[-1] = int(tail) if tail > 0 else gpus_per_server

        sA, sB, sC = split_integer_evenly(num_servers, 3)
        phase_list = np.asarray(([0] * sA) + ([1] * sB) + ([2] * sC), dtype=int)
        rng.shuffle(phase_list)

        bound_policy = activation_strategy.for_model(
            num_servers=num_servers,
            phase_list=phase_list,
            rng=rng,
        )

        if isinstance(stagger_range, int):
            stagger_offsets = rng.integers(low=0, high=max(stagger_range, 1), size=num_servers)
        else:
            stagger_offsets = rng.uniform(0.0, max(float(stagger_range), 1e-9), size=num_servers)

        amplitude_scales = rng.uniform(
            float(amplitude_scale_range[0]),
            float(amplitude_scale_range[1]),
            size=num_servers,
        )

        return cls(
            num_servers=num_servers,
            total_gpus=total_gpus,
            gpus_per_replica=gpus_per_replica,
            gpus_per_server_list=gpus_per_server_list,
            phase_list=phase_list,
            activation_policy=bound_policy,
            stagger_offsets=stagger_offsets,
            amplitude_scales=amplitude_scales,
            noise_fraction=float(noise_fraction),
        )

@dataclass(frozen=True)
class AugmentedPower:
    """Result of power augmentation for one simulation timestep.

    Attributes:
        power_w: Three-phase total power (watts), including base load.
        power_by_model_w: Per-model total active power (watts).
        active_replicas_by_model: Per-model active replica count.
    """

    power_w: ThreePhase
    power_by_model_w: dict[str, float] = field(default_factory=dict)
    active_replicas_by_model: dict[str, int] = field(default_factory=dict)

class PowerAugmenter:
    """Scales per-GPU power through server layouts to three-phase power.

    Given per-GPU power values for each server (one value per server per
    model), applies per-server scaling, noise, activation masking, and
    phase summation to produce datacenter-level three-phase power.

    This class is backend-agnostic. The offline datacenter feeds it
    template-indexed values; the online datacenter can feed it
    live-measured values.

    Args:
        layouts: Per-model server layouts.
        base_w_per_phase: Constant base load per phase (watts).
        seed: Random seed for noise RNG.
    """

    def __init__(
        self,
        layouts: dict[str, ServerLayout],
        base_w_per_phase: float = 0.0,
        seed: int = 0,
    ) -> None:
        self._layouts = layouts
        self._base_w_per_phase = float(base_w_per_phase)
        self._seed = int(seed)
        self._rng = np.random.default_rng(self._seed)

    def step(
        self,
        per_gpu_by_model: dict[str, np.ndarray],
        t: float,
    ) -> AugmentedPower:
        """Augment per-server per-GPU power to three-phase power.

        Args:
            per_gpu_by_model: Mapping of model label to per-GPU power
                array of shape `(num_servers,)`. Only models with active
                replicas should be included.
            t: Current simulation time (seconds), passed to activation
                policies.

        Returns:
            `AugmentedPower` with three-phase power, per-model power,
            and per-model active replica counts.
        """
        phase_power = np.full(3, self._base_w_per_phase, dtype=float)
        power_by_model: dict[str, float] = {}
        active_replicas_by_model: dict[str, int] = {}

        for label, per_gpu in per_gpu_by_model.items():
            layout = self._layouts[label]

            server_powers = per_gpu * layout.gpus_per_server_list * layout.amplitude_scales
            if layout.noise_fraction > 0:
                levels = np.maximum(server_powers, 1.0)
                server_powers = (
                    server_powers + self._rng.normal(0.0, 1.0, size=layout.num_servers) * layout.noise_fraction * levels
                )
            server_powers = np.maximum(server_powers, 0.0)

            active_indices = layout.activation_policy.active_indices(t)
            active_powers = server_powers[active_indices]
            active_phases = layout.phase_list[active_indices]

            model_phase_power = np.zeros(3, dtype=float)
            np.add.at(model_phase_power, active_phases, active_powers)
            phase_power += model_phase_power

            power_by_model[label] = float(np.sum(active_powers))
            active_gpus = int(np.sum(layout.gpus_per_server_list[active_indices]))
            active_replicas_by_model[label] = active_gpus // layout.gpus_per_replica

        return AugmentedPower(
            power_w=ThreePhase(
                a=float(phase_power[0]),
                b=float(phase_power[1]),
                c=float(phase_power[2]),
            ),
            power_by_model_w=power_by_model,
            active_replicas_by_model=active_replicas_by_model,
        )

    def reset(self) -> None:
        """Re-seed the noise RNG to its initial state."""
        self._rng = np.random.default_rng(self._seed)