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, *, name: str) -> None:
        self._name = name
        self._state: DCStateT | None = None
        self._history: list[DCStateT] = []
        self._dc_base_init = DatacenterBackend._INIT_SENTINEL

    @property
    def name(self) -> str:
        """Human-readable name for logging and display."""
        return self._name

    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 applies commands via `command.target.apply_control()`.

    Attributes:
        target: Datacenter backend this command targets.
    """

    target: DatacenterBackend

    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: Datacenter backend this command targets.
    """

    batch_size_by_model: dict[str, int]
    ramp_up_rate_by_model: dict[str, float] = field(default_factory=dict)
    target: DatacenterBackend = field(kw_only=True)

@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: Datacenter backend this command targets.
    """

    model_label: str
    replica_delta: int
    target: DatacenterBackend = field(kw_only=True)

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
    [`ModelDeployment`][openg2g.datacenter.config.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 [`InferenceModelSpec`][openg2g.datacenter.config.InferenceModelSpec] (model identity)
    with the initial batch size. Replica counts (including runtime ramps)
    live on [`ReplicaSchedule`][openg2g.datacenter.config.ReplicaSchedule]
    and are passed separately via `OfflineWorkload.replica_schedules`.

    Attributes:
        spec: The model specification.
        initial_batch_size: Starting batch size for this deployment.
            Must be in `spec.feasible_batch_sizes`.
    """

    spec: InferenceModelSpec
    initial_batch_size: int

    def __post_init__(self) -> None:
        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)

