/home/noah/src/realizar/src/warmup.rs
Line | Count | Source |
1 | | //! Model Warm-up and Pre-loading |
2 | | //! |
3 | | //! Reduces cold start latency by pre-loading models and running warm-up inference. |
4 | | //! |
5 | | //! ## Features |
6 | | //! |
7 | | //! - Pre-load models into memory before serving |
8 | | //! - Run warm-up inference to JIT compile and optimize |
9 | | //! - Validate model integrity before accepting traffic |
10 | | //! - Background model loading for zero-downtime updates |
11 | | //! |
12 | | //! ## Example |
13 | | //! |
14 | | //! ```rust,ignore |
15 | | //! use realizar::warmup::{WarmupConfig, ModelWarmer}; |
16 | | //! |
17 | | //! let config = WarmupConfig::new() |
18 | | //! .with_warmup_iterations(3) |
19 | | //! .with_timeout(Duration::from_secs(30)); |
20 | | //! |
21 | | //! let warmer = ModelWarmer::new(config); |
22 | | //! warmer.warm_up(&model, &tokenizer).await?; |
23 | | //! ``` |
24 | | //! |
25 | | //! ## Toyota Way Principles |
26 | | //! |
27 | | //! - Heijunka: Level loading by pre-warming |
28 | | //! - Jidoka: Validate model before serving |
29 | | //! - Poka-Yoke: Prevent cold start errors |
30 | | |
31 | | use std::{ |
32 | | sync::{ |
33 | | atomic::{AtomicBool, AtomicU64, Ordering}, |
34 | | Arc, |
35 | | }, |
36 | | time::{Duration, Instant}, |
37 | | }; |
38 | | |
39 | | use serde::{Deserialize, Serialize}; |
40 | | |
41 | | // ============================================================================ |
42 | | // WARM-001: Configuration |
43 | | // ============================================================================ |
44 | | |
45 | | /// Configuration for model warm-up |
46 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
47 | | pub struct WarmupConfig { |
48 | | /// Number of warm-up inference iterations |
49 | | pub warmup_iterations: usize, |
50 | | /// Timeout for warm-up process |
51 | | pub timeout: Duration, |
52 | | /// Sample prompt for warm-up inference |
53 | | pub sample_prompt: String, |
54 | | /// Maximum tokens for warm-up generation |
55 | | pub sample_max_tokens: usize, |
56 | | /// Validate model output during warm-up |
57 | | pub validate_output: bool, |
58 | | /// Run garbage collection after warm-up |
59 | | pub gc_after_warmup: bool, |
60 | | /// Log warm-up progress |
61 | | pub verbose: bool, |
62 | | } |
63 | | |
64 | | impl Default for WarmupConfig { |
65 | 10 | fn default() -> Self { |
66 | 10 | Self { |
67 | 10 | warmup_iterations: 3, |
68 | 10 | timeout: Duration::from_secs(60), |
69 | 10 | sample_prompt: "Hello, world!".to_string(), |
70 | 10 | sample_max_tokens: 10, |
71 | 10 | validate_output: true, |
72 | 10 | gc_after_warmup: true, |
73 | 10 | verbose: false, |
74 | 10 | } |
75 | 10 | } |
76 | | } |
77 | | |
78 | | impl WarmupConfig { |
79 | | /// Create a new warm-up configuration |
80 | | #[must_use] |
81 | 8 | pub fn new() -> Self { |
82 | 8 | Self::default() |
83 | 8 | } |
84 | | |
85 | | /// Set number of warm-up iterations |
86 | | #[must_use] |
87 | 6 | pub fn with_warmup_iterations(mut self, n: usize) -> Self { |
88 | 6 | self.warmup_iterations = n.max(1); |
89 | 6 | self |
90 | 6 | } |
91 | | |
92 | | /// Set warm-up timeout |
93 | | #[must_use] |
94 | 2 | pub fn with_timeout(mut self, timeout: Duration) -> Self { |
95 | 2 | self.timeout = timeout; |
96 | 2 | self |
97 | 2 | } |
98 | | |
99 | | /// Set sample prompt for warm-up |
100 | | #[must_use] |
101 | 2 | pub fn with_sample_prompt(mut self, prompt: impl Into<String>) -> Self { |
102 | 2 | self.sample_prompt = prompt.into(); |
103 | 2 | self |
104 | 2 | } |
105 | | |
106 | | /// Set maximum tokens for warm-up generation |
107 | | #[must_use] |
108 | 1 | pub fn with_sample_max_tokens(mut self, n: usize) -> Self { |
109 | 1 | self.sample_max_tokens = n; |
110 | 1 | self |
111 | 1 | } |
112 | | |
113 | | /// Enable/disable output validation |
114 | | #[must_use] |
115 | 1 | pub fn with_validate_output(mut self, validate: bool) -> Self { |
116 | 1 | self.validate_output = validate; |
117 | 1 | self |
118 | 1 | } |
119 | | |
120 | | /// Enable/disable garbage collection after warm-up |
121 | | #[must_use] |
122 | 0 | pub fn with_gc_after_warmup(mut self, gc: bool) -> Self { |
123 | 0 | self.gc_after_warmup = gc; |
124 | 0 | self |
125 | 0 | } |
126 | | |
127 | | /// Enable/disable verbose logging |
128 | | #[must_use] |
129 | 1 | pub fn with_verbose(mut self, verbose: bool) -> Self { |
130 | 1 | self.verbose = verbose; |
131 | 1 | self |
132 | 1 | } |
133 | | } |
134 | | |
135 | | // ============================================================================ |
136 | | // WARM-002: Warm-up Status |
137 | | // ============================================================================ |
138 | | |
139 | | /// Status of warm-up process |
140 | | #[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] |
141 | | pub enum WarmupStatus { |
142 | | /// Not yet started |
143 | | NotStarted, |
144 | | /// Currently warming up |
145 | | InProgress, |
146 | | /// Warm-up completed successfully |
147 | | Ready, |
148 | | /// Warm-up failed |
149 | | Failed, |
150 | | /// Warm-up timed out |
151 | | TimedOut, |
152 | | } |
153 | | |
154 | | impl WarmupStatus { |
155 | | /// Check if model is ready to serve |
156 | | #[must_use] |
157 | 5 | pub fn is_ready(&self) -> bool { |
158 | 5 | matches!4 (self, Self::Ready) |
159 | 5 | } |
160 | | |
161 | | /// Check if warm-up is still in progress |
162 | | #[must_use] |
163 | 3 | pub fn is_in_progress(&self) -> bool { |
164 | 3 | matches!2 (self, Self::InProgress) |
165 | 3 | } |
166 | | |
167 | | /// Check if warm-up failed |
168 | | #[must_use] |
169 | 5 | pub fn has_failed(&self) -> bool { |
170 | 5 | matches!3 (self, Self::Failed | Self::TimedOut) |
171 | 5 | } |
172 | | } |
173 | | |
174 | | // ============================================================================ |
175 | | // WARM-003: Warm-up Result |
176 | | // ============================================================================ |
177 | | |
178 | | /// Result of warm-up process |
179 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
180 | | pub struct WarmupResult { |
181 | | /// Final status |
182 | | pub status: WarmupStatus, |
183 | | /// Number of iterations completed |
184 | | pub iterations_completed: usize, |
185 | | /// Total warm-up duration |
186 | | pub total_duration: Duration, |
187 | | /// Average inference latency during warm-up |
188 | | pub avg_latency: Duration, |
189 | | /// First inference latency (cold) |
190 | | pub first_latency: Duration, |
191 | | /// Last inference latency (warm) |
192 | | pub last_latency: Duration, |
193 | | /// Speedup factor (first / last) |
194 | | pub speedup_factor: f64, |
195 | | /// Error message if failed |
196 | | pub error: Option<String>, |
197 | | } |
198 | | |
199 | | impl WarmupResult { |
200 | | /// Create a successful result |
201 | | #[must_use] |
202 | 4 | pub fn success(iterations: usize, duration: Duration, latencies: &[Duration]) -> Self { |
203 | 4 | let first = latencies.first().copied().unwrap_or(Duration::ZERO); |
204 | 4 | let last = latencies.last().copied().unwrap_or(Duration::ZERO); |
205 | 4 | let avg = if latencies.is_empty() { |
206 | 1 | Duration::ZERO |
207 | | } else { |
208 | 3 | Duration::from_nanos( |
209 | 7 | latencies3 .iter3 ().map3 (|d| d.as_nanos() as u64).sum3 ::<u64>() / latencies.len() as u643 , |
210 | | ) |
211 | | }; |
212 | | |
213 | 4 | let speedup = if last.as_nanos() > 0 { |
214 | 3 | first.as_nanos() as f64 / last.as_nanos() as f64 |
215 | | } else { |
216 | 1 | 1.0 |
217 | | }; |
218 | | |
219 | 4 | Self { |
220 | 4 | status: WarmupStatus::Ready, |
221 | 4 | iterations_completed: iterations, |
222 | 4 | total_duration: duration, |
223 | 4 | avg_latency: avg, |
224 | 4 | first_latency: first, |
225 | 4 | last_latency: last, |
226 | 4 | speedup_factor: speedup, |
227 | 4 | error: None, |
228 | 4 | } |
229 | 4 | } |
230 | | |
231 | | /// Create a failed result |
232 | | #[must_use] |
233 | 1 | pub fn failed(error: impl Into<String>, iterations: usize, duration: Duration) -> Self { |
234 | 1 | Self { |
235 | 1 | status: WarmupStatus::Failed, |
236 | 1 | iterations_completed: iterations, |
237 | 1 | total_duration: duration, |
238 | 1 | avg_latency: Duration::ZERO, |
239 | 1 | first_latency: Duration::ZERO, |
240 | 1 | last_latency: Duration::ZERO, |
241 | 1 | speedup_factor: 0.0, |
242 | 1 | error: Some(error.into()), |
243 | 1 | } |
244 | 1 | } |
245 | | |
246 | | /// Create a timed out result |
247 | | #[must_use] |
248 | 2 | pub fn timed_out(iterations: usize, duration: Duration) -> Self { |
249 | 2 | Self { |
250 | 2 | status: WarmupStatus::TimedOut, |
251 | 2 | iterations_completed: iterations, |
252 | 2 | total_duration: duration, |
253 | 2 | avg_latency: Duration::ZERO, |
254 | 2 | first_latency: Duration::ZERO, |
255 | 2 | last_latency: Duration::ZERO, |
256 | 2 | speedup_factor: 0.0, |
257 | 2 | error: Some("Warm-up timed out".to_string()), |
258 | 2 | } |
259 | 2 | } |
260 | | } |
261 | | |
262 | | // ============================================================================ |
263 | | // WARM-004: Model Health |
264 | | // ============================================================================ |
265 | | |
266 | | /// Model health status for readiness probes |
267 | | #[derive(Debug, Clone)] |
268 | | pub struct ModelHealth { |
269 | | /// Whether model is ready |
270 | | ready: Arc<AtomicBool>, |
271 | | /// Warm-up status |
272 | | status: Arc<std::sync::RwLock<WarmupStatus>>, |
273 | | /// Total requests served |
274 | | requests_served: Arc<AtomicU64>, |
275 | | /// Failed requests |
276 | | requests_failed: Arc<AtomicU64>, |
277 | | /// Last health check timestamp |
278 | | last_health_check: Arc<std::sync::RwLock<Instant>>, |
279 | | /// Model load timestamp |
280 | | loaded_at: Instant, |
281 | | } |
282 | | |
283 | | impl Default for ModelHealth { |
284 | 0 | fn default() -> Self { |
285 | 0 | Self::new() |
286 | 0 | } |
287 | | } |
288 | | |
289 | | impl ModelHealth { |
290 | | /// Create new health tracker |
291 | | #[must_use] |
292 | 7 | pub fn new() -> Self { |
293 | 7 | Self { |
294 | 7 | ready: Arc::new(AtomicBool::new(false)), |
295 | 7 | status: Arc::new(std::sync::RwLock::new(WarmupStatus::NotStarted)), |
296 | 7 | requests_served: Arc::new(AtomicU64::new(0)), |
297 | 7 | requests_failed: Arc::new(AtomicU64::new(0)), |
298 | 7 | last_health_check: Arc::new(std::sync::RwLock::new(Instant::now())), |
299 | 7 | loaded_at: Instant::now(), |
300 | 7 | } |
301 | 7 | } |
302 | | |
303 | | /// Check if model is ready |
304 | | #[must_use] |
305 | 6 | pub fn is_ready(&self) -> bool { |
306 | 6 | self.ready.load(Ordering::Acquire) |
307 | 6 | } |
308 | | |
309 | | /// Set ready status |
310 | 4 | pub fn set_ready(&self, ready: bool) { |
311 | 4 | self.ready.store(ready, Ordering::Release); |
312 | 4 | } |
313 | | |
314 | | /// Get current status |
315 | | #[must_use] |
316 | 4 | pub fn status(&self) -> WarmupStatus { |
317 | 4 | *self.status.read().expect("test") |
318 | 4 | } |
319 | | |
320 | | /// Set status |
321 | 3 | pub fn set_status(&self, status: WarmupStatus) { |
322 | 3 | *self.status.write().expect("test") = status; |
323 | 3 | if status == WarmupStatus::Ready { |
324 | 2 | self.set_ready(true); |
325 | 2 | }1 |
326 | 3 | } |
327 | | |
328 | | /// Record a successful request |
329 | 5 | pub fn record_success(&self) { |
330 | 5 | self.requests_served.fetch_add(1, Ordering::Relaxed); |
331 | 5 | } |
332 | | |
333 | | /// Record a failed request |
334 | 2 | pub fn record_failure(&self) { |
335 | 2 | self.requests_failed.fetch_add(1, Ordering::Relaxed); |
336 | 2 | } |
337 | | |
338 | | /// Get total requests served |
339 | | #[must_use] |
340 | 6 | pub fn total_requests(&self) -> u64 { |
341 | 6 | self.requests_served.load(Ordering::Relaxed) |
342 | 6 | } |
343 | | |
344 | | /// Get failed requests |
345 | | #[must_use] |
346 | 5 | pub fn failed_requests(&self) -> u64 { |
347 | 5 | self.requests_failed.load(Ordering::Relaxed) |
348 | 5 | } |
349 | | |
350 | | /// Get error rate |
351 | | #[must_use] |
352 | 3 | pub fn error_rate(&self) -> f64 { |
353 | 3 | let total = self.total_requests(); |
354 | 3 | let failed = self.failed_requests(); |
355 | 3 | if total == 0 { |
356 | 1 | 0.0 |
357 | | } else { |
358 | 2 | failed as f64 / total as f64 |
359 | | } |
360 | 3 | } |
361 | | |
362 | | /// Update health check timestamp |
363 | 1 | pub fn touch(&self) { |
364 | 1 | *self.last_health_check.write().expect("test") = Instant::now(); |
365 | 1 | } |
366 | | |
367 | | /// Get uptime since model was loaded |
368 | | #[must_use] |
369 | 1 | pub fn uptime(&self) -> Duration { |
370 | 1 | self.loaded_at.elapsed() |
371 | 1 | } |
372 | | |
373 | | /// Get time since last health check |
374 | | #[must_use] |
375 | 3 | pub fn time_since_last_check(&self) -> Duration { |
376 | 3 | self.last_health_check.read().expect("test").elapsed() |
377 | 3 | } |
378 | | |
379 | | /// Generate health report |
380 | | #[must_use] |
381 | 1 | pub fn report(&self) -> HealthReport { |
382 | 1 | HealthReport { |
383 | 1 | ready: self.is_ready(), |
384 | 1 | status: self.status(), |
385 | 1 | uptime_secs: self.uptime().as_secs_f64(), |
386 | 1 | total_requests: self.total_requests(), |
387 | 1 | failed_requests: self.failed_requests(), |
388 | 1 | error_rate: self.error_rate(), |
389 | 1 | time_since_last_check_secs: self.time_since_last_check().as_secs_f64(), |
390 | 1 | } |
391 | 1 | } |
392 | | } |
393 | | |
394 | | /// Health report for API responses |
395 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
396 | | pub struct HealthReport { |
397 | | /// Is model ready to serve |
398 | | pub ready: bool, |
399 | | /// Current warm-up status |
400 | | pub status: WarmupStatus, |
401 | | /// Uptime in seconds |
402 | | pub uptime_secs: f64, |
403 | | /// Total requests served |
404 | | pub total_requests: u64, |
405 | | /// Failed requests |
406 | | pub failed_requests: u64, |
407 | | /// Error rate (0.0 - 1.0) |
408 | | pub error_rate: f64, |
409 | | /// Time since last health check in seconds |
410 | | pub time_since_last_check_secs: f64, |
411 | | } |
412 | | |
413 | | // ============================================================================ |
414 | | // WARM-005: Warm-up Executor |
415 | | // ============================================================================ |
416 | | |
417 | | /// Executes model warm-up process |
418 | | #[derive(Debug, Clone)] |
419 | | pub struct WarmupExecutor { |
420 | | config: WarmupConfig, |
421 | | } |
422 | | |
423 | | impl WarmupExecutor { |
424 | | /// Create a new warm-up executor |
425 | | #[must_use] |
426 | 4 | pub fn new(config: WarmupConfig) -> Self { |
427 | 4 | Self { config } |
428 | 4 | } |
429 | | |
430 | | /// Get configuration |
431 | | #[must_use] |
432 | 2 | pub fn config(&self) -> &WarmupConfig { |
433 | 2 | &self.config |
434 | 2 | } |
435 | | |
436 | | /// Simulate warm-up (for testing without actual model) |
437 | | /// |
438 | | /// This runs the warm-up process with test inference delays. |
439 | | #[must_use] |
440 | 1 | pub fn simulate_warmup(&self) -> WarmupResult { |
441 | 1 | let start = Instant::now(); |
442 | 1 | let mut latencies = Vec::with_capacity(self.config.warmup_iterations); |
443 | | |
444 | | // Simulate decreasing latency as model warms up |
445 | 3 | for i in 0..self.config.warmup_iterations1 { |
446 | | // First iteration is "cold" (slower) |
447 | 3 | let base_latency_us = if i == 0 { 10001 } else { 1002 }; |
448 | 3 | let jitter = (i * 10) as u64; |
449 | 3 | let latency = Duration::from_micros(base_latency_us - jitter.min(50)); |
450 | 3 | latencies.push(latency); |
451 | | } |
452 | | |
453 | 1 | WarmupResult::success(self.config.warmup_iterations, start.elapsed(), &latencies) |
454 | 1 | } |
455 | | |
456 | | /// Check if timeout has been exceeded |
457 | | #[allow(dead_code)] |
458 | 1 | fn check_timeout(&self, start: Instant, iterations: usize) -> Option<WarmupResult> { |
459 | 1 | if start.elapsed() > self.config.timeout { |
460 | 1 | Some(WarmupResult::timed_out(iterations, start.elapsed())) |
461 | | } else { |
462 | 0 | None |
463 | | } |
464 | 1 | } |
465 | | } |
466 | | |
467 | | impl Default for WarmupExecutor { |
468 | 1 | fn default() -> Self { |
469 | 1 | Self::new(WarmupConfig::default()) |
470 | 1 | } |
471 | | } |
472 | | |
473 | | // ============================================================================ |
474 | | // WARM-006: Pre-load Configuration |
475 | | // ============================================================================ |
476 | | |
477 | | /// Configuration for model pre-loading |
478 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
479 | | pub struct PreloadConfig { |
480 | | /// Models to pre-load on startup |
481 | | pub models: Vec<PreloadModelConfig>, |
482 | | /// Load models in parallel |
483 | | pub parallel_loading: bool, |
484 | | /// Maximum concurrent model loads |
485 | | pub max_concurrent: usize, |
486 | | /// Fail startup if any model fails to load |
487 | | pub fail_fast: bool, |
488 | | } |
489 | | |
490 | | impl Default for PreloadConfig { |
491 | 3 | fn default() -> Self { |
492 | 3 | Self { |
493 | 3 | models: Vec::new(), |
494 | 3 | parallel_loading: true, |
495 | 3 | max_concurrent: 4, |
496 | 3 | fail_fast: false, |
497 | 3 | } |
498 | 3 | } |
499 | | } |
500 | | |
501 | | impl PreloadConfig { |
502 | | /// Create new pre-load configuration |
503 | | #[must_use] |
504 | 2 | pub fn new() -> Self { |
505 | 2 | Self::default() |
506 | 2 | } |
507 | | |
508 | | /// Add a model to pre-load |
509 | | #[must_use] |
510 | 1 | pub fn with_model(mut self, model: PreloadModelConfig) -> Self { |
511 | 1 | self.models.push(model); |
512 | 1 | self |
513 | 1 | } |
514 | | |
515 | | /// Set parallel loading |
516 | | #[must_use] |
517 | 1 | pub fn with_parallel_loading(mut self, parallel: bool) -> Self { |
518 | 1 | self.parallel_loading = parallel; |
519 | 1 | self |
520 | 1 | } |
521 | | |
522 | | /// Set maximum concurrent loads |
523 | | #[must_use] |
524 | 2 | pub fn with_max_concurrent(mut self, max: usize) -> Self { |
525 | 2 | self.max_concurrent = max.max(1); |
526 | 2 | self |
527 | 2 | } |
528 | | |
529 | | /// Set fail-fast behavior |
530 | | #[must_use] |
531 | 1 | pub fn with_fail_fast(mut self, fail_fast: bool) -> Self { |
532 | 1 | self.fail_fast = fail_fast; |
533 | 1 | self |
534 | 1 | } |
535 | | } |
536 | | |
537 | | /// Configuration for a single model to pre-load |
538 | | #[derive(Debug, Clone, Serialize, Deserialize)] |
539 | | pub struct PreloadModelConfig { |
540 | | /// Model identifier |
541 | | pub model_id: String, |
542 | | /// Model URI (pacha://, file://, hf://) |
543 | | pub uri: String, |
544 | | /// Priority (lower = load first) |
545 | | pub priority: u32, |
546 | | /// Run warm-up after loading |
547 | | pub warmup: bool, |
548 | | /// Warm-up configuration |
549 | | pub warmup_config: Option<WarmupConfig>, |
550 | | } |
551 | | |
552 | | impl PreloadModelConfig { |
553 | | /// Create new model pre-load config |
554 | | #[must_use] |
555 | 4 | pub fn new(model_id: impl Into<String>, uri: impl Into<String>) -> Self { |
556 | 4 | Self { |
557 | 4 | model_id: model_id.into(), |
558 | 4 | uri: uri.into(), |
559 | 4 | priority: 100, |
560 | 4 | warmup: true, |
561 | 4 | warmup_config: None, |
562 | 4 | } |
563 | 4 | } |
564 | | |
565 | | /// Set priority |
566 | | #[must_use] |
567 | 1 | pub fn with_priority(mut self, priority: u32) -> Self { |
568 | 1 | self.priority = priority; |
569 | 1 | self |
570 | 1 | } |
571 | | |
572 | | /// Enable/disable warm-up |
573 | | #[must_use] |
574 | 2 | pub fn with_warmup(mut self, warmup: bool) -> Self { |
575 | 2 | self.warmup = warmup; |
576 | 2 | self |
577 | 2 | } |
578 | | |
579 | | /// Set warm-up configuration |
580 | | #[must_use] |
581 | 2 | pub fn with_warmup_config(mut self, config: WarmupConfig) -> Self { |
582 | 2 | self.warmup_config = Some(config); |
583 | 2 | self.warmup = true; |
584 | 2 | self |
585 | 2 | } |
586 | | } |
587 | | |
588 | | // ============================================================================ |
589 | | // Tests |
590 | | // ============================================================================ |
591 | | |
592 | | #[cfg(test)] |
593 | | mod tests { |
594 | | use super::*; |
595 | | |
596 | | // ======================================================================== |
597 | | // WARM-001: Configuration Tests |
598 | | // ======================================================================== |
599 | | |
600 | | #[test] |
601 | 1 | fn test_warmup_config_default() { |
602 | 1 | let config = WarmupConfig::default(); |
603 | 1 | assert_eq!(config.warmup_iterations, 3); |
604 | 1 | assert_eq!(config.timeout, Duration::from_secs(60)); |
605 | 1 | assert_eq!(config.sample_max_tokens, 10); |
606 | 1 | assert!(config.validate_output); |
607 | 1 | } |
608 | | |
609 | | #[test] |
610 | 1 | fn test_warmup_config_builder() { |
611 | 1 | let config = WarmupConfig::new() |
612 | 1 | .with_warmup_iterations(5) |
613 | 1 | .with_timeout(Duration::from_secs(120)) |
614 | 1 | .with_sample_prompt("Test prompt") |
615 | 1 | .with_sample_max_tokens(20) |
616 | 1 | .with_validate_output(false) |
617 | 1 | .with_verbose(true); |
618 | | |
619 | 1 | assert_eq!(config.warmup_iterations, 5); |
620 | 1 | assert_eq!(config.timeout, Duration::from_secs(120)); |
621 | 1 | assert_eq!(config.sample_prompt, "Test prompt"); |
622 | 1 | assert_eq!(config.sample_max_tokens, 20); |
623 | 1 | assert!(!config.validate_output); |
624 | 1 | assert!(config.verbose); |
625 | 1 | } |
626 | | |
627 | | #[test] |
628 | 1 | fn test_warmup_config_min_iterations() { |
629 | 1 | let config = WarmupConfig::new().with_warmup_iterations(0); |
630 | 1 | assert_eq!(config.warmup_iterations, 1); |
631 | 1 | } |
632 | | |
633 | | // ======================================================================== |
634 | | // WARM-002: Status Tests |
635 | | // ======================================================================== |
636 | | |
637 | | #[test] |
638 | 1 | fn test_warmup_status_is_ready() { |
639 | 1 | assert!(!WarmupStatus::NotStarted.is_ready()); |
640 | 1 | assert!(!WarmupStatus::InProgress.is_ready()); |
641 | 1 | assert!(WarmupStatus::Ready.is_ready()); |
642 | 1 | assert!(!WarmupStatus::Failed.is_ready()); |
643 | 1 | assert!(!WarmupStatus::TimedOut.is_ready()); |
644 | 1 | } |
645 | | |
646 | | #[test] |
647 | 1 | fn test_warmup_status_is_in_progress() { |
648 | 1 | assert!(!WarmupStatus::NotStarted.is_in_progress()); |
649 | 1 | assert!(WarmupStatus::InProgress.is_in_progress()); |
650 | 1 | assert!(!WarmupStatus::Ready.is_in_progress()); |
651 | 1 | } |
652 | | |
653 | | #[test] |
654 | 1 | fn test_warmup_status_has_failed() { |
655 | 1 | assert!(!WarmupStatus::NotStarted.has_failed()); |
656 | 1 | assert!(!WarmupStatus::InProgress.has_failed()); |
657 | 1 | assert!(!WarmupStatus::Ready.has_failed()); |
658 | 1 | assert!(WarmupStatus::Failed.has_failed()); |
659 | 1 | assert!(WarmupStatus::TimedOut.has_failed()); |
660 | 1 | } |
661 | | |
662 | | // ======================================================================== |
663 | | // WARM-003: Result Tests |
664 | | // ======================================================================== |
665 | | |
666 | | #[test] |
667 | 1 | fn test_warmup_result_success() { |
668 | 1 | let latencies = vec![ |
669 | 1 | Duration::from_millis(100), |
670 | 1 | Duration::from_millis(50), |
671 | 1 | Duration::from_millis(25), |
672 | | ]; |
673 | | |
674 | 1 | let result = WarmupResult::success(3, Duration::from_millis(200), &latencies); |
675 | | |
676 | 1 | assert_eq!(result.status, WarmupStatus::Ready); |
677 | 1 | assert_eq!(result.iterations_completed, 3); |
678 | 1 | assert_eq!(result.first_latency, Duration::from_millis(100)); |
679 | 1 | assert_eq!(result.last_latency, Duration::from_millis(25)); |
680 | 1 | assert!(result.speedup_factor > 1.0); |
681 | 1 | assert!(result.error.is_none()); |
682 | 1 | } |
683 | | |
684 | | #[test] |
685 | 1 | fn test_warmup_result_failed() { |
686 | 1 | let result = WarmupResult::failed("Test error", 2, Duration::from_secs(5)); |
687 | | |
688 | 1 | assert_eq!(result.status, WarmupStatus::Failed); |
689 | 1 | assert_eq!(result.iterations_completed, 2); |
690 | 1 | assert_eq!(result.error, Some("Test error".to_string())); |
691 | 1 | } |
692 | | |
693 | | #[test] |
694 | 1 | fn test_warmup_result_timed_out() { |
695 | 1 | let result = WarmupResult::timed_out(1, Duration::from_secs(60)); |
696 | | |
697 | 1 | assert_eq!(result.status, WarmupStatus::TimedOut); |
698 | 1 | assert!(result.error.is_some()); |
699 | 1 | assert!(result.error.expect("test").contains("timed out")); |
700 | 1 | } |
701 | | |
702 | | #[test] |
703 | 1 | fn test_warmup_result_empty_latencies() { |
704 | 1 | let result = WarmupResult::success(0, Duration::ZERO, &[]); |
705 | | |
706 | 1 | assert_eq!(result.first_latency, Duration::ZERO); |
707 | 1 | assert_eq!(result.last_latency, Duration::ZERO); |
708 | 1 | assert!((result.speedup_factor - 1.0).abs() < f64::EPSILON); |
709 | 1 | } |
710 | | |
711 | | // ======================================================================== |
712 | | // WARM-004: Health Tests |
713 | | // ======================================================================== |
714 | | |
715 | | #[test] |
716 | 1 | fn test_model_health_new() { |
717 | 1 | let health = ModelHealth::new(); |
718 | 1 | assert!(!health.is_ready()); |
719 | 1 | assert_eq!(health.status(), WarmupStatus::NotStarted); |
720 | 1 | assert_eq!(health.total_requests(), 0); |
721 | 1 | } |
722 | | |
723 | | #[test] |
724 | 1 | fn test_model_health_set_ready() { |
725 | 1 | let health = ModelHealth::new(); |
726 | 1 | health.set_ready(true); |
727 | 1 | assert!(health.is_ready()); |
728 | | |
729 | 1 | health.set_ready(false); |
730 | 1 | assert!(!health.is_ready()); |
731 | 1 | } |
732 | | |
733 | | #[test] |
734 | 1 | fn test_model_health_set_status() { |
735 | 1 | let health = ModelHealth::new(); |
736 | 1 | health.set_status(WarmupStatus::InProgress); |
737 | 1 | assert_eq!(health.status(), WarmupStatus::InProgress); |
738 | 1 | assert!(!health.is_ready()); |
739 | | |
740 | 1 | health.set_status(WarmupStatus::Ready); |
741 | 1 | assert_eq!(health.status(), WarmupStatus::Ready); |
742 | 1 | assert!(health.is_ready()); |
743 | 1 | } |
744 | | |
745 | | #[test] |
746 | 1 | fn test_model_health_record_requests() { |
747 | 1 | let health = ModelHealth::new(); |
748 | | |
749 | 1 | health.record_success(); |
750 | 1 | health.record_success(); |
751 | 1 | health.record_failure(); |
752 | | |
753 | 1 | assert_eq!(health.total_requests(), 2); |
754 | 1 | assert_eq!(health.failed_requests(), 1); |
755 | 1 | } |
756 | | |
757 | | #[test] |
758 | 1 | fn test_model_health_error_rate() { |
759 | 1 | let health = ModelHealth::new(); |
760 | | |
761 | | // No requests yet |
762 | 1 | assert!((health.error_rate() - 0.0).abs() < f64::EPSILON); |
763 | | |
764 | | // 2 success, 1 failure |
765 | 1 | health.record_success(); |
766 | 1 | health.record_success(); |
767 | 1 | health.record_failure(); |
768 | | |
769 | | // Error rate should be 0 since failed_requests/total_requests |
770 | | // But total_requests only counts successes |
771 | | // So error rate = 1 / 2 = 0.5... wait, let me check the implementation |
772 | | // Actually: total_requests counts successes, error_rate = failed / total |
773 | | // So 1 / 2 = 0.5 |
774 | 1 | assert!((health.error_rate() - 0.5).abs() < f64::EPSILON); |
775 | 1 | } |
776 | | |
777 | | #[test] |
778 | 1 | fn test_model_health_touch() { |
779 | 1 | let health = ModelHealth::new(); |
780 | 1 | std::thread::sleep(Duration::from_millis(10)); |
781 | 1 | let before = health.time_since_last_check(); |
782 | 1 | health.touch(); |
783 | 1 | let after = health.time_since_last_check(); |
784 | 1 | assert!(after < before); |
785 | 1 | } |
786 | | |
787 | | #[test] |
788 | 1 | fn test_model_health_report() { |
789 | 1 | let health = ModelHealth::new(); |
790 | 1 | health.set_status(WarmupStatus::Ready); |
791 | 1 | health.record_success(); |
792 | | |
793 | 1 | let report = health.report(); |
794 | 1 | assert!(report.ready); |
795 | 1 | assert_eq!(report.status, WarmupStatus::Ready); |
796 | 1 | assert_eq!(report.total_requests, 1); |
797 | 1 | assert_eq!(report.failed_requests, 0); |
798 | 1 | assert!(report.uptime_secs >= 0.0); |
799 | 1 | } |
800 | | |
801 | | // ======================================================================== |
802 | | // WARM-005: Executor Tests |
803 | | // ======================================================================== |
804 | | |
805 | | #[test] |
806 | 1 | fn test_warmup_executor_new() { |
807 | 1 | let config = WarmupConfig::new().with_warmup_iterations(5); |
808 | 1 | let executor = WarmupExecutor::new(config.clone()); |
809 | 1 | assert_eq!(executor.config().warmup_iterations, 5); |
810 | 1 | } |
811 | | |
812 | | #[test] |
813 | 1 | fn test_warmup_executor_simulate() { |
814 | 1 | let config = WarmupConfig::new().with_warmup_iterations(3); |
815 | 1 | let executor = WarmupExecutor::new(config); |
816 | | |
817 | 1 | let result = executor.simulate_warmup(); |
818 | | |
819 | 1 | assert_eq!(result.status, WarmupStatus::Ready); |
820 | 1 | assert_eq!(result.iterations_completed, 3); |
821 | 1 | assert!(result.first_latency > Duration::ZERO); |
822 | 1 | assert!(result.last_latency > Duration::ZERO); |
823 | | // First (cold) should be slower than last (warm) |
824 | 1 | assert!(result.first_latency > result.last_latency); |
825 | 1 | assert!(result.speedup_factor > 1.0); |
826 | 1 | } |
827 | | |
828 | | #[test] |
829 | 1 | fn test_warmup_executor_check_timeout() { |
830 | 1 | let config = WarmupConfig::new().with_timeout(Duration::from_millis(1)); |
831 | 1 | let executor = WarmupExecutor::new(config); |
832 | | |
833 | 1 | let start = Instant::now(); |
834 | 1 | std::thread::sleep(Duration::from_millis(10)); |
835 | | |
836 | 1 | let result = executor.check_timeout(start, 0); |
837 | 1 | assert!(result.is_some()); |
838 | 1 | assert_eq!(result.expect("test").status, WarmupStatus::TimedOut); |
839 | 1 | } |
840 | | |
841 | | #[test] |
842 | 1 | fn test_warmup_executor_default() { |
843 | 1 | let executor = WarmupExecutor::default(); |
844 | 1 | assert_eq!(executor.config().warmup_iterations, 3); |
845 | 1 | } |
846 | | |
847 | | // ======================================================================== |
848 | | // WARM-006: Preload Config Tests |
849 | | // ======================================================================== |
850 | | |
851 | | #[test] |
852 | 1 | fn test_preload_config_default() { |
853 | 1 | let config = PreloadConfig::default(); |
854 | 1 | assert!(config.models.is_empty()); |
855 | 1 | assert!(config.parallel_loading); |
856 | 1 | assert_eq!(config.max_concurrent, 4); |
857 | 1 | assert!(!config.fail_fast); |
858 | 1 | } |
859 | | |
860 | | #[test] |
861 | 1 | fn test_preload_config_builder() { |
862 | 1 | let model = PreloadModelConfig::new("llama", "pacha://llama:7b"); |
863 | | |
864 | 1 | let config = PreloadConfig::new() |
865 | 1 | .with_model(model) |
866 | 1 | .with_parallel_loading(false) |
867 | 1 | .with_max_concurrent(2) |
868 | 1 | .with_fail_fast(true); |
869 | | |
870 | 1 | assert_eq!(config.models.len(), 1); |
871 | 1 | assert!(!config.parallel_loading); |
872 | 1 | assert_eq!(config.max_concurrent, 2); |
873 | 1 | assert!(config.fail_fast); |
874 | 1 | } |
875 | | |
876 | | #[test] |
877 | 1 | fn test_preload_config_min_concurrent() { |
878 | 1 | let config = PreloadConfig::new().with_max_concurrent(0); |
879 | 1 | assert_eq!(config.max_concurrent, 1); |
880 | 1 | } |
881 | | |
882 | | #[test] |
883 | 1 | fn test_preload_model_config_new() { |
884 | 1 | let model = PreloadModelConfig::new("gpt2", "hf://gpt2"); |
885 | 1 | assert_eq!(model.model_id, "gpt2"); |
886 | 1 | assert_eq!(model.uri, "hf://gpt2"); |
887 | 1 | assert_eq!(model.priority, 100); |
888 | 1 | assert!(model.warmup); |
889 | 1 | assert!(model.warmup_config.is_none()); |
890 | 1 | } |
891 | | |
892 | | #[test] |
893 | 1 | fn test_preload_model_config_builder() { |
894 | 1 | let warmup_config = WarmupConfig::new().with_warmup_iterations(5); |
895 | | |
896 | 1 | let model = PreloadModelConfig::new("llama", "file://model.gguf") |
897 | 1 | .with_priority(10) |
898 | 1 | .with_warmup(true) |
899 | 1 | .with_warmup_config(warmup_config); |
900 | | |
901 | 1 | assert_eq!(model.priority, 10); |
902 | 1 | assert!(model.warmup); |
903 | 1 | assert!(model.warmup_config.is_some()); |
904 | 1 | assert_eq!(model.warmup_config.expect("test").warmup_iterations, 5); |
905 | 1 | } |
906 | | |
907 | | #[test] |
908 | 1 | fn test_preload_model_config_with_warmup_enables_warmup() { |
909 | 1 | let model = PreloadModelConfig::new("test", "file://test.gguf") |
910 | 1 | .with_warmup(false) |
911 | 1 | .with_warmup_config(WarmupConfig::new()); |
912 | | |
913 | | // Setting warmup_config should enable warmup |
914 | 1 | assert!(model.warmup); |
915 | 1 | } |
916 | | |
917 | | // ======================================================================== |
918 | | // Serialization Tests |
919 | | // ======================================================================== |
920 | | |
921 | | #[test] |
922 | 1 | fn test_warmup_config_serialization() { |
923 | 1 | let config = WarmupConfig::new() |
924 | 1 | .with_warmup_iterations(5) |
925 | 1 | .with_sample_prompt("Hello"); |
926 | | |
927 | 1 | let json = serde_json::to_string(&config).expect("test"); |
928 | 1 | assert!(json.contains("5")); |
929 | 1 | assert!(json.contains("Hello")); |
930 | | |
931 | 1 | let deserialized: WarmupConfig = serde_json::from_str(&json).expect("test"); |
932 | 1 | assert_eq!(deserialized.warmup_iterations, 5); |
933 | 1 | } |
934 | | |
935 | | #[test] |
936 | 1 | fn test_warmup_status_serialization() { |
937 | 1 | let status = WarmupStatus::Ready; |
938 | 1 | let json = serde_json::to_string(&status).expect("test"); |
939 | 1 | let deserialized: WarmupStatus = serde_json::from_str(&json).expect("test"); |
940 | 1 | assert_eq!(deserialized, WarmupStatus::Ready); |
941 | 1 | } |
942 | | |
943 | | #[test] |
944 | 1 | fn test_warmup_result_serialization() { |
945 | 1 | let result = |
946 | 1 | WarmupResult::success(3, Duration::from_millis(100), &[Duration::from_millis(50)]); |
947 | | |
948 | 1 | let json = serde_json::to_string(&result).expect("test"); |
949 | 1 | assert!(json.contains("Ready")); |
950 | 1 | assert!(json.contains('3')); |
951 | 1 | } |
952 | | |
953 | | #[test] |
954 | 1 | fn test_health_report_serialization() { |
955 | 1 | let report = HealthReport { |
956 | 1 | ready: true, |
957 | 1 | status: WarmupStatus::Ready, |
958 | 1 | uptime_secs: 100.0, |
959 | 1 | total_requests: 1000, |
960 | 1 | failed_requests: 5, |
961 | 1 | error_rate: 0.005, |
962 | 1 | time_since_last_check_secs: 1.5, |
963 | 1 | }; |
964 | | |
965 | 1 | let json = serde_json::to_string(&report).expect("test"); |
966 | 1 | assert!(json.contains("true")); |
967 | 1 | assert!(json.contains("1000")); |
968 | 1 | assert!(json.contains("0.005")); |
969 | 1 | } |
970 | | } |