forge_orchestration/
builder.rs

1//! ForgeBuilder for configuring and constructing Forge instances
2//!
3//! ## Table of Contents
4//! - **ForgeBuilder**: Builder pattern for Forge configuration
5//! - **ForgeConfig**: Complete configuration struct
6
7use crate::autoscaler::{Autoscaler, AutoscalerConfig};
8use crate::error::Result;
9use crate::metrics::ForgeMetrics;
10use crate::moe::{BoxedMoERouter, DefaultMoERouter, MoERouter};
11use crate::networking::{HttpServerConfig, QuicConfig};
12use crate::nomad::NomadClient;
13use crate::runtime::Forge;
14use crate::storage::{BoxedStateStore, FileStore, MemoryStore};
15use crate::types::Region;
16use std::path::PathBuf;
17use std::sync::Arc;
18use tracing::info;
19
20/// Complete Forge configuration
21#[derive(Debug, Clone)]
22pub struct ForgeConfig {
23    /// Nomad API endpoint
24    pub nomad_api: Option<String>,
25    /// Nomad ACL token
26    pub nomad_token: Option<String>,
27    /// etcd endpoints
28    pub etcd_endpoints: Vec<String>,
29    /// Store path for file-based persistence
30    pub store_path: Option<PathBuf>,
31    /// HTTP server config
32    pub http_config: HttpServerConfig,
33    /// QUIC config
34    pub quic_config: QuicConfig,
35    /// Autoscaler config
36    pub autoscaler_config: AutoscalerConfig,
37    /// Federation regions
38    pub federation_regions: Vec<Region>,
39    /// Enable metrics
40    pub metrics_enabled: bool,
41    /// Node name
42    pub node_name: String,
43    /// Datacenter
44    pub datacenter: String,
45}
46
47impl Default for ForgeConfig {
48    fn default() -> Self {
49        Self {
50            nomad_api: None,
51            nomad_token: None,
52            etcd_endpoints: Vec::new(),
53            store_path: None,
54            http_config: HttpServerConfig::default(),
55            quic_config: QuicConfig::default(),
56            autoscaler_config: AutoscalerConfig::default(),
57            federation_regions: Vec::new(),
58            metrics_enabled: true,
59            node_name: hostname::get()
60                .map(|h| h.to_string_lossy().to_string())
61                .unwrap_or_else(|_| "forge-node".to_string()),
62            datacenter: "dc1".to_string(),
63        }
64    }
65}
66
67/// Builder for constructing Forge instances
68pub struct ForgeBuilder {
69    config: ForgeConfig,
70    router: Option<BoxedMoERouter>,
71    store: Option<BoxedStateStore>,
72}
73
74impl ForgeBuilder {
75    /// Create a new ForgeBuilder with default configuration
76    pub fn new() -> Self {
77        Self {
78            config: ForgeConfig::default(),
79            router: None,
80            store: None,
81        }
82    }
83
84    /// Set the Nomad API endpoint
85    pub fn with_nomad_api(mut self, url: impl Into<String>) -> Self {
86        self.config.nomad_api = Some(url.into());
87        self
88    }
89
90    /// Set the Nomad ACL token
91    pub fn with_nomad_token(mut self, token: impl Into<String>) -> Self {
92        self.config.nomad_token = Some(token.into());
93        self
94    }
95
96    /// Set etcd endpoints
97    pub fn with_etcd_endpoints(mut self, endpoints: Vec<impl Into<String>>) -> Self {
98        self.config.etcd_endpoints = endpoints.into_iter().map(|e| e.into()).collect();
99        self
100    }
101
102    /// Set file store path for local storage
103    pub fn with_store_path(mut self, path: impl Into<PathBuf>) -> Self {
104        self.config.store_path = Some(path.into());
105        self
106    }
107
108    /// Set the MoE router
109    pub fn with_moe_router<R: MoERouter + 'static>(mut self, router: R) -> Self {
110        self.router = Some(Arc::new(router));
111        self
112    }
113
114    /// Set a custom state store
115    pub fn with_store(mut self, store: BoxedStateStore) -> Self {
116        self.store = Some(store);
117        self
118    }
119
120    /// Set autoscaler configuration
121    pub fn with_autoscaler(mut self, config: AutoscalerConfig) -> Self {
122        self.config.autoscaler_config = config;
123        self
124    }
125
126    /// Set HTTP server configuration
127    pub fn with_http_config(mut self, config: HttpServerConfig) -> Self {
128        self.config.http_config = config;
129        self
130    }
131
132    /// Set HTTP bind address
133    pub fn with_http_addr(mut self, addr: &str) -> Result<Self> {
134        self.config.http_config = self.config.http_config.with_addr_str(addr)?;
135        Ok(self)
136    }
137
138    /// Set QUIC configuration
139    pub fn with_quic_config(mut self, config: QuicConfig) -> Self {
140        self.config.quic_config = config;
141        self
142    }
143
144    /// Enable multi-region federation
145    pub fn with_federation(mut self, regions: Vec<impl Into<Region>>) -> Self {
146        self.config.federation_regions = regions.into_iter().map(|r| r.into()).collect();
147        self
148    }
149
150    /// Set node name
151    pub fn with_node_name(mut self, name: impl Into<String>) -> Self {
152        self.config.node_name = name.into();
153        self
154    }
155
156    /// Set datacenter
157    pub fn with_datacenter(mut self, dc: impl Into<String>) -> Self {
158        self.config.datacenter = dc.into();
159        self
160    }
161
162    /// Enable or disable metrics
163    pub fn with_metrics(mut self, enabled: bool) -> Self {
164        self.config.metrics_enabled = enabled;
165        self
166    }
167
168    /// Build the Forge instance
169    pub fn build(self) -> Result<Forge> {
170        info!(
171            node = %self.config.node_name,
172            dc = %self.config.datacenter,
173            "Building Forge instance"
174        );
175
176        // Create Nomad client if configured
177        let nomad = match &self.config.nomad_api {
178            Some(url) => {
179                let mut client = NomadClient::new(url)?;
180                if let Some(token) = &self.config.nomad_token {
181                    client = client.with_token(token);
182                }
183                Some(client)
184            }
185            None => None,
186        };
187
188        // Create or use provided state store
189        let store: BoxedStateStore = match self.store {
190            Some(s) => s,
191            None => {
192                if let Some(path) = &self.config.store_path {
193                    Arc::new(FileStore::open(path)?) as BoxedStateStore
194                } else {
195                    Arc::new(MemoryStore::new()) as BoxedStateStore
196                }
197            }
198        };
199
200        // Create MoE router
201        let router = self.router.unwrap_or_else(|| Arc::new(DefaultMoERouter::new()));
202
203        // Create autoscaler
204        let autoscaler = Autoscaler::new(self.config.autoscaler_config.clone())?;
205
206        // Create metrics
207        let metrics = if self.config.metrics_enabled {
208            Some(Arc::new(ForgeMetrics::new()?))
209        } else {
210            None
211        };
212
213        Ok(Forge::new(
214            self.config,
215            nomad,
216            store,
217            router,
218            autoscaler,
219            metrics,
220        ))
221    }
222}
223
224impl Default for ForgeBuilder {
225    fn default() -> Self {
226        Self::new()
227    }
228}
229
230#[cfg(test)]
231mod tests {
232    use super::*;
233
234    #[test]
235    fn test_builder_default() {
236        let forge = ForgeBuilder::new().build();
237        assert!(forge.is_ok());
238    }
239
240    #[test]
241    fn test_builder_with_nomad() {
242        let forge = ForgeBuilder::new()
243            .with_nomad_api("http://localhost:4646")
244            .with_nomad_token("secret-token")
245            .build();
246        assert!(forge.is_ok());
247    }
248
249    #[test]
250    fn test_builder_with_autoscaler() {
251        let config = AutoscalerConfig::default()
252            .upscale_threshold(0.9)
253            .downscale_threshold(0.2);
254
255        let forge = ForgeBuilder::new().with_autoscaler(config).build();
256        assert!(forge.is_ok());
257    }
258
259    #[test]
260    fn test_builder_with_custom_router() {
261        use crate::moe::RoundRobinMoERouter;
262
263        let forge = ForgeBuilder::new()
264            .with_moe_router(RoundRobinMoERouter::new())
265            .build();
266        assert!(forge.is_ok());
267    }
268}
forge_orchestration/builder.rs

forge_orchestration/
builder.rs
