1use crate::error::{ErrorData, Result};
11use crate::instance_catalog::{is_same_architecture_machine, Architecture};
12use crate::resource::{ResourceDefinition, ResourceOutputsDefinition, ResourceRef};
13use crate::{PermissionProfile, Platform, ResourceType};
14use alien_error::AlienError;
15use bon::Builder;
16use serde::{Deserialize, Serialize};
17use std::any::Any;
18use std::collections::BTreeMap;
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 architecture: Option<Architecture>,
52 #[serde(skip_serializing_if = "Option::is_none")]
54 pub gpu: Option<GpuSpec>,
55}
56
57#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
59#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
60#[serde(rename_all = "camelCase")]
61pub struct ComputeChoiceRange {
62 pub min: u32,
64 pub max: u32,
66 pub default: u32,
68}
69
70impl ComputeChoiceRange {
71 pub fn contains(&self, value: u32) -> bool {
73 self.min <= value && value <= self.max
74 }
75}
76
77#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
79#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
80#[serde(rename_all = "camelCase", tag = "type")]
81pub enum CapacityGroupScalePolicy {
82 Fixed {
84 machines: ComputeChoiceRange,
86 },
87 Autoscale {
89 min: ComputeChoiceRange,
91 max: ComputeChoiceRange,
93 },
94}
95
96impl CapacityGroupScalePolicy {
97 pub fn from_selected_bounds(min_size: u32, max_size: u32) -> Self {
99 if min_size == max_size {
100 Self::Fixed {
101 machines: ComputeChoiceRange {
102 min: min_size,
103 max: max_size,
104 default: min_size,
105 },
106 }
107 } else {
108 Self::Autoscale {
109 min: ComputeChoiceRange {
110 min: min_size,
111 max: min_size,
112 default: min_size,
113 },
114 max: ComputeChoiceRange {
115 min: max_size,
116 max: max_size,
117 default: max_size,
118 },
119 }
120 }
121 }
122
123 pub fn default_min_size(&self) -> u32 {
125 match self {
126 Self::Fixed { machines } => machines.default,
127 Self::Autoscale { min, .. } => min.default,
128 }
129 }
130
131 pub fn default_max_size(&self) -> u32 {
133 match self {
134 Self::Fixed { machines } => machines.default,
135 Self::Autoscale { max, .. } => max.default,
136 }
137 }
138}
139
140#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
146#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
147#[serde(rename_all = "camelCase")]
148pub struct CapacityGroup {
149 pub group_id: String,
151 #[serde(skip_serializing_if = "Option::is_none")]
155 pub instance_type: Option<String>,
156 #[serde(skip_serializing_if = "Option::is_none")]
158 pub profile: Option<MachineProfile>,
159 pub min_size: u32,
161 pub max_size: u32,
163 #[serde(skip_serializing_if = "Option::is_none")]
168 pub scale_policy: Option<CapacityGroupScalePolicy>,
169 #[serde(skip_serializing_if = "Option::is_none")]
175 pub nested_virtualization: Option<bool>,
176}
177
178#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Builder)]
215#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
216#[serde(rename_all = "camelCase", deny_unknown_fields)]
217#[builder(start_fn = new)]
218pub struct ComputeCluster {
219 #[builder(start_fn)]
222 pub id: String,
223
224 #[builder(field)]
227 pub capacity_groups: Vec<CapacityGroup>,
228
229 #[serde(default, skip_serializing_if = "Option::is_none")]
232 pub node_permissions: Option<PermissionProfile>,
233
234 #[serde(default, skip_serializing_if = "Option::is_none")]
236 pub node_permissions_platforms: Option<Vec<Platform>>,
237
238 #[serde(default, skip_serializing_if = "Option::is_none")]
241 pub dynamic_container_pool: Option<String>,
242
243 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
246 #[builder(default)]
247 pub selected_failure_domains: BTreeMap<String, Vec<String>>,
248
249 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
252 #[builder(default)]
253 pub failure_domain_spread: BTreeMap<String, u8>,
254
255 #[serde(skip_serializing_if = "Option::is_none")]
259 pub container_cidr: Option<String>,
260}
261
262impl ComputeCluster {
263 pub const RESOURCE_TYPE: ResourceType = ResourceType::from_static("compute-cluster");
265
266 pub fn id(&self) -> &str {
268 &self.id
269 }
270
271 pub fn container_cidr(&self) -> &str {
273 self.container_cidr.as_deref().unwrap_or("10.244.0.0/16")
274 }
275}
276
277impl<S: compute_cluster_builder::State> ComputeClusterBuilder<S> {
278 pub fn capacity_group(mut self, group: CapacityGroup) -> Self {
280 self.capacity_groups.push(group);
281 self
282 }
283}
284
285#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
287#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
288#[serde(rename_all = "camelCase")]
289pub struct CapacityGroupStatus {
290 pub group_id: String,
292 pub current_machines: u32,
294 pub desired_machines: u32,
296 pub instance_type: String,
298}
299
300#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
302#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
303#[serde(rename_all = "camelCase")]
304pub struct ComputeClusterOutputs {
305 pub cluster_id: String,
307 pub horizon_ready: bool,
309 pub capacity_group_statuses: Vec<CapacityGroupStatus>,
311 pub total_machines: u32,
313}
314
315impl ResourceOutputsDefinition for ComputeClusterOutputs {
316 fn get_resource_type(&self) -> ResourceType {
317 ComputeCluster::RESOURCE_TYPE.clone()
318 }
319
320 fn as_any(&self) -> &dyn Any {
321 self
322 }
323
324 fn box_clone(&self) -> Box<dyn ResourceOutputsDefinition> {
325 Box::new(self.clone())
326 }
327
328 fn outputs_eq(&self, other: &dyn ResourceOutputsDefinition) -> bool {
329 other.as_any().downcast_ref::<ComputeClusterOutputs>() == Some(self)
330 }
331
332 fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
333 serde_json::to_value(self)
334 }
335}
336
337impl ResourceDefinition for ComputeCluster {
338 fn get_resource_type(&self) -> ResourceType {
339 Self::RESOURCE_TYPE
340 }
341
342 fn id(&self) -> &str {
343 &self.id
344 }
345
346 fn replace_after_failed_create_is_safe(&self) -> bool {
347 false
350 }
351
352 fn get_dependencies(&self) -> Vec<ResourceRef> {
353 Vec::new()
359 }
360
361 fn validate_update(&self, new_config: &dyn ResourceDefinition) -> Result<()> {
362 let new_cluster = new_config
363 .as_any()
364 .downcast_ref::<ComputeCluster>()
365 .ok_or_else(|| {
366 AlienError::new(ErrorData::UnexpectedResourceType {
367 resource_id: self.id.clone(),
368 expected: Self::RESOURCE_TYPE,
369 actual: new_config.get_resource_type(),
370 })
371 })?;
372
373 if self.id != new_cluster.id {
374 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
375 resource_id: self.id.clone(),
376 reason: "the 'id' field is immutable".to_string(),
377 }));
378 }
379
380 if self.container_cidr.is_some()
382 && new_cluster.container_cidr.is_some()
383 && self.container_cidr != new_cluster.container_cidr
384 {
385 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
386 resource_id: self.id.clone(),
387 reason: "the 'containerCidr' field is immutable once set".to_string(),
388 }));
389 }
390
391 for new_group in &new_cluster.capacity_groups {
393 if let Some(existing_group) = self
394 .capacity_groups
395 .iter()
396 .find(|g| g.group_id == new_group.group_id)
397 {
398 if let (Some(old), Some(new)) = (
402 existing_group.instance_type.as_deref(),
403 new_group.instance_type.as_deref(),
404 ) {
405 if old != new
406 && ![Platform::Aws, Platform::Gcp, Platform::Azure]
407 .into_iter()
408 .any(|platform| is_same_architecture_machine(platform, old, new))
409 {
410 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
411 resource_id: self.id.clone(),
412 reason: format!(
413 "capacity group '{}' can't change machine from '{old}' to '{new}': only a catalog machine in the same cloud and of the same CPU architecture can replace it",
414 new_group.group_id
415 ),
416 }));
417 }
418 }
419 }
420 }
421
422 Ok(())
423 }
424
425 fn as_any(&self) -> &dyn Any {
426 self
427 }
428
429 fn as_any_mut(&mut self) -> &mut dyn Any {
430 self
431 }
432
433 fn box_clone(&self) -> Box<dyn ResourceDefinition> {
434 Box::new(self.clone())
435 }
436
437 fn resource_eq(&self, other: &dyn ResourceDefinition) -> bool {
438 other.as_any().downcast_ref::<ComputeCluster>() == Some(self)
439 }
440
441 fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
442 serde_json::to_value(self)
443 }
444}
445
446#[cfg(test)]
447mod tests {
448 use super::*;
449
450 #[test]
451 fn test_compute_cluster_creation() {
452 let cluster = ComputeCluster::new("compute".to_string())
453 .capacity_group(CapacityGroup {
454 group_id: "general".to_string(),
455 instance_type: Some("m7g.xlarge".to_string()),
456 profile: None,
457 min_size: 1,
458 max_size: 5,
459 scale_policy: None,
460 nested_virtualization: None,
461 })
462 .build();
463
464 assert_eq!(cluster.id(), "compute");
465 assert_eq!(cluster.capacity_groups.len(), 1);
466 assert_eq!(cluster.capacity_groups[0].group_id, "general");
467 assert_eq!(cluster.container_cidr(), "10.244.0.0/16");
468 }
469
470 #[test]
471 fn test_compute_cluster_multiple_capacity_groups() {
472 let cluster = ComputeCluster::new("multi-pool".to_string())
473 .capacity_group(CapacityGroup {
474 group_id: "general".to_string(),
475 instance_type: Some("m7g.xlarge".to_string()),
476 profile: None,
477 min_size: 1,
478 max_size: 3,
479 scale_policy: None,
480 nested_virtualization: None,
481 })
482 .capacity_group(CapacityGroup {
483 group_id: "gpu".to_string(),
484 instance_type: Some("g5.xlarge".to_string()),
485 profile: Some(MachineProfile {
486 cpu: "4.0".to_string(),
487 memory_bytes: 17179869184, ephemeral_storage_bytes: 214748364800, architecture: None,
490 gpu: Some(GpuSpec {
491 gpu_type: "nvidia-a10g".to_string(),
492 count: 1,
493 }),
494 }),
495 min_size: 0,
496 max_size: 2,
497 scale_policy: None,
498 nested_virtualization: None,
499 })
500 .build();
501
502 assert_eq!(cluster.capacity_groups.len(), 2);
503 assert_eq!(cluster.capacity_groups[0].group_id, "general");
504 assert_eq!(cluster.capacity_groups[1].group_id, "gpu");
505 assert!(cluster.capacity_groups[1]
506 .profile
507 .as_ref()
508 .unwrap()
509 .gpu
510 .is_some());
511 }
512
513 #[test]
514 fn test_compute_cluster_custom_cidr() {
515 let cluster = ComputeCluster::new("custom-net".to_string())
516 .container_cidr("172.30.0.0/16".to_string())
517 .capacity_group(CapacityGroup {
518 group_id: "general".to_string(),
519 instance_type: None,
520 profile: None,
521 min_size: 1,
522 max_size: 5,
523 scale_policy: None,
524 nested_virtualization: None,
525 })
526 .build();
527
528 assert_eq!(cluster.container_cidr(), "172.30.0.0/16");
529 }
530
531 #[test]
532 fn test_compute_cluster_validate_update_immutable_id() {
533 let cluster1 = ComputeCluster::new("cluster-1".to_string())
534 .capacity_group(CapacityGroup {
535 group_id: "general".to_string(),
536 instance_type: None,
537 profile: None,
538 min_size: 1,
539 max_size: 5,
540 scale_policy: None,
541 nested_virtualization: None,
542 })
543 .build();
544
545 let cluster2 = ComputeCluster::new("cluster-2".to_string())
546 .capacity_group(CapacityGroup {
547 group_id: "general".to_string(),
548 instance_type: None,
549 profile: None,
550 min_size: 1,
551 max_size: 5,
552 scale_policy: None,
553 nested_virtualization: None,
554 })
555 .build();
556
557 let result = cluster1.validate_update(&cluster2);
558 assert!(result.is_err());
559 }
560
561 #[test]
562 fn test_compute_cluster_validate_update_scale_change() {
563 let cluster1 = ComputeCluster::new("compute".to_string())
564 .capacity_group(CapacityGroup {
565 group_id: "general".to_string(),
566 instance_type: Some("m7g.xlarge".to_string()),
567 profile: None,
568 min_size: 1,
569 max_size: 5,
570 scale_policy: None,
571 nested_virtualization: None,
572 })
573 .build();
574
575 let cluster2 = ComputeCluster::new("compute".to_string())
576 .capacity_group(CapacityGroup {
577 group_id: "general".to_string(),
578 instance_type: Some("m7g.xlarge".to_string()),
579 profile: None,
580 min_size: 2,
581 max_size: 10,
582 scale_policy: None,
583 nested_virtualization: None,
584 })
585 .build();
586
587 let result = cluster1.validate_update(&cluster2);
589 assert!(result.is_ok());
590 }
591
592 #[test]
593 fn machine_changes_stay_within_one_cloud_and_architecture() {
594 let cluster = |machine: &str| {
595 ComputeCluster::new("compute".to_string())
596 .capacity_group(CapacityGroup {
597 group_id: "general".to_string(),
598 instance_type: Some(machine.to_string()),
599 profile: None,
600 min_size: 1,
601 max_size: 5,
602 scale_policy: None,
603 nested_virtualization: None,
604 })
605 .build()
606 };
607 for (old, new) in [
608 ("t4g.small", "t4g.medium"),
609 ("e2-medium", "n2-standard-4"),
610 ("Standard_D2s_v5", "Standard_D4s_v5"),
611 ] {
612 cluster(old)
613 .validate_update(&cluster(new))
614 .expect("same cloud and architecture");
615 }
616 for (old, new) in [
617 ("t4g.small", "m7i.large"),
618 ("t4g.small", "t4g.unknown"),
619 ("n2-standard-2", "m7i.large"),
620 ("Standard_D2s_v5", "n2-standard-4"),
621 ] {
622 let error = cluster(old)
623 .validate_update(&cluster(new))
624 .expect_err("needs setup");
625 assert_eq!(error.code, "INVALID_RESOURCE_UPDATE", "{old} -> {new}");
626 }
627 }
628
629 #[test]
630 fn node_permissions_are_inline_optional_and_roundtrip_exactly() {
631 let mut cluster = ComputeCluster::new("compute".to_string()).build();
632 let absent = serde_json::to_value(&cluster).unwrap();
633 assert!(absent.get("nodePermissions").is_none());
634 assert_eq!(
635 serde_json::from_value::<ComputeCluster>(absent).unwrap(),
636 cluster
637 );
638 cluster.node_permissions =
639 Some(PermissionProfile::new().resource("objects", ["storage/data-read"]));
640 let json = serde_json::to_value(&cluster).unwrap();
641 assert_eq!(
642 json["nodePermissions"],
643 serde_json::json!({"objects": ["storage/data-read"]})
644 );
645 assert_eq!(
646 serde_json::from_value::<ComputeCluster>(json).unwrap(),
647 cluster
648 );
649 assert_eq!(ResourceDefinition::get_permissions(&cluster), None);
650 }
651
652 #[test]
653 fn test_compute_cluster_serialization() {
654 let cluster = ComputeCluster::new("test-cluster".to_string())
655 .capacity_group(CapacityGroup {
656 group_id: "general".to_string(),
657 instance_type: Some("m7g.xlarge".to_string()),
658 profile: None,
659 min_size: 1,
660 max_size: 5,
661 scale_policy: None,
662 nested_virtualization: None,
663 })
664 .build();
665
666 let json = serde_json::to_string(&cluster).unwrap();
667 let deserialized: ComputeCluster = serde_json::from_str(&json).unwrap();
668 assert_eq!(cluster, deserialized);
669 }
670}