ferrum_engine/
recurrent_state.rs

1//! In-memory recurrent-state manager used for lifecycle and integration tests.
2//!
3//! This manager owns metadata handles and capacity accounting only. Backend
4//! implementations that carry actual CUDA/Metal tensors should provide their
5//! own handle type and downcast payload.
6
7use ferrum_interfaces::{
8    RecurrentStateHandle, RecurrentStateHandleStats, RecurrentStateManager,
9    RecurrentStateManagerStats, RecurrentStateSpec,
10};
11use ferrum_types::{FerrumError, RequestId, Result};
12use parking_lot::Mutex;
13use std::{
14    any::Any,
15    collections::HashMap,
16    sync::{
17        atomic::{AtomicBool, AtomicU64, Ordering},
18        Arc,
19    },
20    time::Instant,
21};
22
23#[derive(Debug, Clone)]
24pub struct InMemoryRecurrentStateConfig {
25    pub total_memory_bytes: usize,
26    pub total_batch_slots: usize,
27}
28
29impl Default for InMemoryRecurrentStateConfig {
30    fn default() -> Self {
31        Self {
32            total_memory_bytes: usize::MAX,
33            total_batch_slots: usize::MAX,
34        }
35    }
36}
37
38#[derive(Debug, Clone)]
39pub struct InMemoryRecurrentStateHandle {
40    spec: RecurrentStateSpec,
41    memory_bytes: usize,
42    cache_id: String,
43    valid: Arc<AtomicBool>,
44    created_at: Instant,
45}
46
47impl InMemoryRecurrentStateHandle {
48    fn new(spec: RecurrentStateSpec) -> Self {
49        let memory_bytes = spec.estimated_memory_bytes();
50        let cache_id = format!("recurrent-state-{}", spec.request_id);
51        Self {
52            spec,
53            memory_bytes,
54            cache_id,
55            valid: Arc::new(AtomicBool::new(true)),
56            created_at: Instant::now(),
57        }
58    }
59
60    fn invalidate(&self) {
61        self.valid.store(false, Ordering::Relaxed);
62    }
63}
64
65impl RecurrentStateHandle for InMemoryRecurrentStateHandle {
66    fn request_id(&self) -> RequestId {
67        self.spec.request_id.clone()
68    }
69
70    fn device(&self) -> ferrum_types::Device {
71        self.spec.device.clone()
72    }
73
74    fn num_layers(&self) -> usize {
75        self.spec.num_layers
76    }
77
78    fn state_bytes(&self) -> usize {
79        self.memory_bytes
80    }
81
82    fn clone_handle(&self) -> Result<Arc<dyn RecurrentStateHandle>> {
83        Ok(Arc::new(self.clone()))
84    }
85
86    fn as_any(&self) -> &dyn Any {
87        self
88    }
89
90    fn stats(&self) -> RecurrentStateHandleStats {
91        RecurrentStateHandleStats {
92            memory_bytes: self.memory_bytes,
93            state_tensors: self.spec.tensors.len(),
94            batch_slots: self.spec.max_batch_slots,
95            last_access: self.created_at,
96        }
97    }
98
99    fn is_valid(&self) -> bool {
100        self.valid.load(Ordering::Relaxed)
101    }
102
103    fn cache_id(&self) -> String {
104        self.cache_id.clone()
105    }
106}
107
108#[derive(Debug)]
109pub struct InMemoryRecurrentStateManager {
110    config: InMemoryRecurrentStateConfig,
111    handles: Mutex<HashMap<RequestId, Arc<InMemoryRecurrentStateHandle>>>,
112    allocation_count: AtomicU64,
113    allocation_failures: AtomicU64,
114}
115
116impl InMemoryRecurrentStateManager {
117    pub fn new(config: InMemoryRecurrentStateConfig) -> Self {
118        Self {
119            config,
120            handles: Mutex::new(HashMap::new()),
121            allocation_count: AtomicU64::new(0),
122            allocation_failures: AtomicU64::new(0),
123        }
124    }
125
126    fn used_memory_bytes_locked(
127        handles: &HashMap<RequestId, Arc<InMemoryRecurrentStateHandle>>,
128    ) -> usize {
129        handles.values().map(|handle| handle.memory_bytes).sum()
130    }
131
132    fn used_batch_slots_locked(
133        handles: &HashMap<RequestId, Arc<InMemoryRecurrentStateHandle>>,
134    ) -> usize {
135        handles
136            .values()
137            .map(|handle| handle.spec.max_batch_slots)
138            .sum()
139    }
140}
141
142#[async_trait::async_trait]
143impl RecurrentStateManager for InMemoryRecurrentStateManager {
144    async fn allocate(&self, spec: &RecurrentStateSpec) -> Result<Arc<dyn RecurrentStateHandle>> {
145        let mut handles = self.handles.lock();
146        if handles.contains_key(&spec.request_id) {
147            self.allocation_failures.fetch_add(1, Ordering::Relaxed);
148            return Err(FerrumError::already_exists(format!(
149                "recurrent state already allocated for {}",
150                spec.request_id
151            )));
152        }
153
154        let projected_memory =
155            Self::used_memory_bytes_locked(&handles).saturating_add(spec.estimated_memory_bytes());
156        let projected_slots =
157            Self::used_batch_slots_locked(&handles).saturating_add(spec.max_batch_slots);
158        if projected_memory > self.config.total_memory_bytes
159            || projected_slots > self.config.total_batch_slots
160        {
161            self.allocation_failures.fetch_add(1, Ordering::Relaxed);
162            return Err(FerrumError::resource_exhausted(
163                "insufficient recurrent-state capacity",
164            ));
165        }
166
167        let handle = Arc::new(InMemoryRecurrentStateHandle::new(spec.clone()));
168        handles.insert(spec.request_id.clone(), handle.clone());
169        self.allocation_count.fetch_add(1, Ordering::Relaxed);
170        Ok(handle)
171    }
172
173    async fn deallocate(&self, request_id: RequestId) -> Result<()> {
174        if let Some(handle) = self.handles.lock().remove(&request_id) {
175            handle.invalidate();
176        }
177        Ok(())
178    }
179
180    fn can_allocate(&self, spec: &RecurrentStateSpec) -> bool {
181        let handles = self.handles.lock();
182        if handles.contains_key(&spec.request_id) {
183            return false;
184        }
185        Self::used_memory_bytes_locked(&handles).saturating_add(spec.estimated_memory_bytes())
186            <= self.config.total_memory_bytes
187            && Self::used_batch_slots_locked(&handles).saturating_add(spec.max_batch_slots)
188                <= self.config.total_batch_slots
189    }
190
191    fn get_handle(&self, request_id: RequestId) -> Option<Arc<dyn RecurrentStateHandle>> {
192        self.handles
193            .lock()
194            .get(&request_id)
195            .map(|handle| handle.clone() as Arc<dyn RecurrentStateHandle>)
196    }
197
198    fn list_handles(&self) -> Vec<(RequestId, Arc<dyn RecurrentStateHandle>)> {
199        self.handles
200            .lock()
201            .iter()
202            .map(|(request_id, handle)| {
203                (
204                    request_id.clone(),
205                    handle.clone() as Arc<dyn RecurrentStateHandle>,
206                )
207            })
208            .collect()
209    }
210
211    fn stats(&self) -> RecurrentStateManagerStats {
212        let handles = self.handles.lock();
213        RecurrentStateManagerStats {
214            total_memory_bytes: self.config.total_memory_bytes,
215            used_memory_bytes: Self::used_memory_bytes_locked(&handles),
216            active_states: handles.len(),
217            active_state_tensors: handles
218                .values()
219                .map(|handle| handle.spec.tensors.len())
220                .sum(),
221            total_batch_slots: self.config.total_batch_slots,
222            used_batch_slots: Self::used_batch_slots_locked(&handles),
223            allocation_count: self.allocation_count.load(Ordering::Relaxed),
224            allocation_failures: self.allocation_failures.load(Ordering::Relaxed),
225            eviction_count: 0,
226        }
227    }
228
229    async fn reset(&self) -> Result<()> {
230        let mut handles = self.handles.lock();
231        for handle in handles.values() {
232            handle.invalidate();
233        }
234        handles.clear();
235        Ok(())
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use ferrum_interfaces::{RecurrentStateManager, RecurrentStateTensorSpec};
243    use ferrum_types::{DataType, Device};
244
245    fn spec(request_id: RequestId) -> RecurrentStateSpec {
246        RecurrentStateSpec {
247            request_id,
248            num_layers: 2,
249            tensors: vec![
250                RecurrentStateTensorSpec::new(0, "delta_state", vec![8, 16], DataType::BF16),
251                RecurrentStateTensorSpec::new(1, "delta_state", vec![8, 16], DataType::BF16),
252            ],
253            device: Device::CPU,
254            max_batch_slots: 1,
255        }
256    }
257
258    #[tokio::test]
259    async fn in_memory_manager_allocates_and_deallocates_state() {
260        let manager = InMemoryRecurrentStateManager::new(InMemoryRecurrentStateConfig {
261            total_memory_bytes: 4096,
262            total_batch_slots: 4,
263        });
264        let request_id = RequestId::new();
265        let spec = spec(request_id.clone());
266
267        let handle = manager.allocate(&spec).await.unwrap();
268
269        assert_eq!(handle.request_id(), request_id);
270        assert_eq!(handle.state_bytes(), 512);
271        assert!(handle.is_valid());
272        assert_eq!(manager.stats().active_states, 1);
273        assert_eq!(manager.stats().used_memory_bytes, 512);
274
275        manager.deallocate(request_id.clone()).await.unwrap();
276
277        assert!(!handle.is_valid());
278        assert!(manager.get_handle(request_id).is_none());
279        assert_eq!(manager.stats().active_states, 0);
280    }
281
282    #[tokio::test]
283    async fn in_memory_manager_rejects_duplicate_and_capacity_overcommit() {
284        let manager = InMemoryRecurrentStateManager::new(InMemoryRecurrentStateConfig {
285            total_memory_bytes: 512,
286            total_batch_slots: 1,
287        });
288        let request_id = RequestId::new();
289        let first = spec(request_id.clone());
290        let second = spec(RequestId::new());
291
292        manager.allocate(&first).await.unwrap();
293
294        let duplicate = manager.allocate(&first).await.unwrap_err();
295        assert!(matches!(duplicate, FerrumError::AlreadyExists { .. }));
296
297        let overcommit = manager.allocate(&second).await.unwrap_err();
298        assert!(matches!(overcommit, FerrumError::ResourceExhausted { .. }));
299        assert_eq!(manager.stats().allocation_failures, 2);
300    }
301
302    #[tokio::test]
303    async fn in_memory_manager_reset_invalidates_all_handles() {
304        let manager = InMemoryRecurrentStateManager::new(InMemoryRecurrentStateConfig {
305            total_memory_bytes: 4096,
306            total_batch_slots: 4,
307        });
308        let handle_a = manager.allocate(&spec(RequestId::new())).await.unwrap();
309        let handle_b = manager.allocate(&spec(RequestId::new())).await.unwrap();
310
311        manager.reset().await.unwrap();
312
313        assert!(!handle_a.is_valid());
314        assert!(!handle_b.is_valid());
315        assert!(manager.list_handles().is_empty());
316    }
317}
ferrum_engine/recurrent_state.rs

ferrum_engine/
recurrent_state.rs
