aisimulate_core/replay/core/
mod.rs

1// SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4//! Router-neutral contracts for offline replay.
5//!
6//! Runtime adapters and concrete policies live one level up so this directory
7//! can later become a standalone crate without deployment-specific APIs.
8
9use anyhow::Result;
10use uuid::Uuid;
11
12pub mod round_robin;
13
14pub trait RequestIdentity {
15    fn request_id(&self) -> Option<Uuid>;
16
17    /// An authored attention-DP preference. Policies that do not provide
18    /// rank-affinity semantics may ignore it.
19    fn preferred_dp_rank(&self) -> Option<u32> {
20        None
21    }
22}
23
24#[derive(Debug)]
25pub struct ReadyArrival<Request, Metadata> {
26    pub request: Request,
27    pub arrival_time_ms: f64,
28    pub metadata: Metadata,
29    pub session_id: Option<String>,
30    pub turn_index: Option<usize>,
31}
32
33pub trait AdmissionSource {
34    type Request;
35    type Metadata;
36
37    fn next_ready_time_ms(&mut self) -> Option<f64>;
38    fn drain_ready(
39        &mut self,
40        now_ms: f64,
41        cluster_in_flight: usize,
42    ) -> Result<Vec<ReadyArrival<Self::Request, Self::Metadata>>>;
43    fn on_output_token(&mut self, request_id: Uuid, token_id: u32) -> Result<()>;
44    fn on_terminal(&mut self, request_id: Uuid, now_ms: f64, rejected: bool) -> Result<()>;
45    fn is_drained(&self) -> bool;
46    fn total_requests(&self) -> usize;
47}
48
49#[derive(Debug, Clone, Copy, PartialEq, Eq)]
50pub struct PlacementCacheSample {
51    pub overlap_blocks: u32,
52    pub isl_blocks: u32,
53}
54
55#[derive(Debug, Clone, Copy, PartialEq, Eq)]
56pub struct Placement {
57    pub request_id: Uuid,
58    pub scheduler_id: usize,
59    pub reported_overlap_tokens: usize,
60    pub cache_sample: Option<PlacementCacheSample>,
61}
62
63#[derive(Debug)]
64pub enum PlacementDecision {
65    Immediate(Placement),
66    Queued,
67}
68
69#[derive(Debug)]
70pub struct PlacementEffects {
71    pub decision: PlacementDecision,
72    pub released: Vec<Placement>,
73}
74
75#[derive(Debug, Clone, PartialEq, Eq)]
76pub struct WorkerTopology {
77    pub worker_id: usize,
78    pub scheduler_ids: Vec<usize>,
79}
80
81pub trait PlacementPolicy<Request> {
82    type Metadata;
83    type Observation;
84
85    fn place(
86        &mut self,
87        request: &Request,
88        metadata: Self::Metadata,
89        session_id: Option<String>,
90        now_ms: f64,
91    ) -> Result<PlacementEffects>;
92    fn observe(&mut self, observation: Self::Observation, now_ms: f64) -> Result<Vec<Placement>>;
93    fn cancel_pending(&mut self, request_id: Uuid) -> bool;
94    fn request_terminal(&mut self, request_id: Uuid, now_ms: f64) -> Result<Vec<Placement>>;
95    fn prefill_completed(&mut self, request_id: Uuid, now_ms: f64) -> Result<Vec<Placement>>;
96    fn pending_count(&self) -> usize;
97    fn worker_ready(&mut self, worker: WorkerTopology, now_ms: f64) -> Result<Vec<Placement>>;
98    fn worker_draining(&mut self, worker: WorkerTopology, now_ms: f64) -> Result<Vec<Placement>>;
99    fn worker_removed(&mut self, worker: WorkerTopology, now_ms: f64) -> Result<Vec<Placement>>;
100    fn topology_settled(&mut self, now_ms: f64) -> Result<Vec<Placement>>;
101}
102
103#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
104pub(crate) struct EngineProgress {
105    pub(crate) made_progress: bool,
106    pub(crate) had_raw_observations: bool,
107}
108
109pub trait EngineEventBatch: Default {
110    fn is_empty(&self) -> bool;
111    fn append(&mut self, other: Self);
112}
113
114impl EngineEventBatch for () {
115    #[inline]
116    fn is_empty(&self) -> bool {
117        true
118    }
119
120    #[inline]
121    fn append(&mut self, _other: Self) {}
122}
123
124#[derive(Debug, Default)]
125pub struct NoEngineEvents;
aisimulate_core/replay/core/mod.rs

aisimulate_core/replay/core/
mod.rs
