1use crate::error::{ErrorData, Result};
13use crate::resource::{ResourceDefinition, ResourceOutputsDefinition, ResourceRef};
14use crate::ResourceType;
15use alien_error::AlienError;
16use bon::Builder;
17use serde::{Deserialize, Serialize};
18use std::any::Any;
19use std::fmt::Debug;
20
21#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
23#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
24#[serde(rename_all = "camelCase")]
25pub struct GpuSpec {
26 #[serde(rename = "type")]
28 pub gpu_type: String,
29 pub count: u32,
31}
32
33#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
39#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
40#[serde(rename_all = "camelCase")]
41pub struct MachineProfile {
42 pub cpu: String,
45 pub memory_bytes: u64,
47 pub ephemeral_storage_bytes: u64,
49 #[serde(skip_serializing_if = "Option::is_none")]
51 pub gpu: Option<GpuSpec>,
52}
53
54#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
60#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
61#[serde(rename_all = "camelCase")]
62pub struct CapacityGroup {
63 pub group_id: String,
65 #[serde(skip_serializing_if = "Option::is_none")]
68 pub instance_type: Option<String>,
69 #[serde(skip_serializing_if = "Option::is_none")]
71 pub profile: Option<MachineProfile>,
72 pub min_size: u32,
74 pub max_size: u32,
76}
77
78#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
89#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
90#[serde(rename_all = "camelCase")]
91pub struct TemplateInputs {
92 pub horizond_download_base_url: String,
94 pub horizon_api_url: String,
96 #[serde(skip_serializing_if = "Option::is_none")]
100 pub horizond_binary_hash: Option<String>,
101 #[serde(skip_serializing_if = "Option::is_none")]
103 pub monitoring_logs_endpoint: Option<String>,
104 #[serde(skip_serializing_if = "Option::is_none")]
108 pub monitoring_metrics_endpoint: Option<String>,
109 #[serde(skip_serializing_if = "Option::is_none")]
112 pub monitoring_auth_hash: Option<String>,
113 #[serde(skip_serializing_if = "Option::is_none")]
116 pub monitoring_metrics_auth_hash: Option<String>,
117}
118
119#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Builder)]
147#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
148#[serde(rename_all = "camelCase", deny_unknown_fields)]
149#[builder(start_fn = new)]
150pub struct ContainerCluster {
151 #[builder(start_fn)]
154 pub id: String,
155
156 #[builder(field)]
159 pub capacity_groups: Vec<CapacityGroup>,
160
161 #[serde(skip_serializing_if = "Option::is_none")]
165 pub container_cidr: Option<String>,
166
167 #[builder(skip)]
170 #[serde(skip_serializing_if = "Option::is_none")]
171 pub template_inputs: Option<TemplateInputs>,
172}
173
174impl ContainerCluster {
175 pub const RESOURCE_TYPE: ResourceType = ResourceType::from_static("container-cluster");
177
178 pub fn id(&self) -> &str {
180 &self.id
181 }
182
183 pub fn container_cidr(&self) -> &str {
185 self.container_cidr.as_deref().unwrap_or("10.244.0.0/16")
186 }
187}
188
189impl<S: container_cluster_builder::State> ContainerClusterBuilder<S> {
190 pub fn capacity_group(mut self, group: CapacityGroup) -> Self {
192 self.capacity_groups.push(group);
193 self
194 }
195}
196
197#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
199#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
200#[serde(rename_all = "camelCase")]
201pub struct CapacityGroupStatus {
202 pub group_id: String,
204 pub current_machines: u32,
206 pub desired_machines: u32,
208 pub instance_type: String,
210}
211
212#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
214#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
215#[serde(rename_all = "camelCase")]
216pub struct ContainerClusterOutputs {
217 pub cluster_id: String,
219 pub horizon_ready: bool,
221 pub capacity_group_statuses: Vec<CapacityGroupStatus>,
223 pub total_machines: u32,
225}
226
227impl ResourceOutputsDefinition for ContainerClusterOutputs {
228 fn get_resource_type(&self) -> ResourceType {
229 ContainerCluster::RESOURCE_TYPE.clone()
230 }
231
232 fn as_any(&self) -> &dyn Any {
233 self
234 }
235
236 fn box_clone(&self) -> Box<dyn ResourceOutputsDefinition> {
237 Box::new(self.clone())
238 }
239
240 fn outputs_eq(&self, other: &dyn ResourceOutputsDefinition) -> bool {
241 other.as_any().downcast_ref::<ContainerClusterOutputs>() == Some(self)
242 }
243
244 fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
245 serde_json::to_value(self)
246 }
247}
248
249impl ResourceDefinition for ContainerCluster {
250 fn get_resource_type(&self) -> ResourceType {
251 Self::RESOURCE_TYPE
252 }
253
254 fn id(&self) -> &str {
255 &self.id
256 }
257
258 fn get_dependencies(&self) -> Vec<ResourceRef> {
259 Vec::new()
265 }
266
267 fn validate_update(&self, new_config: &dyn ResourceDefinition) -> Result<()> {
268 let new_cluster = new_config
269 .as_any()
270 .downcast_ref::<ContainerCluster>()
271 .ok_or_else(|| {
272 AlienError::new(ErrorData::UnexpectedResourceType {
273 resource_id: self.id.clone(),
274 expected: Self::RESOURCE_TYPE,
275 actual: new_config.get_resource_type(),
276 })
277 })?;
278
279 if self.id != new_cluster.id {
280 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
281 resource_id: self.id.clone(),
282 reason: "the 'id' field is immutable".to_string(),
283 }));
284 }
285
286 if self.container_cidr.is_some()
288 && new_cluster.container_cidr.is_some()
289 && self.container_cidr != new_cluster.container_cidr
290 {
291 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
292 resource_id: self.id.clone(),
293 reason: "the 'containerCidr' field is immutable once set".to_string(),
294 }));
295 }
296
297 for new_group in &new_cluster.capacity_groups {
299 if let Some(existing_group) = self
300 .capacity_groups
301 .iter()
302 .find(|g| g.group_id == new_group.group_id)
303 {
304 if existing_group.instance_type.is_some()
306 && new_group.instance_type.is_some()
307 && existing_group.instance_type != new_group.instance_type
308 {
309 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
310 resource_id: self.id.clone(),
311 reason: format!(
312 "instance type for capacity group '{}' is immutable",
313 new_group.group_id
314 ),
315 }));
316 }
317 }
318 }
319
320 Ok(())
321 }
322
323 fn as_any(&self) -> &dyn Any {
324 self
325 }
326
327 fn as_any_mut(&mut self) -> &mut dyn Any {
328 self
329 }
330
331 fn box_clone(&self) -> Box<dyn ResourceDefinition> {
332 Box::new(self.clone())
333 }
334
335 fn resource_eq(&self, other: &dyn ResourceDefinition) -> bool {
336 other.as_any().downcast_ref::<ContainerCluster>() == Some(self)
337 }
338
339 fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
340 serde_json::to_value(self)
341 }
342}
343
344#[cfg(test)]
345mod tests {
346 use super::*;
347
348 #[test]
349 fn test_container_cluster_creation() {
350 let cluster = ContainerCluster::new("compute".to_string())
351 .capacity_group(CapacityGroup {
352 group_id: "general".to_string(),
353 instance_type: Some("m7g.xlarge".to_string()),
354 profile: None,
355 min_size: 1,
356 max_size: 5,
357 })
358 .build();
359
360 assert_eq!(cluster.id(), "compute");
361 assert_eq!(cluster.capacity_groups.len(), 1);
362 assert_eq!(cluster.capacity_groups[0].group_id, "general");
363 assert_eq!(cluster.container_cidr(), "10.244.0.0/16");
364 }
365
366 #[test]
367 fn test_container_cluster_multiple_capacity_groups() {
368 let cluster = ContainerCluster::new("multi-pool".to_string())
369 .capacity_group(CapacityGroup {
370 group_id: "general".to_string(),
371 instance_type: Some("m7g.xlarge".to_string()),
372 profile: None,
373 min_size: 1,
374 max_size: 3,
375 })
376 .capacity_group(CapacityGroup {
377 group_id: "gpu".to_string(),
378 instance_type: Some("g5.xlarge".to_string()),
379 profile: Some(MachineProfile {
380 cpu: "4.0".to_string(),
381 memory_bytes: 17179869184, ephemeral_storage_bytes: 214748364800, gpu: Some(GpuSpec {
384 gpu_type: "nvidia-a10g".to_string(),
385 count: 1,
386 }),
387 }),
388 min_size: 0,
389 max_size: 2,
390 })
391 .build();
392
393 assert_eq!(cluster.capacity_groups.len(), 2);
394 assert_eq!(cluster.capacity_groups[0].group_id, "general");
395 assert_eq!(cluster.capacity_groups[1].group_id, "gpu");
396 assert!(cluster.capacity_groups[1]
397 .profile
398 .as_ref()
399 .unwrap()
400 .gpu
401 .is_some());
402 }
403
404 #[test]
405 fn test_container_cluster_custom_cidr() {
406 let cluster = ContainerCluster::new("custom-net".to_string())
407 .container_cidr("172.30.0.0/16".to_string())
408 .capacity_group(CapacityGroup {
409 group_id: "general".to_string(),
410 instance_type: None,
411 profile: None,
412 min_size: 1,
413 max_size: 5,
414 })
415 .build();
416
417 assert_eq!(cluster.container_cidr(), "172.30.0.0/16");
418 }
419
420 #[test]
421 fn test_container_cluster_validate_update_immutable_id() {
422 let cluster1 = ContainerCluster::new("cluster-1".to_string())
423 .capacity_group(CapacityGroup {
424 group_id: "general".to_string(),
425 instance_type: None,
426 profile: None,
427 min_size: 1,
428 max_size: 5,
429 })
430 .build();
431
432 let cluster2 = ContainerCluster::new("cluster-2".to_string())
433 .capacity_group(CapacityGroup {
434 group_id: "general".to_string(),
435 instance_type: None,
436 profile: None,
437 min_size: 1,
438 max_size: 5,
439 })
440 .build();
441
442 let result = cluster1.validate_update(&cluster2);
443 assert!(result.is_err());
444 }
445
446 #[test]
447 fn test_container_cluster_validate_update_scale_change() {
448 let cluster1 = ContainerCluster::new("compute".to_string())
449 .capacity_group(CapacityGroup {
450 group_id: "general".to_string(),
451 instance_type: Some("m7g.xlarge".to_string()),
452 profile: None,
453 min_size: 1,
454 max_size: 5,
455 })
456 .build();
457
458 let cluster2 = ContainerCluster::new("compute".to_string())
459 .capacity_group(CapacityGroup {
460 group_id: "general".to_string(),
461 instance_type: Some("m7g.xlarge".to_string()),
462 profile: None,
463 min_size: 2,
464 max_size: 10,
465 })
466 .build();
467
468 let result = cluster1.validate_update(&cluster2);
470 assert!(result.is_ok());
471 }
472
473 #[test]
474 fn test_container_cluster_serialization() {
475 let cluster = ContainerCluster::new("test-cluster".to_string())
476 .capacity_group(CapacityGroup {
477 group_id: "general".to_string(),
478 instance_type: Some("m7g.xlarge".to_string()),
479 profile: None,
480 min_size: 1,
481 max_size: 5,
482 })
483 .build();
484
485 let json = serde_json::to_string(&cluster).unwrap();
486 let deserialized: ContainerCluster = serde_json::from_str(&json).unwrap();
487 assert_eq!(cluster, deserialized);
488 }
489}