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
178impl CapacityGroup {
179 pub fn aws_root_volume_size_gib(&self) -> i32 {
181 self.aws_filesystem_size_gib()
182 .saturating_add(12)
183 .max(20)
184 .min(i32::MAX as u64) as i32
185 }
186
187 pub fn aws_runtime_data_volume_size_gib(&self) -> Option<i32> {
189 (self.nested_virtualization == Some(true))
190 .then(|| self.aws_filesystem_size_gib().max(20).min(i32::MAX as u64) as i32)
191 }
192
193 fn aws_filesystem_size_gib(&self) -> u64 {
194 self.profile
195 .as_ref()
196 .map_or(0, |profile| profile.ephemeral_storage_bytes)
197 .div_ceil(1024 * 1024 * 1024)
198 .saturating_mul(5)
199 .div_ceil(4)
200 }
201}
202
203#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Builder)]
240#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
241#[serde(rename_all = "camelCase", deny_unknown_fields)]
242#[builder(start_fn = new)]
243pub struct ComputeCluster {
244 #[builder(start_fn)]
247 pub id: String,
248
249 #[builder(field)]
252 pub capacity_groups: Vec<CapacityGroup>,
253
254 #[serde(default, skip_serializing_if = "Option::is_none")]
257 pub node_permissions: Option<PermissionProfile>,
258
259 #[serde(default, skip_serializing_if = "Option::is_none")]
261 pub node_permissions_platforms: Option<Vec<Platform>>,
262
263 #[serde(default, skip_serializing_if = "Option::is_none")]
266 pub dynamic_container_pool: Option<String>,
267
268 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
271 #[builder(default)]
272 pub selected_failure_domains: BTreeMap<String, Vec<String>>,
273
274 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
277 #[builder(default)]
278 pub failure_domain_spread: BTreeMap<String, u8>,
279
280 #[serde(skip_serializing_if = "Option::is_none")]
284 pub container_cidr: Option<String>,
285}
286
287impl ComputeCluster {
288 pub const RESOURCE_TYPE: ResourceType = ResourceType::from_static("compute-cluster");
290
291 pub fn id(&self) -> &str {
293 &self.id
294 }
295
296 pub fn container_cidr(&self) -> &str {
298 self.container_cidr.as_deref().unwrap_or("10.244.0.0/16")
299 }
300}
301
302impl<S: compute_cluster_builder::State> ComputeClusterBuilder<S> {
303 pub fn capacity_group(mut self, group: CapacityGroup) -> Self {
305 self.capacity_groups.push(group);
306 self
307 }
308}
309
310#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
312#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
313#[serde(rename_all = "camelCase")]
314pub struct CapacityGroupStatus {
315 pub group_id: String,
317 pub current_machines: u32,
319 pub desired_machines: u32,
321 pub instance_type: String,
323}
324
325#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
327#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
328#[serde(rename_all = "camelCase")]
329pub struct ComputeClusterOutputs {
330 pub cluster_id: String,
332 pub horizon_ready: bool,
334 pub capacity_group_statuses: Vec<CapacityGroupStatus>,
336 pub total_machines: u32,
338}
339
340impl ResourceOutputsDefinition for ComputeClusterOutputs {
341 fn get_resource_type(&self) -> ResourceType {
342 ComputeCluster::RESOURCE_TYPE.clone()
343 }
344
345 fn as_any(&self) -> &dyn Any {
346 self
347 }
348
349 fn box_clone(&self) -> Box<dyn ResourceOutputsDefinition> {
350 Box::new(self.clone())
351 }
352
353 fn outputs_eq(&self, other: &dyn ResourceOutputsDefinition) -> bool {
354 other.as_any().downcast_ref::<ComputeClusterOutputs>() == Some(self)
355 }
356
357 fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
358 serde_json::to_value(self)
359 }
360}
361
362impl ResourceDefinition for ComputeCluster {
363 fn get_resource_type(&self) -> ResourceType {
364 Self::RESOURCE_TYPE
365 }
366
367 fn id(&self) -> &str {
368 &self.id
369 }
370
371 fn replace_after_failed_create_is_safe(&self) -> bool {
372 false
375 }
376
377 fn get_dependencies(&self) -> Vec<ResourceRef> {
378 Vec::new()
384 }
385
386 fn validate_update(&self, new_config: &dyn ResourceDefinition) -> Result<()> {
387 let new_cluster = new_config
388 .as_any()
389 .downcast_ref::<ComputeCluster>()
390 .ok_or_else(|| {
391 AlienError::new(ErrorData::UnexpectedResourceType {
392 resource_id: self.id.clone(),
393 expected: Self::RESOURCE_TYPE,
394 actual: new_config.get_resource_type(),
395 })
396 })?;
397
398 if self.id != new_cluster.id {
399 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
400 resource_id: self.id.clone(),
401 reason: "the 'id' field is immutable".to_string(),
402 }));
403 }
404
405 if self.container_cidr.is_some()
407 && new_cluster.container_cidr.is_some()
408 && self.container_cidr != new_cluster.container_cidr
409 {
410 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
411 resource_id: self.id.clone(),
412 reason: "the 'containerCidr' field is immutable once set".to_string(),
413 }));
414 }
415
416 for new_group in &new_cluster.capacity_groups {
418 if let Some(existing_group) = self
419 .capacity_groups
420 .iter()
421 .find(|g| g.group_id == new_group.group_id)
422 {
423 if let (Some(old), Some(new)) = (
427 existing_group.instance_type.as_deref(),
428 new_group.instance_type.as_deref(),
429 ) {
430 if old != new
431 && ![Platform::Aws, Platform::Gcp, Platform::Azure]
432 .into_iter()
433 .any(|platform| is_same_architecture_machine(platform, old, new))
434 {
435 return Err(AlienError::new(ErrorData::InvalidResourceUpdate {
436 resource_id: self.id.clone(),
437 reason: format!(
438 "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",
439 new_group.group_id
440 ),
441 }));
442 }
443 }
444 }
445 }
446
447 Ok(())
448 }
449
450 fn as_any(&self) -> &dyn Any {
451 self
452 }
453
454 fn as_any_mut(&mut self) -> &mut dyn Any {
455 self
456 }
457
458 fn box_clone(&self) -> Box<dyn ResourceDefinition> {
459 Box::new(self.clone())
460 }
461
462 fn resource_eq(&self, other: &dyn ResourceDefinition) -> bool {
463 other.as_any().downcast_ref::<ComputeCluster>() == Some(self)
464 }
465
466 fn to_json_value(&self) -> serde_json::Result<serde_json::Value> {
467 serde_json::to_value(self)
468 }
469}
470
471#[cfg(test)]
472mod tests {
473 use super::*;
474
475 #[test]
476 fn aws_disk_sizing_preserves_rounding_minimums_and_saturation() {
477 let gib = 1024 * 1024 * 1024;
478 for (bytes, root, data) in [
479 (0, 20, 20),
480 (gib, 20, 20),
481 (20 * gib, 37, 25),
482 (20 * gib + 1, 39, 27),
483 (u64::MAX, i32::MAX, i32::MAX),
484 ] {
485 let mut group = CapacityGroup {
486 group_id: "general".into(),
487 instance_type: None,
488 profile: Some(MachineProfile {
489 cpu: "4.0".into(),
490 memory_bytes: 8 * gib,
491 ephemeral_storage_bytes: bytes,
492 architecture: None,
493 gpu: None,
494 }),
495 min_size: 1,
496 max_size: 1,
497 scale_policy: None,
498 nested_virtualization: None,
499 };
500 assert_eq!(group.aws_root_volume_size_gib(), root);
501 assert_eq!(group.aws_runtime_data_volume_size_gib(), None);
502 group.nested_virtualization = Some(true);
503 assert_eq!(group.aws_root_volume_size_gib(), root);
504 assert_eq!(group.aws_runtime_data_volume_size_gib(), Some(data));
505 }
506 }
507
508 #[test]
509 fn test_compute_cluster_creation() {
510 let cluster = ComputeCluster::new("compute".to_string())
511 .capacity_group(CapacityGroup {
512 group_id: "general".to_string(),
513 instance_type: Some("m7g.xlarge".to_string()),
514 profile: None,
515 min_size: 1,
516 max_size: 5,
517 scale_policy: None,
518 nested_virtualization: None,
519 })
520 .build();
521
522 assert_eq!(cluster.id(), "compute");
523 assert_eq!(cluster.capacity_groups.len(), 1);
524 assert_eq!(cluster.capacity_groups[0].group_id, "general");
525 assert_eq!(cluster.container_cidr(), "10.244.0.0/16");
526 }
527
528 #[test]
529 fn test_compute_cluster_multiple_capacity_groups() {
530 let cluster = ComputeCluster::new("multi-pool".to_string())
531 .capacity_group(CapacityGroup {
532 group_id: "general".to_string(),
533 instance_type: Some("m7g.xlarge".to_string()),
534 profile: None,
535 min_size: 1,
536 max_size: 3,
537 scale_policy: None,
538 nested_virtualization: None,
539 })
540 .capacity_group(CapacityGroup {
541 group_id: "gpu".to_string(),
542 instance_type: Some("g5.xlarge".to_string()),
543 profile: Some(MachineProfile {
544 cpu: "4.0".to_string(),
545 memory_bytes: 17179869184, ephemeral_storage_bytes: 214748364800, architecture: None,
548 gpu: Some(GpuSpec {
549 gpu_type: "nvidia-a10g".to_string(),
550 count: 1,
551 }),
552 }),
553 min_size: 0,
554 max_size: 2,
555 scale_policy: None,
556 nested_virtualization: None,
557 })
558 .build();
559
560 assert_eq!(cluster.capacity_groups.len(), 2);
561 assert_eq!(cluster.capacity_groups[0].group_id, "general");
562 assert_eq!(cluster.capacity_groups[1].group_id, "gpu");
563 assert!(cluster.capacity_groups[1]
564 .profile
565 .as_ref()
566 .unwrap()
567 .gpu
568 .is_some());
569 }
570
571 #[test]
572 fn test_compute_cluster_custom_cidr() {
573 let cluster = ComputeCluster::new("custom-net".to_string())
574 .container_cidr("172.30.0.0/16".to_string())
575 .capacity_group(CapacityGroup {
576 group_id: "general".to_string(),
577 instance_type: None,
578 profile: None,
579 min_size: 1,
580 max_size: 5,
581 scale_policy: None,
582 nested_virtualization: None,
583 })
584 .build();
585
586 assert_eq!(cluster.container_cidr(), "172.30.0.0/16");
587 }
588
589 #[test]
590 fn test_compute_cluster_validate_update_immutable_id() {
591 let cluster1 = ComputeCluster::new("cluster-1".to_string())
592 .capacity_group(CapacityGroup {
593 group_id: "general".to_string(),
594 instance_type: None,
595 profile: None,
596 min_size: 1,
597 max_size: 5,
598 scale_policy: None,
599 nested_virtualization: None,
600 })
601 .build();
602
603 let cluster2 = ComputeCluster::new("cluster-2".to_string())
604 .capacity_group(CapacityGroup {
605 group_id: "general".to_string(),
606 instance_type: None,
607 profile: None,
608 min_size: 1,
609 max_size: 5,
610 scale_policy: None,
611 nested_virtualization: None,
612 })
613 .build();
614
615 let result = cluster1.validate_update(&cluster2);
616 assert!(result.is_err());
617 }
618
619 #[test]
620 fn test_compute_cluster_validate_update_scale_change() {
621 let cluster1 = ComputeCluster::new("compute".to_string())
622 .capacity_group(CapacityGroup {
623 group_id: "general".to_string(),
624 instance_type: Some("m7g.xlarge".to_string()),
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 cluster2 = ComputeCluster::new("compute".to_string())
634 .capacity_group(CapacityGroup {
635 group_id: "general".to_string(),
636 instance_type: Some("m7g.xlarge".to_string()),
637 profile: None,
638 min_size: 2,
639 max_size: 10,
640 scale_policy: None,
641 nested_virtualization: None,
642 })
643 .build();
644
645 let result = cluster1.validate_update(&cluster2);
647 assert!(result.is_ok());
648 }
649
650 #[test]
651 fn machine_changes_stay_within_one_cloud_and_architecture() {
652 let cluster = |machine: &str| {
653 ComputeCluster::new("compute".to_string())
654 .capacity_group(CapacityGroup {
655 group_id: "general".to_string(),
656 instance_type: Some(machine.to_string()),
657 profile: None,
658 min_size: 1,
659 max_size: 5,
660 scale_policy: None,
661 nested_virtualization: None,
662 })
663 .build()
664 };
665 for (old, new) in [
666 ("t4g.small", "t4g.medium"),
667 ("e2-medium", "n2-standard-4"),
668 ("Standard_D2s_v5", "Standard_D4s_v5"),
669 ] {
670 cluster(old)
671 .validate_update(&cluster(new))
672 .expect("same cloud and architecture");
673 }
674 for (old, new) in [
675 ("t4g.small", "m7i.large"),
676 ("t4g.small", "t4g.unknown"),
677 ("n2-standard-2", "m7i.large"),
678 ("Standard_D2s_v5", "n2-standard-4"),
679 ] {
680 let error = cluster(old)
681 .validate_update(&cluster(new))
682 .expect_err("needs setup");
683 assert_eq!(error.code, "INVALID_RESOURCE_UPDATE", "{old} -> {new}");
684 }
685 }
686
687 #[test]
688 fn node_permissions_are_inline_optional_and_roundtrip_exactly() {
689 let mut cluster = ComputeCluster::new("compute".to_string()).build();
690 let absent = serde_json::to_value(&cluster).unwrap();
691 assert!(absent.get("nodePermissions").is_none());
692 assert_eq!(
693 serde_json::from_value::<ComputeCluster>(absent).unwrap(),
694 cluster
695 );
696 cluster.node_permissions =
697 Some(PermissionProfile::new().resource("objects", ["storage/data-read"]));
698 let json = serde_json::to_value(&cluster).unwrap();
699 assert_eq!(
700 json["nodePermissions"],
701 serde_json::json!({"objects": ["storage/data-read"]})
702 );
703 assert_eq!(
704 serde_json::from_value::<ComputeCluster>(json).unwrap(),
705 cluster
706 );
707 assert_eq!(ResourceDefinition::get_permissions(&cluster), None);
708 }
709
710 #[test]
711 fn test_compute_cluster_serialization() {
712 let cluster = ComputeCluster::new("test-cluster".to_string())
713 .capacity_group(CapacityGroup {
714 group_id: "general".to_string(),
715 instance_type: Some("m7g.xlarge".to_string()),
716 profile: None,
717 min_size: 1,
718 max_size: 5,
719 scale_policy: None,
720 nested_virtualization: None,
721 })
722 .build();
723
724 let json = serde_json::to_string(&cluster).unwrap();
725 let deserialized: ComputeCluster = serde_json::from_str(&json).unwrap();
726 assert_eq!(cluster, deserialized);
727 }
728}