forge_orchestration/
runtime.rs

1//! Forge runtime and control plane
2//!
3//! ## Table of Contents
4//! - **Forge**: Main runtime struct
5//! - **ForgeHandle**: Handle for interacting with running Forge
6
7use crate::autoscaler::{Autoscaler, MetricsSnapshot, ScalingDecision};
8use crate::builder::ForgeConfig;
9use crate::error::Result;
10use crate::job::Job;
11use crate::metrics::ForgeMetrics;
12use crate::moe::{BoxedMoERouter, RouteResult};
13use crate::networking::{HttpServer, HttpState};
14use crate::nomad::NomadClient;
15use crate::scheduler::{reconcile::Reconciler, NodeResources};
16use crate::storage::{keys, store_get_json, store_set_json, BoxedStateStore};
17use crate::types::{Expert, NodeId, Shard, ShardId};
18use axum::{
19    extract::{Path, State},
20    routing::{delete, get, post},
21    Json, Router,
22};
23use dashmap::DashMap;
24use std::sync::Arc;
25use std::time::Instant;
26use tokio::sync::{broadcast, RwLock};
27use tracing::{error, info, warn};
28
29/// Runtime state
30#[derive(Debug, Clone, Copy, PartialEq, Eq)]
31pub enum RuntimeState {
32    /// Not started
33    Stopped,
34    /// Starting up
35    Starting,
36    /// Running normally
37    Running,
38    /// Shutting down
39    ShuttingDown,
40}
41
42/// Main Forge runtime
43pub struct Forge {
44    config: ForgeConfig,
45    state: Arc<RwLock<RuntimeState>>,
46    node_id: NodeId,
47    start_time: Option<Instant>,
48
49    // Core components
50    nomad: Option<NomadClient>,
51    store: BoxedStateStore,
52    router: BoxedMoERouter,
53    autoscaler: Autoscaler,
54    metrics: Option<Arc<ForgeMetrics>>,
55
56    // Runtime state
57    jobs: DashMap<String, Job>,
58    experts: DashMap<usize, Expert>,
59    shards: DashMap<ShardId, Shard>,
60
61    // Nodes registered for the embedded scheduler / reconcile loop
62    scheduler_nodes: std::sync::Mutex<Vec<NodeResources>>,
63
64    // Shutdown signal
65    shutdown_tx: broadcast::Sender<()>,
66}
67
68impl Forge {
69    /// Create a new Forge instance (use ForgeBuilder instead)
70    pub(crate) fn new(
71        config: ForgeConfig,
72        nomad: Option<NomadClient>,
73        store: BoxedStateStore,
74        router: BoxedMoERouter,
75        autoscaler: Autoscaler,
76        metrics: Option<Arc<ForgeMetrics>>,
77    ) -> Self {
78        let (shutdown_tx, _) = broadcast::channel(1);
79
80        Self {
81            config,
82            state: Arc::new(RwLock::new(RuntimeState::Stopped)),
83            node_id: NodeId::new(),
84            start_time: None,
85            nomad,
86            store,
87            router,
88            autoscaler,
89            metrics,
90            jobs: DashMap::new(),
91            experts: DashMap::new(),
92            shards: DashMap::new(),
93            scheduler_nodes: std::sync::Mutex::new(Vec::new()),
94            shutdown_tx,
95        }
96    }
97
98    /// Get the node ID
99    pub fn node_id(&self) -> NodeId {
100        self.node_id
101    }
102
103    /// Get current runtime state
104    pub async fn state(&self) -> RuntimeState {
105        *self.state.read().await
106    }
107
108    /// Get uptime in seconds
109    pub fn uptime_secs(&self) -> u64 {
110        self.start_time
111            .map(|t| t.elapsed().as_secs())
112            .unwrap_or(0)
113    }
114
115    /// Get metrics instance
116    pub fn metrics(&self) -> Option<&Arc<ForgeMetrics>> {
117        self.metrics.as_ref()
118    }
119
120    /// Get the state store
121    pub fn store(&self) -> &BoxedStateStore {
122        &self.store
123    }
124
125    /// Get the MoE router
126    pub fn router(&self) -> &BoxedMoERouter {
127        &self.router
128    }
129
130    /// Check if Nomad is configured
131    pub fn has_nomad(&self) -> bool {
132        self.nomad.is_some()
133    }
134
135    /// Register a node available to the embedded scheduler / reconcile loop.
136    ///
137    /// Register nodes before [`Forge::run`], which seeds the reconcile loop with
138    /// the registered set.
139    pub fn register_node(&self, node: NodeResources) {
140        self.scheduler_nodes.lock().unwrap().push(node);
141    }
142
143    /// Build a [`Reconciler`] that shares this runtime's state store and is
144    /// seeded with the registered nodes. The reconciler converges submitted jobs
145    /// (persisted under `keys::JOBS`) onto node assignments.
146    ///
147    /// [`Forge::run`] spawns one automatically; this is exposed for tests and
148    /// advanced embedding. It uses a fresh autoscaler with the same config as the
149    /// runtime's — the reconciler is the autoscaling authority for the loop.
150    pub async fn new_reconciler(&self) -> Result<Reconciler> {
151        let nodes: Vec<NodeResources> = self.scheduler_nodes.lock().unwrap().clone();
152        let autoscaler = Arc::new(Autoscaler::new(self.autoscaler.config().clone())?);
153        let mut reconciler = Reconciler::new(self.store.clone(), autoscaler);
154        for node in nodes {
155            reconciler.register_node(node);
156        }
157        reconciler.bootstrap().await?;
158        Ok(reconciler)
159    }
160
161    /// Run the Forge control plane
162    pub async fn run(mut self) -> Result<()> {
163        {
164            let mut state = self.state.write().await;
165            *state = RuntimeState::Starting;
166        }
167
168        info!(
169            node_id = %self.node_id,
170            node_name = %self.config.node_name,
171            "Starting Forge control plane"
172        );
173
174        self.start_time = Some(Instant::now());
175
176        // Verify Nomad connectivity if configured
177        if let Some(nomad) = &self.nomad {
178            match nomad.health().await {
179                Ok(true) => info!("Nomad connection verified"),
180                Ok(false) => warn!("Nomad returned unhealthy status"),
181                Err(e) => warn!(error = %e, "Failed to connect to Nomad"),
182            }
183        }
184
185        // Load existing state from store
186        self.load_state().await?;
187
188        {
189            let mut state = self.state.write().await;
190            *state = RuntimeState::Running;
191        }
192
193        info!("Forge control plane running");
194
195        // Spawn the reconcile control loop. It shares the state store (so it sees
196        // jobs submitted via `submit_job`) and the nodes registered via
197        // `register_node`; with no nodes registered it idles harmlessly.
198        match self.new_reconciler().await {
199            Ok(reconciler) => {
200                let shutdown_rx = self.shutdown_tx.subscribe();
201                tokio::spawn(async move {
202                    if let Err(e) = reconciler.run(shutdown_rx).await {
203                        error!(error = %e, "reconcile loop exited with error");
204                    }
205                });
206                info!("Reconcile control loop spawned");
207            }
208            Err(e) => warn!(error = %e, "Failed to start reconcile loop"),
209        }
210
211        // Create shared state for HTTP handlers
212        let forge_state = Arc::new(RwLock::new(ForgeHttpState {
213            jobs: self.jobs.clone(),
214            metrics: self.metrics.clone(),
215        }));
216
217        // Build HTTP router
218        let http_router = self.build_http_router(forge_state);
219
220        // Start HTTP server
221        let http_server = HttpServer::new(self.config.http_config.clone())
222            .with_router(http_router);
223
224        // Run until shutdown
225        let mut shutdown_rx = self.shutdown_tx.subscribe();
226
227        tokio::select! {
228            result = http_server.serve() => {
229                if let Err(e) = result {
230                    error!(error = %e, "HTTP server error");
231                }
232            }
233            _ = shutdown_rx.recv() => {
234                info!("Shutdown signal received");
235            }
236        }
237
238        self.shutdown().await?;
239
240        Ok(())
241    }
242
243    /// Shutdown the control plane
244    pub async fn shutdown(&self) -> Result<()> {
245        {
246            let mut state = self.state.write().await;
247            if *state == RuntimeState::Stopped {
248                return Ok(());
249            }
250            *state = RuntimeState::ShuttingDown;
251        }
252
253        info!("Shutting down Forge control plane");
254
255        // Save state
256        self.save_state().await?;
257
258        // Send shutdown signal
259        let _ = self.shutdown_tx.send(());
260
261        {
262            let mut state = self.state.write().await;
263            *state = RuntimeState::Stopped;
264        }
265
266        info!("Forge control plane stopped");
267        Ok(())
268    }
269
270    /// Signal shutdown
271    pub fn signal_shutdown(&self) {
272        let _ = self.shutdown_tx.send(());
273    }
274
275    /// Subscribe to shutdown signal
276    pub fn shutdown_receiver(&self) -> broadcast::Receiver<()> {
277        self.shutdown_tx.subscribe()
278    }
279
280    // State management
281
282    async fn load_state(&self) -> Result<()> {
283        // Load jobs from store
284        let job_keys = self.store.list_prefix(keys::JOBS).await?;
285        for key in job_keys {
286            if let Some(job) = store_get_json::<Job>(self.store.as_ref(), &key).await? {
287                self.jobs.insert(job.id.clone(), job);
288            }
289        }
290
291        info!(jobs = self.jobs.len(), "Loaded state from store");
292        Ok(())
293    }
294
295    async fn save_state(&self) -> Result<()> {
296        // Save all jobs
297        for entry in self.jobs.iter() {
298            let key = keys::job(&entry.key());
299            store_set_json(self.store.as_ref(), &key, entry.value()).await?;
300        }
301
302        info!(jobs = self.jobs.len(), "Saved state to store");
303        Ok(())
304    }
305
306    // Job management
307
308    /// Submit a job
309    pub async fn submit_job(&self, job: Job) -> Result<String> {
310        let job_id = job.id.clone();
311
312        // Submit to Nomad if configured
313        if let Some(nomad) = &self.nomad {
314            nomad.submit_job(&job).await?;
315        }
316
317        // Store locally
318        let key = keys::job(&job_id);
319        store_set_json(self.store.as_ref(), &key, &job).await?;
320        self.jobs.insert(job_id.clone(), job);
321
322        if let Some(metrics) = &self.metrics {
323            metrics.record_job_submitted();
324        }
325
326        info!(job_id = %job_id, "Job submitted");
327        Ok(job_id)
328    }
329
330    /// Get a job by ID
331    pub fn get_job(&self, job_id: &str) -> Option<Job> {
332        self.jobs.get(job_id).map(|e| e.value().clone())
333    }
334
335    /// List all jobs
336    pub fn list_jobs(&self) -> Vec<Job> {
337        self.jobs.iter().map(|e| e.value().clone()).collect()
338    }
339
340    /// Stop a job
341    pub async fn stop_job(&self, job_id: &str, purge: bool) -> Result<()> {
342        // Stop in Nomad if configured
343        if let Some(nomad) = &self.nomad {
344            nomad.stop_job(job_id, purge).await?;
345        }
346
347        // Remove locally
348        if purge {
349            let key = keys::job(job_id);
350            self.store.delete(&key).await?;
351            self.jobs.remove(job_id);
352        }
353
354        if let Some(metrics) = &self.metrics {
355            metrics.record_job_completed(true);
356        }
357
358        info!(job_id = %job_id, purge = purge, "Job stopped");
359        Ok(())
360    }
361
362    /// Scale a job's task group
363    pub async fn scale_job(&self, job_id: &str, group: &str, count: u32) -> Result<()> {
364        // Scale in Nomad if configured
365        if let Some(nomad) = &self.nomad {
366            nomad
367                .scale_job(job_id, group, count, Some("Manual scale"))
368                .await?;
369        }
370
371        // Update local state
372        if let Some(mut job) = self.jobs.get_mut(job_id) {
373            for g in &mut job.groups {
374                if g.name == group {
375                    g.scaling.desired = count;
376                    break;
377                }
378            }
379        }
380
381        if let Some(metrics) = &self.metrics {
382            let direction = "manual";
383            metrics.record_scale_event(job_id, direction);
384            metrics.set_instances(job_id, group, count as f64);
385        }
386
387        info!(job_id = %job_id, group = %group, count = count, "Job scaled");
388        Ok(())
389    }
390
391    // MoE routing
392
393    /// Route an input to an expert
394    pub async fn route(&self, input: &str) -> RouteResult {
395        let experts: Vec<Expert> = self.experts.iter().map(|e| e.value().clone()).collect();
396
397        let result = if experts.is_empty() {
398            self.router.route(input, 8).await
399        } else {
400            self.router.route_with_experts(input, &experts).await
401        };
402
403        if let Some(metrics) = &self.metrics {
404            metrics.record_route(self.router.name(), result.expert_index, 0.001);
405        }
406
407        result
408    }
409
410    /// Register an expert
411    pub fn register_expert(&self, expert: Expert) {
412        info!(index = expert.index, node = %expert.node, "Expert registered");
413        self.experts.insert(expert.index, expert);
414    }
415
416    /// Update expert load
417    pub fn update_expert_load(&self, index: usize, load: f64) {
418        if let Some(mut expert) = self.experts.get_mut(&index) {
419            expert.update_load(load);
420        }
421    }
422
423    // Autoscaling
424
425    /// Evaluate autoscaling for a job
426    pub async fn evaluate_scaling(
427        &self,
428        job_id: &str,
429        cpu: f64,
430        memory: f64,
431        instances: u32,
432    ) -> ScalingDecision {
433        let metrics = MetricsSnapshot::new(cpu, memory, instances);
434        let decision = self.autoscaler.evaluate(job_id, metrics).await;
435
436        // Execute the decision (previously it was computed and discarded). Apply
437        // the delta to each of the job's groups via `scale_job`, which updates
438        // Nomad (if configured), local state, and metrics. The autoscaler already
439        // enforced hysteresis before returning a scaling action.
440        if decision.is_scaling() {
441            let groups: Vec<(String, u32)> = self
442                .jobs
443                .get(job_id)
444                .map(|j| {
445                    j.groups
446                        .iter()
447                        .map(|g| (g.name.clone(), g.scaling.desired))
448                        .collect()
449                })
450                .unwrap_or_default();
451
452            for (group, current) in groups {
453                let target = match &decision {
454                    ScalingDecision::ScaleUp(n) => current.saturating_add(*n),
455                    ScalingDecision::ScaleDown(n) => current.saturating_sub(*n),
456                    ScalingDecision::ScaleTo(c) => *c,
457                    ScalingDecision::NoChange => current,
458                };
459                if target != current {
460                    if let Err(e) = self.scale_job(job_id, &group, target).await {
461                        warn!(job_id = %job_id, group = %group, error = %e, "autoscale execution failed");
462                    }
463                }
464            }
465        }
466
467        decision
468    }
469
470    // HTTP router
471
472    fn build_http_router(&self, state: Arc<RwLock<ForgeHttpState>>) -> Router {
473        let state = HttpState { app: state };
474
475        Router::new()
476            .route("/health", get(health_handler))
477            .route("/ready", get(ready_handler))
478            .route("/api/v1/jobs", get(list_jobs_handler))
479            .route("/api/v1/jobs", post(submit_job_handler))
480            .route("/api/v1/jobs/:id", get(get_job_handler))
481            .route("/api/v1/jobs/:id", delete(stop_job_handler))
482            .route("/metrics", get(metrics_handler))
483            .with_state(state)
484    }
485}
486
487// HTTP state for handlers
488struct ForgeHttpState {
489    jobs: DashMap<String, Job>,
490    metrics: Option<Arc<ForgeMetrics>>,
491}
492
493// HTTP handlers
494
495async fn health_handler() -> Json<serde_json::Value> {
496    Json(serde_json::json!({
497        "status": "healthy",
498        "version": env!("CARGO_PKG_VERSION")
499    }))
500}
501
502async fn ready_handler() -> axum::http::StatusCode {
503    axum::http::StatusCode::OK
504}
505
506async fn list_jobs_handler(
507    State(state): State<HttpState<ForgeHttpState>>,
508) -> Json<Vec<Job>> {
509    let app = state.app.read().await;
510    let jobs: Vec<Job> = app.jobs.iter().map(|e| e.value().clone()).collect();
511    Json(jobs)
512}
513
514async fn get_job_handler(
515    State(state): State<HttpState<ForgeHttpState>>,
516    Path(id): Path<String>,
517) -> std::result::Result<Json<Job>, axum::http::StatusCode> {
518    let app = state.app.read().await;
519    app.jobs
520        .get(&id)
521        .map(|e| Json(e.value().clone()))
522        .ok_or(axum::http::StatusCode::NOT_FOUND)
523}
524
525async fn submit_job_handler(
526    State(state): State<HttpState<ForgeHttpState>>,
527    Json(job): Json<Job>,
528) -> std::result::Result<Json<serde_json::Value>, axum::http::StatusCode> {
529    let app = state.app.read().await;
530    let job_id = job.id.clone();
531    app.jobs.insert(job_id.clone(), job);
532    Ok(Json(serde_json::json!({ "job_id": job_id })))
533}
534
535async fn stop_job_handler(
536    State(state): State<HttpState<ForgeHttpState>>,
537    Path(id): Path<String>,
538) -> axum::http::StatusCode {
539    let app = state.app.read().await;
540    if app.jobs.remove(&id).is_some() {
541        axum::http::StatusCode::NO_CONTENT
542    } else {
543        axum::http::StatusCode::NOT_FOUND
544    }
545}
546
547async fn metrics_handler(
548    State(state): State<HttpState<ForgeHttpState>>,
549) -> std::result::Result<String, axum::http::StatusCode> {
550    let app = state.app.read().await;
551    match &app.metrics {
552        Some(m) => m
553            .gather_text()
554            .map_err(|_| axum::http::StatusCode::INTERNAL_SERVER_ERROR),
555        None => Err(axum::http::StatusCode::NOT_FOUND),
556    }
557}
558
559#[cfg(test)]
560mod tests {
561    use super::*;
562    use crate::builder::ForgeBuilder;
563    use crate::job::{Driver, Task};
564
565    #[tokio::test]
566    async fn test_forge_creation() {
567        let forge = ForgeBuilder::new().build().unwrap();
568        assert_eq!(forge.state().await, RuntimeState::Stopped);
569    }
570
571    #[tokio::test]
572    async fn test_job_management() {
573        let forge = ForgeBuilder::new().build().unwrap();
574
575        let job = Job::new("test-job").with_group(
576            "api",
577            Task::new("server")
578                .driver(Driver::Exec)
579                .command("/bin/server"),
580        );
581
582        let job_id = forge.submit_job(job).await.unwrap();
583        assert!(forge.get_job(&job_id).is_some());
584
585        let jobs = forge.list_jobs();
586        assert_eq!(jobs.len(), 1);
587
588        forge.stop_job(&job_id, true).await.unwrap();
589        assert!(forge.get_job(&job_id).is_none());
590    }
591
592    #[tokio::test]
593    async fn test_routing() {
594        let forge = ForgeBuilder::new().build().unwrap();
595
596        let result = forge.route("test-input").await;
597        assert!(result.expert_index < 8);
598    }
599
600    #[tokio::test]
601    async fn test_expert_registration() {
602        let forge = ForgeBuilder::new().build().unwrap();
603
604        let expert = Expert::new(0, NodeId::new());
605        forge.register_expert(expert);
606
607        forge.update_expert_load(0, 0.5);
608    }
609
610    #[tokio::test]
611    async fn test_evaluate_scaling_executes_decision() {
612        let forge = ForgeBuilder::new().build().unwrap();
613        let mut job = Job::new("svc").with_group(
614            "api",
615            Task::new("s").driver(Driver::Exec).command("/bin/s"),
616        );
617        job.groups[0].scaling = crate::job::ScalingConfig::new(1, 10).with_desired(2);
618        let id = forge.submit_job(job).await.unwrap();
619
620        // High utilization -> ScaleUp, and the decision is actually applied.
621        let decision = forge.evaluate_scaling(&id, 0.95, 0.5, 2).await;
622        assert!(matches!(decision, ScalingDecision::ScaleUp(_)));
623
624        let updated = forge.get_job(&id).unwrap();
625        assert_eq!(updated.groups[0].scaling.desired, 3, "scale decision must be executed");
626    }
627
628    #[tokio::test]
629    async fn test_reconciler_integration_schedules_jobs() {
630        let forge = ForgeBuilder::new().build().unwrap();
631        forge.register_node(NodeResources::new(NodeId::new(), 8000, 8192));
632
633        let job = Job::new("svc").with_group(
634            "api",
635            Task::new("s")
636                .driver(Driver::Exec)
637                .command("/bin/s")
638                .resources(1000, 1024),
639        );
640        forge.submit_job(job).await.unwrap();
641
642        // The reconciler shares the store, so it sees the submitted job and binds it.
643        let mut reconciler = forge.new_reconciler().await.unwrap();
644        let report = reconciler.reconcile_once().await.unwrap();
645        assert!(report.scheduled >= 1, "reconciler should schedule the submitted replica");
646        assert!(reconciler.bound_count() >= 1);
647    }
648}
forge_orchestration/runtime.rs

forge_orchestration/
runtime.rs
