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 #[serde(default, skip_serializing_if = "Option::is_none")]
59 pub data_disk: Option<DataDiskRequirement>,
60}
61
62#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
65#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
66#[serde(rename_all = "camelCase")]
67pub struct DataDiskRequirement {
68 pub size_bytes: u64,
70}
71
72#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
74#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
75#[serde(rename_all = "camelCase")]
76pub struct ComputeChoiceRange {
77 pub min: u32,
79 pub max: u32,
81 pub default: u32,
83}
84
85impl ComputeChoiceRange {
86 pub fn contains(&self, value: u32) -> bool {
88 self.min <= value && value <= self.max
89 }
90}
91
92#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
94#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
95#[serde(rename_all = "camelCase", tag = "type")]
96pub enum CapacityGroupScalePolicy {
97 Fixed {
99 machines: ComputeChoiceRange,
101 },
102 Autoscale {
104 min: ComputeChoiceRange,
106 max: ComputeChoiceRange,
108 },
109}
110
111impl CapacityGroupScalePolicy {
112 pub fn from_selected_bounds(min_size: u32, max_size: u32) -> Self {
114 if min_size == max_size {
115 Self::Fixed {
116 machines: ComputeChoiceRange {
117 min: min_size,
118 max: max_size,
119 default: min_size,
120 },
121 }
122 } else {
123 Self::Autoscale {
124 min: ComputeChoiceRange {
125 min: min_size,
126 max: min_size,
127 default: min_size,
128 },
129 max: ComputeChoiceRange {
130 min: max_size,
131 max: max_size,
132 default: max_size,
133 },
134 }
135 }
136 }
137
138 pub fn default_min_size(&self) -> u32 {
140 match self {
141 Self::Fixed { machines } => machines.default,
142 Self::Autoscale { min, .. } => min.default,
143 }
144 }
145
146 pub fn default_max_size(&self) -> u32 {
148 match self {
149 Self::Fixed { machines } => machines.default,
150 Self::Autoscale { max, .. } => max.default,
151 }
152 }
153}
154
155#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
161#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
162#[serde(rename_all = "camelCase")]
163pub struct CapacityGroup {
164 pub group_id: String,
166 #[serde(skip_serializing_if = "Option::is_none")]
170 pub instance_type: Option<String>,
171 #[serde(skip_serializing_if = "Option::is_none")]
173 pub profile: Option<MachineProfile>,
174 pub min_size: u32,
176 pub max_size: u32,
178 #[serde(skip_serializing_if = "Option::is_none")]
183 pub scale_policy: Option<CapacityGroupScalePolicy>,
184 #[serde(skip_serializing_if = "Option::is_none")]
190 pub nested_virtualization: Option<bool>,
191}
192
193impl CapacityGroup {
194 pub fn aws_root_volume_size_gib(&self) -> i32 {
196 self.aws_filesystem_size_gib()
197 .saturating_add(12)
198 .max(20)
199 .min(i32::MAX as u64) as i32
200 }
201
202 pub fn aws_runtime_data_volume_size_gib(&self) -> Option<i32> {
204 (self.nested_virtualization == Some(true))
205 .then(|| self.aws_filesystem_size_gib().max(20).min(i32::MAX as u64) as i32)
206 }
207
208 fn aws_filesystem_size_gib(&self) -> u64 {
209 self.profile
210 .as_ref()
211 .map_or(0, |profile| profile.ephemeral_storage_bytes)
212 .div_ceil(1024 * 1024 * 1024)
213 .saturating_mul(5)
214 .div_ceil(4)
215 }
216}
217
218#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Builder)]
256#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
257#[serde(rename_all = "camelCase", deny_unknown_fields)]
258#[builder(start_fn = new)]
259pub struct ComputeCluster {
260 #[builder(start_fn)]
263 pub id: String,
264
265 #[builder(field)]
268 pub capacity_groups: Vec<CapacityGroup>,
269
270 #[serde(default, skip_serializing_if = "Option::is_none")]
273 pub node_permissions: Option<PermissionProfile>,
274
275 #[serde(default, skip_serializing_if = "Option::is_none")]
277 pub node_permissions_platforms: Option<Vec<Platform>>,
278
279 #[serde(default, skip_serializing_if = "Option::is_none")]
282 pub dynamic_container_pool: Option<String>,
283
284 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
287 #[builder(default)]
288 pub selected_failure_domains: BTreeMap<String, Vec<String>>,
289
290 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
293 #[builder(default)]
294 pub failure_domain_spread: BTreeMap<String, u8>,
295
296 #[serde(skip_serializing_if = "Option::is_none")]
300 pub container_cidr: Option<String>,
301}
302
303impl ComputeCluster {
304 pub const RESOURCE_TYPE: ResourceType = ResourceType::from_static("compute-cluster");
306
307 pub fn id(&self) -> &str {
309 &self.id
310 }
311
312 pub fn container_cidr(&self) -> &str {
314 self.container_cidr.as_deref().unwrap_or("10.244.0.0/16")
315 }
316}
317
318impl<S: compute_cluster_builder::State> ComputeClusterBuilder<S> {
319 pub fn capacity_group(mut self, group: CapacityGroup) -> Self {
321 self.capacity_groups.push(group);
322 self
323 }
324}
325
326#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
328#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
329#[serde(rename_all = "camelCase")]
330pub struct CapacityGroupStatus {
331 pub group_id: String,
333 pub current_machines: u32,
335 pub desired_machines: u32,
337 pub instance_type: String,
339}
340
341#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
343#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
344#[serde(rename_all = "camelCase")]
345pub struct ComputeClusterOutputs {
346 pub cluster_id: String,
348 pub horizon_ready: bool,
350 pub capacity_group_statuses: Vec<CapacityGroupStatus>,
352 pub total_machines: u32,
354}
355
356impl ResourceOutputsDefinition for ComputeClusterOutputs {
357 fn get_resource_type(&self) -> ResourceType {
358 ComputeCluster::RESOURCE_TYPE.clone()
359 }
360
361 fn as_any(&self) -> &dyn Any {
362 self
363 }
364
365 fn box_clone(&self) -> Box<dyn ResourceOutputsDefinition> {
366 Box::new(self.clone())
367 }
368
369 fn outputs_eq(&self, other: &dyn ResourceOutputsDefinition) -> bool {
370 other.as_any().downcast_ref::<ComputeClusterOutputs>() == Some(self)
371 }
372
373 fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
374 serde_json::to_value(self)
375 }
376}
377
378impl ResourceDefinition for ComputeCluster {
379 fn get_resource_type(&self) -> ResourceType {
380 Self::RESOURCE_TYPE
381 }
382
383 fn id(&self) -> &str {
384 &self.id
385 }
386
387 fn replace_after_failed_create_is_safe(&self) -> bool {
388 false
391 }
392
393 fn get_dependencies(&self) -> Vec<ResourceRef> {
394 Vec::new()
400 }
401
402 fn validate_update(&self, new_config: &dyn ResourceDefinition) -> Result<()> {
403 let new_cluster = new_config
404 .as_any()
405 .downcast_ref::<ComputeCluster>()
406 .ok_or_else(|| {
407 AlienError::new(ErrorData::UnexpectedResourceType {
408 resource_id: self.id.clone(),
409 expected: Self::RESOURCE_TYPE,
410 actual: new_config.get_resource_type(),
411 })
412 })?;
413
414 if self.id != new_cluster.id {
415 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
416 resource_id: self.id.clone(),
417 reason: "the 'id' field is immutable".to_string(),
418 }));
419 }
420
421 if self.container_cidr.is_some()
423 && new_cluster.container_cidr.is_some()
424 && self.container_cidr != new_cluster.container_cidr
425 {
426 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
427 resource_id: self.id.clone(),
428 reason: "the 'containerCidr' field is immutable once set".to_string(),
429 }));
430 }
431
432 for new_group in &new_cluster.capacity_groups {
434 if let Some(existing_group) = self
435 .capacity_groups
436 .iter()
437 .find(|g| g.group_id == new_group.group_id)
438 {
439 if let (Some(old), Some(new)) = (
443 existing_group.instance_type.as_deref(),
444 new_group.instance_type.as_deref(),
445 ) {
446 if old != new
447 && ![Platform::Aws, Platform::Gcp, Platform::Azure]
448 .into_iter()
449 .any(|platform| is_same_architecture_machine(platform, old, new))
450 {
451 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
452 resource_id: self.id.clone(),
453 reason: format!(
454 "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",
455 new_group.group_id
456 ),
457 }));
458 }
459 }
460 }
461 }
462
463 Ok(())
464 }
465
466 fn as_any(&self) -> &dyn Any {
467 self
468 }
469
470 fn as_any_mut(&mut self) -> &mut dyn Any {
471 self
472 }
473
474 fn box_clone(&self) -> Box<dyn ResourceDefinition> {
475 Box::new(self.clone())
476 }
477
478 fn resource_eq(&self, other: &dyn ResourceDefinition) -> bool {
479 other.as_any().downcast_ref::<ComputeCluster>() == Some(self)
480 }
481
482 fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
483 serde_json::to_value(self)
484 }
485}
486
487#[cfg(test)]
488mod tests {
489 use super::*;
490
491 #[test]
492 fn aws_disk_sizing_preserves_rounding_minimums_and_saturation() {
493 let gib = 1024 * 1024 * 1024;
494 for (bytes, root, data) in [
495 (0, 20, 20),
496 (gib, 20, 20),
497 (20 * gib, 37, 25),
498 (20 * gib + 1, 39, 27),
499 (u64::MAX, i32::MAX, i32::MAX),
500 ] {
501 let mut group = CapacityGroup {
502 group_id: "general".into(),
503 instance_type: None,
504 profile: Some(MachineProfile {
505 cpu: "4.0".into(),
506 memory_bytes: 8 * gib,
507 ephemeral_storage_bytes: bytes,
508 architecture: None,
509 gpu: None,
510 data_disk: None,
511 }),
512 min_size: 1,
513 max_size: 1,
514 scale_policy: None,
515 nested_virtualization: None,
516 };
517 assert_eq!(group.aws_root_volume_size_gib(), root);
518 assert_eq!(group.aws_runtime_data_volume_size_gib(), None);
519 group.nested_virtualization = Some(true);
520 assert_eq!(group.aws_root_volume_size_gib(), root);
521 assert_eq!(group.aws_runtime_data_volume_size_gib(), Some(data));
522 }
523 }
524
525 #[test]
526 fn test_compute_cluster_creation() {
527 let cluster = ComputeCluster::new("compute".to_string())
528 .capacity_group(CapacityGroup {
529 group_id: "general".to_string(),
530 instance_type: Some("m7g.xlarge".to_string()),
531 profile: None,
532 min_size: 1,
533 max_size: 5,
534 scale_policy: None,
535 nested_virtualization: None,
536 })
537 .build();
538
539 assert_eq!(cluster.id(), "compute");
540 assert_eq!(cluster.capacity_groups.len(), 1);
541 assert_eq!(cluster.capacity_groups[0].group_id, "general");
542 assert_eq!(cluster.container_cidr(), "10.244.0.0/16");
543 }
544
545 #[test]
546 fn test_compute_cluster_multiple_capacity_groups() {
547 let cluster = ComputeCluster::new("multi-pool".to_string())
548 .capacity_group(CapacityGroup {
549 group_id: "general".to_string(),
550 instance_type: Some("m7g.xlarge".to_string()),
551 profile: None,
552 min_size: 1,
553 max_size: 3,
554 scale_policy: None,
555 nested_virtualization: None,
556 })
557 .capacity_group(CapacityGroup {
558 group_id: "gpu".to_string(),
559 instance_type: Some("g5.xlarge".to_string()),
560 profile: Some(MachineProfile {
561 cpu: "4.0".to_string(),
562 memory_bytes: 17179869184, ephemeral_storage_bytes: 214748364800, architecture: None,
565 gpu: Some(GpuSpec {
566 gpu_type: "nvidia-a10g".to_string(),
567 count: 1,
568 }),
569 data_disk: None,
570 }),
571 min_size: 0,
572 max_size: 2,
573 scale_policy: None,
574 nested_virtualization: None,
575 })
576 .build();
577
578 assert_eq!(cluster.capacity_groups.len(), 2);
579 assert_eq!(cluster.capacity_groups[0].group_id, "general");
580 assert_eq!(cluster.capacity_groups[1].group_id, "gpu");
581 assert!(cluster.capacity_groups[1]
582 .profile
583 .as_ref()
584 .unwrap()
585 .gpu
586 .is_some());
587 }
588
589 #[test]
590 fn test_compute_cluster_custom_cidr() {
591 let cluster = ComputeCluster::new("custom-net".to_string())
592 .container_cidr("172.30.0.0/16".to_string())
593 .capacity_group(CapacityGroup {
594 group_id: "general".to_string(),
595 instance_type: None,
596 profile: None,
597 min_size: 1,
598 max_size: 5,
599 scale_policy: None,
600 nested_virtualization: None,
601 })
602 .build();
603
604 assert_eq!(cluster.container_cidr(), "172.30.0.0/16");
605 }
606
607 #[test]
608 fn test_compute_cluster_validate_update_immutable_id() {
609 let cluster1 = ComputeCluster::new("cluster-1".to_string())
610 .capacity_group(CapacityGroup {
611 group_id: "general".to_string(),
612 instance_type: None,
613 profile: None,
614 min_size: 1,
615 max_size: 5,
616 scale_policy: None,
617 nested_virtualization: None,
618 })
619 .build();
620
621 let cluster2 = ComputeCluster::new("cluster-2".to_string())
622 .capacity_group(CapacityGroup {
623 group_id: "general".to_string(),
624 instance_type: None,
625 profile: None,
626 min_size: 1,
627 max_size: 5,
628 scale_policy: None,
629 nested_virtualization: None,
630 })
631 .build();
632
633 let result = cluster1.validate_update(&cluster2);
634 assert!(result.is_err());
635 }
636
637 #[test]
638 fn test_compute_cluster_validate_update_scale_change() {
639 let cluster1 = ComputeCluster::new("compute".to_string())
640 .capacity_group(CapacityGroup {
641 group_id: "general".to_string(),
642 instance_type: Some("m7g.xlarge".to_string()),
643 profile: None,
644 min_size: 1,
645 max_size: 5,
646 scale_policy: None,
647 nested_virtualization: None,
648 })
649 .build();
650
651 let cluster2 = ComputeCluster::new("compute".to_string())
652 .capacity_group(CapacityGroup {
653 group_id: "general".to_string(),
654 instance_type: Some("m7g.xlarge".to_string()),
655 profile: None,
656 min_size: 2,
657 max_size: 10,
658 scale_policy: None,
659 nested_virtualization: None,
660 })
661 .build();
662
663 let result = cluster1.validate_update(&cluster2);
665 assert!(result.is_ok());
666 }
667
668 #[test]
669 fn machine_changes_stay_within_one_cloud_and_architecture() {
670 let cluster = |machine: &str| {
671 ComputeCluster::new("compute".to_string())
672 .capacity_group(CapacityGroup {
673 group_id: "general".to_string(),
674 instance_type: Some(machine.to_string()),
675 profile: None,
676 min_size: 1,
677 max_size: 5,
678 scale_policy: None,
679 nested_virtualization: None,
680 })
681 .build()
682 };
683 for (old, new) in [
684 ("t4g.small", "t4g.medium"),
685 ("e2-medium", "n2-standard-4"),
686 ("Standard_D2s_v5", "Standard_D4s_v5"),
687 ] {
688 cluster(old)
689 .validate_update(&cluster(new))
690 .expect("same cloud and architecture");
691 }
692 for (old, new) in [
693 ("t4g.small", "m7i.large"),
694 ("t4g.small", "t4g.unknown"),
695 ("n2-standard-2", "m7i.large"),
696 ("Standard_D2s_v5", "n2-standard-4"),
697 ] {
698 let error = cluster(old)
699 .validate_update(&cluster(new))
700 .expect_err("needs setup");
701 assert_eq!(error.code, "INVALID_RESOURCE_UPDATE", "{old} -> {new}");
702 }
703 }
704
705 #[test]
706 fn node_permissions_are_inline_optional_and_roundtrip_exactly() {
707 let mut cluster = ComputeCluster::new("compute".to_string()).build();
708 let absent = serde_json::to_value(&cluster).unwrap();
709 assert!(absent.get("nodePermissions").is_none());
710 assert_eq!(
711 serde_json::from_value::<ComputeCluster>(absent).unwrap(),
712 cluster
713 );
714 cluster.node_permissions =
715 Some(PermissionProfile::new().resource("objects", ["storage/data-read"]));
716 let json = serde_json::to_value(&cluster).unwrap();
717 assert_eq!(
718 json["nodePermissions"],
719 serde_json::json!({"objects": ["storage/data-read"]})
720 );
721 assert_eq!(
722 serde_json::from_value::<ComputeCluster>(json).unwrap(),
723 cluster
724 );
725 assert_eq!(ResourceDefinition::get_permissions(&cluster), None);
726 }
727
728 #[test]
729 fn test_compute_cluster_serialization() {
730 let cluster = ComputeCluster::new("test-cluster".to_string())
731 .capacity_group(CapacityGroup {
732 group_id: "general".to_string(),
733 instance_type: Some("m7g.xlarge".to_string()),
734 profile: None,
735 min_size: 1,
736 max_size: 5,
737 scale_policy: None,
738 nested_virtualization: None,
739 })
740 .build();
741
742 let json = serde_json::to_string(&cluster).unwrap();
743 let deserialized: ComputeCluster = serde_json::from_str(&json).unwrap();
744 assert_eq!(cluster, deserialized);
745 }
746}