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

forge_orchestration/
builder.rs