class ReplicaSchedule:
    """Per-model replica count over time.

    Specifies the initial replica count and optional linear ramps to
    new target counts over time windows. Method-chaining API:

    ```python
    ReplicaSchedule(initial=720).ramp_to(144, t_start=2500, t_end=3000)

    ReplicaSchedule(initial=720)
        .ramp_to(144, t_start=2500, t_end=3000)
        .ramp_to(200, t_start=3200, t_end=3400)
    ```

    Semantics: before the first ramp, the active count equals `initial`.
    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.

    Attributes:
        initial: Replica count before any ramp.
    """

    __slots__ = ("_initial", "_ramps")

    def __init__(self, *, initial: int) -> None:
        if initial < 0:
            raise ValueError(f"ReplicaSchedule initial must be >= 0, got {initial}.")
        self._initial = initial
        self._ramps: tuple[tuple[int, float, float], ...] = ()

    @property
    def initial(self) -> int:
        """Replica count before any ramp."""
        return self._initial

    def ramp_to(self, target: int, *, t_start: float, t_end: float) -> ReplicaSchedule:
        """Add a linear ramp to `target` over `[t_start, t_end]`.

        Ramps may be specified in any order; they are sorted by `t_start`
        in the returned schedule. The new ramp must not overlap any
        existing ramp's window (touching at the boundary is allowed:
        `prev.t_end == new.t_start`).

        Args:
            target: Target replica count after the ramp completes.
            t_start: Global simulation time when the ramp begins (seconds).
            t_end: Global simulation time when the ramp ends (seconds).

        Returns:
            A new `ReplicaSchedule` with the ramp added and ramps sorted
            by `t_start`.

        Raises:
            ValueError: If `target` is negative, `t_end < t_start`, or
                the new ramp overlaps any existing ramp window.
        """
        if target < 0:
            raise ValueError(f"ramp_to target must be >= 0, got {target}.")
        if t_end < t_start:
            raise ValueError(f"t_end ({t_end}) must be >= t_start ({t_start}).")
        # Reject overlap with existing ramps: [t_start, t_end] must not
        # overlap any existing [s, e]. Touching boundaries (t_end == s or
        # t_start == e) is allowed.
        for _target, s, e in self._ramps:
            if t_start < e and t_end > s:
                raise ValueError(f"ramp_to window [{t_start}, {t_end}] overlaps existing ramp window [{s}, {e}].")
        new = ReplicaSchedule(initial=self._initial)
        new._ramps = tuple(
            sorted(
                (*self._ramps, (int(target), float(t_start), float(t_end))),
                key=lambda r: r[1],
            )
        )
        return new

    def max_count(self) -> int:
        """Return the maximum replica count (initial or any ramp target)."""
        if not self._ramps:
            return self._initial
        return max(self._initial, *(target for target, _, _ in self._ramps))

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

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

        Returns:
            Active replica count(s), same shape as *t*.
        """
        if isinstance(t, np.ndarray):
            vfunc = np.vectorize(self._count_scalar, otypes=[float])
            return vfunc(t)
        return float(self._count_scalar(float(t)))

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

    def __repr__(self) -> str:
        parts = [f"ReplicaSchedule(initial={self._initial})"]
        for target, t_start, t_end in self._ramps:
            parts.append(f".ramp_to({target}, t_start={t_start}, t_end={t_end})")
        return "".join(parts)

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

    def __bool__(self) -> bool:
        return True

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

    @model_validator(mode="after")
    def _validate(self) -> PowerAugmentationConfig:
        lo, hi = self.amplitude_scale_range
        if lo > hi:
            raise ValueError(f"amplitude_scale_range low ({lo}) must be <= high ({hi}).")
        if lo <= 0:
            raise ValueError(f"amplitude_scale_range low ({lo}) must be > 0.")
        if self.noise_fraction < 0:
            raise ValueError(f"noise_fraction must be >= 0, got {self.noise_fraction}.")
        return self

@dataclass
class ServerPool:
    """Shared pool of physical servers for a datacenter.

    All models draw servers from this single pool based on their current
    GPU demand. Each server has model-independent properties (phase
    assignment, stagger offset, amplitude scale). The pool handles
    server-to-model allocation at each timestep via per-model priority
    orderings.

    Attributes:
        num_servers: Total number of physical servers in the pool.
        gpus_per_server: GPUs per physical server.
        phase_list: Phase assignment per server (0=A, 1=B, 2=C).
        stagger_offsets: Per-server offsets for desynchronization.
        amplitude_scales: Per-server power multiplier for inter-server variation.
        noise_fraction: Gaussian noise standard deviation as a fraction of
            per-server power.
        model_priorities: Per-model priority ordering over all servers.
            Each model gets a deterministic permutation of `[0..num_servers)`.
    """

    num_servers: int
    gpus_per_server: int
    phase_list: np.ndarray
    stagger_offsets: np.ndarray
    amplitude_scales: np.ndarray
    noise_fraction: float
    model_priorities: dict[str, np.ndarray]

    def allocate(self, gpu_demands: dict[str, int]) -> dict[str, np.ndarray]:
        """Allocate servers to models with phase-balanced round-robin.

        Each model gets k = ceil(gpu_demand / gpus_per_server) servers.
        Servers are picked by cycling through phases (A, B, C, A, B, C, ...)
        and within each phase, picking the highest-priority unclaimed server.
        This ensures each model's allocation is as phase-balanced as possible.

        Args:
            gpu_demands: Mapping of model label to total GPUs needed.

        Returns:
            Mapping of model label to array of allocated server indices.
        """
        claimed = np.zeros(self.num_servers, dtype=bool)
        result: dict[str, np.ndarray] = {}

        # Pre-group servers by phase, sorted by priority per model
        phase_groups: dict[str, list[list[int]]] = {}
        for label in sorted(gpu_demands):
            priority = self.model_priorities[label]
            groups: list[list[int]] = [[], [], []]
            for idx in priority:
                groups[self.phase_list[idx]].append(idx)
            phase_groups[label] = groups

        for label in sorted(gpu_demands):
            demand = gpu_demands[label]
            if demand <= 0:
                result[label] = np.array([], dtype=int)
                continue
            k = math.ceil(demand / self.gpus_per_server)
            groups = phase_groups[label]
            cursors = [0, 0, 0]
            allocated: list[int] = []
            phase = 0

            while len(allocated) < k:
                # Try each phase starting from current, wrapping around
                found = False
                for attempt in range(3):
                    p = (phase + attempt) % 3
                    while cursors[p] < len(groups[p]):
                        idx = groups[p][cursors[p]]
                        cursors[p] += 1
                        if not claimed[idx]:
                            allocated.append(idx)
                            claimed[idx] = True
                            phase = (p + 1) % 3
                            found = True
                            break
                    if found:
                        break
                if not found:
                    raise RuntimeError(
                        f"ServerPool over-subscribed: model {label!r} requested "
                        f"{k} servers but pool has no more unclaimed servers."
                    )

            result[label] = np.array(allocated, dtype=int)

        return result