1use crate::{
8 instance_catalog::{self, Architecture, WorkloadRequirements},
9 CapacityGroup, CapacityGroupScalePolicy, ComputeChoiceRange, ComputePoolSelection, Container,
10 Daemon, ErrorData, FailureDomainSelection, GpuSpec, MachineProfile, Platform, ResourceSpec,
11 Stack,
12};
13use alien_error::{AlienError, Result};
14use serde::{Deserialize, Serialize};
15use std::collections::HashMap;
16
17#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
19#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
20#[serde(rename_all = "camelCase")]
21pub struct ComputePlan {
22 pub pools: Vec<ComputePoolPlan>,
24}
25
26#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
28#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
29#[serde(rename_all = "camelCase")]
30pub struct ComputePoolPlan {
31 pub pool_id: String,
33 pub workloads: Vec<String>,
35 pub requirements: MachineProfile,
37 pub scale: CapacityGroupScalePolicy,
39 pub selected: ComputePoolSelection,
41 pub recommended: ComputePoolSelection,
43 pub machines: Vec<ComputeMachineOption>,
45 #[serde(default, skip_serializing_if = "Vec::is_empty")]
47 pub errors: Vec<String>,
48}
49
50#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
52#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
53#[serde(rename_all = "camelCase")]
54pub struct ComputeMachineOption {
55 pub machine: String,
57 pub profile: MachineProfile,
59 pub recommended: bool,
61}
62
63pub fn default_persistent_failure_domains() -> FailureDomainSelection {
67 FailureDomainSelection {
68 spread: 1,
69 selected_failure_domains: Vec::new(),
70 }
71}
72
73pub fn plan_compute(
75 stack: &Stack,
76 platform: Platform,
77 selected_settings: Option<&crate::ComputeSettings>,
78) -> Result<ComputePlan, ErrorData> {
79 let mut groups = collect_workload_groups(stack)?;
80 merge_explicit_compute_groups(stack, &mut groups)?;
81
82 let mut pool_ids: Vec<String> = groups.keys().cloned().collect();
83 pool_ids.sort();
84
85 let mut pools = Vec::new();
86 for pool_id in pool_ids {
87 let group = groups.remove(&pool_id).expect("pool id came from map keys");
88 let requirements = group.requirements;
89 let selected = selected_settings.and_then(|settings| settings.pools.get(&pool_id));
90 let recommended = recommended_selection(
91 platform,
92 &requirements,
93 &group.scale,
94 group.requires_failure_domain,
95 )?;
96 let mut selected_choice = selected.cloned().unwrap_or_else(|| recommended.clone());
97 if selected_choice.failure_domains().is_none() {
98 if let Some(default_failure_domains) = recommended.failure_domains().cloned() {
99 match &mut selected_choice {
100 ComputePoolSelection::Fixed {
101 failure_domains, ..
102 }
103 | ComputePoolSelection::Autoscale {
104 failure_domains, ..
105 } => *failure_domains = Some(default_failure_domains),
106 }
107 }
108 }
109 let errors = validate_compute_pool_selection(
110 platform,
111 &pool_id,
112 &selected_choice,
113 &requirements,
114 &group.scale,
115 );
116 let machines = machine_options(platform, &requirements, selected_choice.machine())?;
117
118 pools.push(ComputePoolPlan {
119 pool_id,
120 workloads: group.workloads,
121 requirements: requirements_to_profile(&requirements),
122 scale: group.scale,
123 selected: selected_choice,
124 recommended,
125 machines,
126 errors,
127 });
128 }
129
130 Ok(ComputePlan { pools })
131}
132
133#[derive(Debug, Clone)]
134struct PlannedGroup {
135 workloads: Vec<String>,
136 requirements: WorkloadRequirements,
137 scale: CapacityGroupScalePolicy,
138 requires_failure_domain: bool,
139}
140
141fn collect_workload_groups(stack: &Stack) -> Result<HashMap<String, PlannedGroup>, ErrorData> {
142 let mut groups: HashMap<String, Vec<Workload>> = HashMap::new();
143
144 for entry in stack.resources.values() {
145 if let Some(container) = entry.config.downcast_ref::<Container>() {
146 groups
147 .entry(
148 container
149 .pool
150 .clone()
151 .unwrap_or_else(|| needed_container_pool(container).to_string()),
152 )
153 .or_default()
154 .push(Workload::from_container(container)?);
155 }
156 if let Some(daemon) = entry.config.downcast_ref::<Daemon>() {
157 if daemon.cluster.is_some() {
158 groups
159 .entry(daemon.pool.clone().unwrap_or_else(|| "general".to_string()))
160 .or_default()
161 .push(Workload::from_daemon(daemon)?);
162 }
163 }
164 }
165
166 let mut planned = HashMap::new();
167 for (pool_id, workloads) in groups {
168 let requirements = aggregate_workloads(&workloads);
169 let requires_failure_domain = workloads
170 .iter()
171 .any(|workload| workload.requires_failure_domain);
172 let min_size = default_min_machines(&requirements);
173 let max_size = default_max_machines(&requirements);
174 planned.insert(
175 pool_id,
176 PlannedGroup {
177 workloads: workloads.into_iter().map(|w| w.id).collect(),
178 scale: CapacityGroupScalePolicy::from_selected_bounds(min_size, max_size),
179 requirements,
180 requires_failure_domain,
181 },
182 );
183 }
184 Ok(planned)
185}
186
187fn merge_explicit_compute_groups(
188 stack: &Stack,
189 groups: &mut HashMap<String, PlannedGroup>,
190) -> Result<(), ErrorData> {
191 for entry in stack.resources.values() {
192 let Some(cluster) = entry.config.downcast_ref::<crate::ComputeCluster>() else {
193 continue;
194 };
195 for group in &cluster.capacity_groups {
196 let explicit_requirements = profile_to_requirements(
197 group.profile.as_ref(),
198 group.nested_virtualization.unwrap_or(false),
199 );
200 let scale = group.scale_policy.clone().unwrap_or_else(|| {
201 CapacityGroupScalePolicy::from_selected_bounds(group.min_size, group.max_size)
202 });
203 groups
204 .entry(group.group_id.clone())
205 .and_modify(|planned| {
206 if group.instance_type.is_none() {
207 merge_requirements(&mut planned.requirements, &explicit_requirements);
208 } else {
209 planned.requirements.architecture = planned
216 .requirements
217 .architecture
218 .or(explicit_requirements.architecture);
219 planned.requirements.nested_virt |= explicit_requirements.nested_virt;
220 }
221 planned.scale = merge_scale_policy(&planned.scale, &scale);
222 })
223 .or_insert_with(|| PlannedGroup {
224 workloads: Vec::new(),
225 scale,
226 requirements: explicit_requirements,
227 requires_failure_domain: false,
228 });
229 }
230 }
231 Ok(())
232}
233
234fn recommended_selection(
235 platform: Platform,
236 requirements: &WorkloadRequirements,
237 scale: &CapacityGroupScalePolicy,
238 requires_failure_domain: bool,
239) -> Result<ComputePoolSelection, ErrorData> {
240 let machine = match platform {
241 Platform::Aws | Platform::Gcp | Platform::Azure => Some(
242 instance_catalog::select_instance_type(platform, requirements)
243 .map_err(|message| {
244 AlienError::new(ErrorData::GenericError {
245 message: format!("Failed to select {platform} machine: {message}"),
246 })
247 })?
248 .instance_type
249 .to_string(),
250 ),
251 Platform::Local | Platform::Kubernetes | Platform::Machines | Platform::Test => None,
252 };
253
254 let failure_domains = (requires_failure_domain
255 && matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure))
256 .then(default_persistent_failure_domains);
257
258 match scale {
259 CapacityGroupScalePolicy::Fixed { machines } => Ok(ComputePoolSelection::Fixed {
260 machines: machines.default.max(1),
261 machine,
262 failure_domains,
263 }),
264 CapacityGroupScalePolicy::Autoscale { min, max } => Ok(ComputePoolSelection::Autoscale {
265 min: min.default,
266 max: max.default.max(min.default),
267 machine,
268 failure_domains,
269 }),
270 }
271}
272
273pub fn validate_compute_pool_selection(
276 platform: Platform,
277 pool_id: &str,
278 selection: &ComputePoolSelection,
279 requirements: &WorkloadRequirements,
280 scale: &CapacityGroupScalePolicy,
281) -> Vec<String> {
282 let mut errors = Vec::new();
283 if let Err(message) = selection.validate() {
284 errors.push(message);
285 }
286 if let Err(message) = validate_selection_against_scale(selection, scale) {
287 errors.push(format!("Pool '{pool_id}' {message}"));
288 }
289 if matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure) {
290 match selection.machine() {
291 Some(machine) => match instance_catalog::find_instance_type(platform, machine) {
292 Some(spec) => {
293 let architecture = requirements.architecture.unwrap_or(spec.architecture);
294 if !instance_satisfies(spec, requirements, architecture) {
295 errors.push(format!(
296 "{} machine '{}' does not satisfy pool '{}' requirements",
297 platform, machine, pool_id
298 ));
299 }
300 }
301 None => errors.push(format!(
302 "Unknown {} machine '{}' for pool '{}'",
303 platform, machine, pool_id
304 )),
305 },
306 None => errors.push(format!(
307 "Pool '{}' requires a provider machine on {}",
308 pool_id, platform
309 )),
310 }
311 }
312 errors
313}
314
315pub fn capacity_group_requirements(group: &CapacityGroup) -> WorkloadRequirements {
317 profile_to_requirements(
318 group.profile.as_ref(),
319 group.nested_virtualization.unwrap_or(false),
320 )
321}
322
323fn machine_options(
324 platform: Platform,
325 requirements: &WorkloadRequirements,
326 selected_machine: Option<&str>,
327) -> Result<Vec<ComputeMachineOption>, ErrorData> {
328 if !matches!(platform, Platform::Aws | Platform::Gcp | Platform::Azure) {
329 return Ok(Vec::new());
330 }
331 let recommended =
332 instance_catalog::select_instance_type(platform, requirements).map_err(|message| {
333 AlienError::new(ErrorData::GenericError {
334 message: format!("Failed to select {platform} machine: {message}"),
335 })
336 })?;
337 let resolved_architecture = requirements
338 .architecture
339 .or_else(|| {
340 selected_machine.and_then(|machine| {
341 instance_catalog::find_instance_type(platform, machine)
342 .map(|spec| spec.architecture)
343 })
344 })
345 .or(recommended.profile.architecture)
346 .ok_or_else(|| {
347 AlienError::new(ErrorData::GenericError {
348 message: format!("Selected {platform} machine has no CPU architecture"),
349 })
350 })?;
351 let recommended = recommended.instance_type.to_string();
352
353 let mut options: Vec<ComputeMachineOption> = instance_catalog::catalog_for_platform(platform)
354 .into_iter()
355 .filter(|spec| instance_satisfies(spec, requirements, resolved_architecture))
356 .map(|spec| ComputeMachineOption {
357 machine: spec.name.to_string(),
358 profile: spec.to_machine_profile_for_storage(requirements.max_ephemeral_storage_bytes),
359 recommended: spec.name == recommended || Some(spec.name) == selected_machine,
360 })
361 .collect();
362 options.sort_by(|a, b| a.machine.cmp(&b.machine));
363 Ok(options)
364}
365
366fn instance_satisfies(
367 spec: &instance_catalog::InstanceTypeSpec,
368 requirements: &WorkloadRequirements,
369 resolved_architecture: Architecture,
370) -> bool {
371 if spec.architecture != resolved_architecture {
372 return false;
373 }
374 if requirements.nested_virt && !spec.is_nested_virt_capable() {
375 return false;
376 }
377 if spec.vcpu < requirements.max_cpu_per_container.ceil() as u32 {
378 return false;
379 }
380 if spec.memory_bytes < requirements.max_memory_per_container {
381 return false;
382 }
383 if spec.has_configurable_ephemeral_storage() {
384 if crate::instance_catalog::max_configurable_ephemeral_storage_bytes(spec.platform)
385 .is_none_or(|max| requirements.max_ephemeral_storage_bytes > max)
386 {
387 return false;
388 }
389 } else if spec.ephemeral_storage_bytes < requirements.max_ephemeral_storage_bytes {
390 return false;
391 }
392 match (&requirements.gpu, spec.gpu) {
393 (Some(required), Some(actual)) => {
394 (required.gpu_type == "any" || required.gpu_type == actual.gpu_type)
395 && actual.count >= required.count
396 }
397 (Some(_), None) => false,
398 (None, _) => true,
399 }
400}
401
402#[derive(Debug, Clone)]
403struct Workload {
404 id: String,
405 cpu: f64,
406 memory_bytes: u64,
407 desired_replicas: f64,
408 max_replicas: f64,
409 ephemeral_storage_bytes: u64,
410 gpu: Option<GpuSpec>,
411 requires_failure_domain: bool,
412}
413
414impl Workload {
415 fn from_container(container: &Container) -> Result<Self, ErrorData> {
416 let cpu = parse_cpu(&container.id, &container.cpu)?;
417 let memory_bytes = parse_memory(&container.id, &container.memory)?;
418 let desired_replicas = container
419 .autoscaling
420 .as_ref()
421 .map(|a| a.desired)
422 .or(container.replicas)
423 .unwrap_or(1) as f64;
424 let max_replicas = container
425 .autoscaling
426 .as_ref()
427 .map(|a| a.max)
428 .or(container.replicas)
429 .unwrap_or(1) as f64;
430 let ephemeral_storage_bytes = container
431 .ephemeral_storage
432 .as_deref()
433 .map(instance_catalog::parse_memory_bytes)
434 .transpose()
435 .map_err(|message| {
436 AlienError::new(ErrorData::GenericError {
437 message: format!(
438 "Failed to parse ephemeral storage for '{}': {message}",
439 container.id
440 ),
441 })
442 })?
443 .unwrap_or(0);
444
445 Ok(Self {
446 id: container.id.clone(),
447 cpu,
448 memory_bytes,
449 desired_replicas,
450 max_replicas,
451 ephemeral_storage_bytes,
452 gpu: container.gpu.as_ref().map(|gpu| GpuSpec {
453 gpu_type: gpu.gpu_type.clone(),
454 count: gpu.count,
455 }),
456 requires_failure_domain: container.stateful && container.persistent_storage.is_some(),
457 })
458 }
459
460 fn from_daemon(daemon: &Daemon) -> Result<Self, ErrorData> {
461 Ok(Self {
462 id: daemon.id.clone(),
463 cpu: parse_cpu(&daemon.id, &daemon.cpu)?,
464 memory_bytes: parse_memory(&daemon.id, &daemon.memory)?,
465 desired_replicas: 1.0,
466 max_replicas: 1.0,
467 ephemeral_storage_bytes: 0,
468 gpu: None,
469 requires_failure_domain: false,
470 })
471 }
472}
473
474fn parse_cpu(resource_id: &str, spec: &ResourceSpec) -> Result<f64, ErrorData> {
475 instance_catalog::parse_cpu(&spec.desired).map_err(|message| {
476 AlienError::new(ErrorData::GenericError {
477 message: format!(
478 "Failed to parse CPU requirement '{}' for '{}': {message}",
479 spec.desired, resource_id
480 ),
481 })
482 })
483}
484
485fn parse_memory(resource_id: &str, spec: &ResourceSpec) -> Result<u64, ErrorData> {
486 instance_catalog::parse_memory_bytes(&spec.desired).map_err(|message| {
487 AlienError::new(ErrorData::GenericError {
488 message: format!(
489 "Failed to parse memory requirement '{}' for '{}': {message}",
490 spec.desired, resource_id
491 ),
492 })
493 })
494}
495
496fn aggregate_workloads(workloads: &[Workload]) -> WorkloadRequirements {
497 let mut requirements = default_requirements();
498 requirements.total_cpu_at_desired = 0.0;
499 requirements.total_memory_bytes_at_desired = 0;
500 requirements.total_cpu_at_max = 0.0;
501 requirements.total_memory_bytes_at_max = 0;
502 requirements.max_cpu_per_container = 0.0;
503 requirements.max_memory_per_container = 0;
504 requirements.max_ephemeral_storage_bytes = 0;
505 requirements.gpu = None;
506
507 for workload in workloads {
508 requirements.total_cpu_at_desired += workload.cpu * workload.desired_replicas;
509 requirements.total_cpu_at_max += workload.cpu * workload.max_replicas;
510 requirements.total_memory_bytes_at_desired +=
511 (workload.memory_bytes as f64 * workload.desired_replicas) as u64;
512 requirements.total_memory_bytes_at_max +=
513 (workload.memory_bytes as f64 * workload.max_replicas) as u64;
514 requirements.max_cpu_per_container = requirements.max_cpu_per_container.max(workload.cpu);
515 requirements.max_memory_per_container = requirements
516 .max_memory_per_container
517 .max(workload.memory_bytes);
518 requirements.max_ephemeral_storage_bytes = requirements
519 .max_ephemeral_storage_bytes
520 .max(workload.ephemeral_storage_bytes);
521 if requirements.gpu.is_none() {
522 requirements.gpu = workload.gpu.clone();
523 }
524 }
525 requirements
526}
527
528fn default_requirements() -> WorkloadRequirements {
529 WorkloadRequirements {
530 total_cpu_at_desired: 1.0,
531 total_memory_bytes_at_desired: 2 * 1024 * 1024 * 1024,
532 total_cpu_at_max: 1.0,
533 total_memory_bytes_at_max: 2 * 1024 * 1024 * 1024,
534 max_cpu_per_container: 1.0,
535 max_memory_per_container: 2 * 1024 * 1024 * 1024,
536 max_ephemeral_storage_bytes: 0,
537 gpu: None,
538 architecture: None,
539 nested_virt: false,
540 }
541}
542
543fn profile_to_requirements(
544 profile: Option<&MachineProfile>,
545 nested_virt: bool,
546) -> WorkloadRequirements {
547 let Some(profile) = profile else {
548 return WorkloadRequirements {
549 nested_virt,
550 ..default_requirements()
551 };
552 };
553 let cpu = instance_catalog::parse_cpu(&profile.cpu).unwrap_or(1.0);
554 WorkloadRequirements {
555 total_cpu_at_desired: cpu,
556 total_memory_bytes_at_desired: profile.memory_bytes,
557 total_cpu_at_max: cpu,
558 total_memory_bytes_at_max: profile.memory_bytes,
559 max_cpu_per_container: cpu,
560 max_memory_per_container: profile.memory_bytes,
561 max_ephemeral_storage_bytes: profile.ephemeral_storage_bytes,
562 gpu: profile.gpu.clone(),
563 architecture: profile.architecture,
564 nested_virt,
565 }
566}
567
568fn merge_requirements(existing: &mut WorkloadRequirements, declared: &WorkloadRequirements) {
569 existing.total_cpu_at_desired = existing
570 .total_cpu_at_desired
571 .max(declared.total_cpu_at_desired);
572 existing.total_memory_bytes_at_desired = existing
573 .total_memory_bytes_at_desired
574 .max(declared.total_memory_bytes_at_desired);
575 existing.total_cpu_at_max = existing.total_cpu_at_max.max(declared.total_cpu_at_max);
576 existing.total_memory_bytes_at_max = existing
577 .total_memory_bytes_at_max
578 .max(declared.total_memory_bytes_at_max);
579 existing.max_cpu_per_container = existing
580 .max_cpu_per_container
581 .max(declared.max_cpu_per_container);
582 existing.max_memory_per_container = existing
583 .max_memory_per_container
584 .max(declared.max_memory_per_container);
585 existing.max_ephemeral_storage_bytes = existing
586 .max_ephemeral_storage_bytes
587 .max(declared.max_ephemeral_storage_bytes);
588 if existing.gpu.is_none() {
589 existing.gpu = declared.gpu.clone();
590 }
591 existing.nested_virt |= declared.nested_virt;
592 if existing.architecture.is_none() {
593 existing.architecture = declared.architecture;
594 }
595}
596
597fn requirements_to_profile(requirements: &WorkloadRequirements) -> MachineProfile {
598 MachineProfile {
599 cpu: requirements.max_cpu_per_container.to_string(),
600 memory_bytes: requirements.max_memory_per_container,
601 ephemeral_storage_bytes: requirements.max_ephemeral_storage_bytes,
602 architecture: requirements.architecture,
603 gpu: requirements.gpu.clone(),
604 }
605}
606
607fn merge_scale_policy(
608 existing: &CapacityGroupScalePolicy,
609 declared: &CapacityGroupScalePolicy,
610) -> CapacityGroupScalePolicy {
611 match (existing, declared) {
612 (
613 CapacityGroupScalePolicy::Fixed {
614 machines: existing_machines,
615 },
616 CapacityGroupScalePolicy::Fixed {
617 machines: declared_machines,
618 },
619 ) => CapacityGroupScalePolicy::Fixed {
620 machines: merge_choice_range(existing_machines, declared_machines),
621 },
622 (_, declared) => declared.clone(),
623 }
624}
625
626fn merge_choice_range(
627 existing: &ComputeChoiceRange,
628 declared: &ComputeChoiceRange,
629) -> ComputeChoiceRange {
630 ComputeChoiceRange {
631 min: existing.min.max(declared.min),
632 max: existing.max.max(declared.max),
633 default: declared.default,
634 }
635}
636
637fn validate_selection_against_scale(
638 selection: &ComputePoolSelection,
639 scale: &CapacityGroupScalePolicy,
640) -> std::result::Result<(), String> {
641 match (selection, scale) {
642 (
643 ComputePoolSelection::Fixed { machines, .. },
644 CapacityGroupScalePolicy::Fixed { machines: allowed },
645 ) => {
646 if allowed.contains(*machines) {
647 Ok(())
648 } else {
649 Err(format!(
650 "fixed machine count {machines} is outside the allowed range {}-{}",
651 allowed.min, allowed.max
652 ))
653 }
654 }
655 (
656 ComputePoolSelection::Autoscale { min, max, .. },
657 CapacityGroupScalePolicy::Autoscale {
658 min: allowed_min,
659 max: allowed_max,
660 },
661 ) => {
662 if !allowed_min.contains(*min) {
663 return Err(format!(
664 "autoscale minimum {min} is outside the allowed range {}-{}",
665 allowed_min.min, allowed_min.max
666 ));
667 }
668 if !allowed_max.contains(*max) {
669 return Err(format!(
670 "autoscale maximum {max} is outside the allowed range {}-{}",
671 allowed_max.min, allowed_max.max
672 ));
673 }
674 Ok(())
675 }
676 (ComputePoolSelection::Fixed { .. }, CapacityGroupScalePolicy::Autoscale { .. }) => {
677 Err("must use autoscale mode".to_string())
678 }
679 (ComputePoolSelection::Autoscale { .. }, CapacityGroupScalePolicy::Fixed { .. }) => {
680 Err("must use fixed mode".to_string())
681 }
682 }
683}
684
685fn needed_container_pool(container: &Container) -> &'static str {
686 if container.stateful && container.persistent_storage.is_some() {
687 return "stateful";
688 }
689 if container.gpu.is_some() {
690 return "gpu";
691 }
692 if let Some(storage) = &container.ephemeral_storage {
693 if instance_catalog::parse_memory_bytes(storage).unwrap_or(0) > 200 * 1024 * 1024 * 1024 {
694 return "storage";
695 }
696 }
697 "general"
698}
699
700fn default_min_machines(requirements: &WorkloadRequirements) -> u32 {
701 if requirements.total_cpu_at_desired > 0.0 || requirements.total_memory_bytes_at_desired > 0 {
702 1
703 } else {
704 0
705 }
706}
707
708fn default_max_machines(requirements: &WorkloadRequirements) -> u32 {
709 let min = default_min_machines(requirements);
710 let by_cpu =
711 (requirements.total_cpu_at_max / requirements.max_cpu_per_container.max(1.0)).ceil() as u32;
712 let by_mem = requirements
713 .total_memory_bytes_at_max
714 .div_ceil(requirements.max_memory_per_container.max(1)) as u32;
715 min.max(by_cpu).max(by_mem).max(1)
716}
717
718#[cfg(test)]
719mod tests {
720 use super::*;
721 use crate::{
722 instance_catalog::Architecture, CapacityGroup, CapacityGroupScalePolicy,
723 ComputeChoiceRange, ComputeCluster, ComputeSettings, ContainerCode, DaemonCode, Resource,
724 ResourceEntry, ResourceLifecycle, Stack, Storage, Worker, WorkerCode,
725 };
726
727 fn stack_with_container() -> Stack {
728 let container = Container::new("api".to_string())
729 .code(ContainerCode::Image {
730 image: "api:latest".to_string(),
731 })
732 .cpu(ResourceSpec {
733 min: "1".to_string(),
734 desired: "2".to_string(),
735 })
736 .memory(ResourceSpec {
737 min: "2Gi".to_string(),
738 desired: "4Gi".to_string(),
739 })
740 .permissions("api".to_string())
741 .build();
742 Stack {
743 id: "test".to_string(),
744 resources: [(
745 "api".to_string(),
746 ResourceEntry {
747 config: Resource::new(container),
748 lifecycle: ResourceLifecycle::Live,
749 dependencies: Vec::new(),
750 remote_access: false,
751 enabled_when: None,
752 },
753 )]
754 .into_iter()
755 .collect(),
756 permissions: crate::permissions::PermissionsConfig::default(),
757 supported_platforms: None,
758 inputs: vec![],
759 dynamic_container_repositories: Vec::new(),
760 dynamic_container_image_resources: Vec::new(),
761 }
762 }
763
764 #[test]
765 fn serverless_workers_and_storage_do_not_allocate_machine_pools() {
766 let mut stack = stack_with_container();
767 stack.resources.clear();
768 for (id, resource) in [
769 (
770 "api",
771 Resource::new(
772 Worker::new("api".to_string())
773 .code(WorkerCode::Image {
774 image: "example/worker:v1".to_string(),
775 })
776 .permissions("execution".to_string())
777 .build(),
778 ),
779 ),
780 (
781 "data",
782 Resource::new(Storage::new("data".to_string()).build()),
783 ),
784 ] {
785 stack.resources.insert(
786 id.to_string(),
787 ResourceEntry {
788 config: resource,
789 lifecycle: ResourceLifecycle::Frozen,
790 dependencies: Vec::new(),
791 remote_access: false,
792 enabled_when: None,
793 },
794 );
795 }
796
797 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
798 let plan = plan_compute(&stack, platform, None).expect("plan should build");
799 assert!(plan.pools.is_empty(), "serverless plan allocated {plan:?}");
800 }
801 }
802
803 #[test]
804 fn empty_stack_does_not_allocate_machine_pools() {
805 let mut stack = stack_with_container();
806 stack.resources.clear();
807 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
808 assert!(plan.pools.is_empty());
809 }
810
811 #[test]
812 fn explicit_cluster_without_workloads_only_plans_declared_pools() {
813 let mut stack = stack_with_container();
814 stack.resources.clear();
815 let cluster = ComputeCluster::new("compute".to_string())
816 .capacity_group(CapacityGroup {
817 group_id: "batch".to_string(),
818 instance_type: Some("t4g.medium".to_string()),
819 profile: None,
820 min_size: 1,
821 max_size: 1,
822 scale_policy: None,
823 nested_virtualization: None,
824 })
825 .build();
826 stack.resources.insert(
827 "compute".to_string(),
828 ResourceEntry {
829 config: Resource::new(cluster),
830 lifecycle: ResourceLifecycle::Frozen,
831 dependencies: Vec::new(),
832 remote_access: false,
833 enabled_when: None,
834 },
835 );
836
837 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
838 assert_eq!(plan.pools.len(), 1);
839 assert_eq!(plan.pools[0].pool_id, "batch");
840 assert!(plan.pools[0].errors.is_empty());
841 }
842
843 #[test]
844 fn cloud_plan_recommends_provider_machine_without_mutating_selection() {
845 let stack = stack_with_container();
846
847 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
848
849 let pool = plan.pools.first().expect("general pool should exist");
850 assert_eq!(pool.pool_id, "general");
851 assert_eq!(pool.workloads, vec!["api"]);
852 assert!(pool.selected.machine().is_some());
853 assert!(pool.machines.iter().any(|machine| machine.recommended));
854 }
855
856 #[test]
857 fn cloud_plan_only_offers_machines_matching_the_default_image_target() {
858 let stack = stack_with_container();
859
860 for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
861 let plan = plan_compute(&stack, platform, None).expect("plan should build");
862 let pool = plan.pools.first().expect("general pool should exist");
863 let expected = instance_catalog::default_architecture(platform)
864 .expect("managed cloud should have a default architecture");
865
866 assert!(pool.machines.iter().all(|machine| {
867 instance_catalog::find_instance_type(platform, &machine.machine)
868 .is_some_and(|spec| spec.architecture == expected)
869 }));
870 }
871 }
872
873 #[test]
874 fn selected_machine_is_preserved_as_static_deployment_choice() {
875 let stack = stack_with_container();
876 let settings = ComputeSettings {
877 pools: [(
878 "general".to_string(),
879 ComputePoolSelection::Fixed {
880 machines: 1,
881 machine: Some("m7g.xlarge".to_string()),
882 failure_domains: None,
883 },
884 )]
885 .into_iter()
886 .collect(),
887 };
888
889 let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
890
891 let pool = plan.pools.first().expect("general pool should exist");
892 assert_eq!(pool.selected.machine(), Some("m7g.xlarge"));
893 assert!(pool.errors.is_empty());
894 }
895
896 #[test]
897 fn configurable_cloud_disk_satisfies_storage_above_fixed_local_catalog() {
898 let mut stack = stack_with_container();
899 let container = stack
900 .resources
901 .get_mut("api")
902 .and_then(|entry| entry.config.downcast_mut::<Container>())
903 .expect("test stack should contain a container");
904 container.ephemeral_storage = Some("8000Gi".to_string());
905 let settings = ComputeSettings {
906 pools: [(
907 "storage".to_string(),
908 ComputePoolSelection::Fixed {
909 machines: 1,
910 machine: Some("m8i.2xlarge".to_string()),
911 failure_domains: None,
912 },
913 )]
914 .into_iter()
915 .collect(),
916 };
917
918 let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
919 let pool = plan.pools.first().expect("general pool should exist");
920 assert!(pool.errors.is_empty());
921 let selected = pool
922 .machines
923 .iter()
924 .find(|option| option.machine == "m8i.2xlarge")
925 .expect("selected machine should remain an option");
926 assert_eq!(
927 selected.profile.ephemeral_storage_bytes,
928 8_000 * 1024 * 1024 * 1024
929 );
930 }
931
932 #[test]
933 fn selected_machine_defines_architecture_when_workloads_do_not() {
934 let stack = stack_with_container();
935 let settings = ComputeSettings {
936 pools: [(
937 "general".to_string(),
938 ComputePoolSelection::Fixed {
939 machines: 1,
940 machine: Some("m7i.xlarge".to_string()),
941 failure_domains: None,
942 },
943 )]
944 .into_iter()
945 .collect(),
946 };
947
948 let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
949
950 let pool = plan.pools.first().expect("general pool should exist");
951 assert!(pool.errors.is_empty());
952 assert!(pool.machines.iter().all(|machine| {
953 instance_catalog::find_instance_type(Platform::Aws, &machine.machine)
954 .is_some_and(|spec| spec.architecture == Architecture::X86_64)
955 }));
956 }
957
958 #[test]
959 fn explicit_capacity_group_requirements_are_merged_with_workloads() {
960 let mut stack = stack_with_container();
961 let cluster = ComputeCluster::new("compute".to_string())
962 .capacity_group(CapacityGroup {
963 group_id: "general".to_string(),
964 instance_type: None,
965 profile: Some(MachineProfile {
966 cpu: "4".to_string(),
967 memory_bytes: 16 * 1024 * 1024 * 1024,
968 ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
969 architecture: Some(Architecture::X86_64),
970 gpu: None,
971 }),
972 min_size: 2,
973 max_size: 5,
974 scale_policy: None,
975 nested_virtualization: Some(true),
976 })
977 .build();
978 stack.resources.insert(
979 "compute".to_string(),
980 ResourceEntry {
981 config: Resource::new(cluster),
982 lifecycle: ResourceLifecycle::Frozen,
983 dependencies: Vec::new(),
984 remote_access: false,
985 enabled_when: None,
986 },
987 );
988
989 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
990
991 let pool = plan.pools.first().expect("general pool should exist");
992 let machine = pool
993 .selected
994 .machine()
995 .expect("AWS selection should include a machine");
996 let spec = instance_catalog::find_instance_type(Platform::Aws, machine)
997 .expect("selected machine should exist in the catalog");
998 assert!(spec.is_nested_virt_capable());
999 assert_eq!(pool.selected.min_size(), 2);
1000 assert_eq!(pool.selected.max_size(), 5);
1001 assert!(pool.errors.is_empty());
1002 }
1003
1004 #[test]
1005 fn materialized_machine_profile_is_not_counted_as_workload_demand() {
1006 let mut stack = stack_with_container();
1007 let selected = instance_catalog::find_instance_type(Platform::Aws, "m7g.xlarge")
1008 .expect("machine should exist in catalog");
1009 let cluster = ComputeCluster::new("compute".to_string())
1010 .capacity_group(CapacityGroup {
1011 group_id: "general".to_string(),
1012 instance_type: Some(selected.name.to_string()),
1013 profile: Some(selected.to_machine_profile()),
1014 min_size: 1,
1015 max_size: 1,
1016 scale_policy: None,
1017 nested_virtualization: None,
1018 })
1019 .build();
1020 stack.resources.insert(
1021 "compute".to_string(),
1022 ResourceEntry {
1023 config: Resource::new(cluster),
1024 lifecycle: ResourceLifecycle::Frozen,
1025 dependencies: Vec::new(),
1026 remote_access: false,
1027 enabled_when: None,
1028 },
1029 );
1030
1031 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
1032 let pool = plan.pools.first().expect("general pool should exist");
1033
1034 assert_eq!(pool.requirements.cpu, "2");
1035 assert_eq!(pool.requirements.memory_bytes, 4 * 1024 * 1024 * 1024);
1036 assert_eq!(pool.recommended.machine(), Some("m7g.xlarge"));
1037 assert!(pool
1038 .machines
1039 .iter()
1040 .any(|option| option.machine == "m7g.xlarge"));
1041 }
1042
1043 #[test]
1044 fn nested_x86_fixed_range_pool_preserves_bounds_and_rejects_graviton() {
1045 let daemon = Daemon::new("vm-runtime-loader".to_string())
1046 .code(DaemonCode::Image {
1047 image: "example.com/vm-runtime:latest".to_string(),
1048 })
1049 .cluster("vm-runtime".to_string())
1050 .cpu(ResourceSpec {
1051 min: "2".to_string(),
1052 desired: "2".to_string(),
1053 })
1054 .memory(ResourceSpec {
1055 min: "4Gi".to_string(),
1056 desired: "4Gi".to_string(),
1057 })
1058 .permissions("loader".to_string())
1059 .build();
1060 let cluster = ComputeCluster::new("vm-runtime".to_string())
1061 .capacity_group(CapacityGroup {
1062 group_id: "general".to_string(),
1063 instance_type: None,
1064 profile: Some(MachineProfile {
1065 cpu: "4".to_string(),
1066 memory_bytes: 16 * 1024 * 1024 * 1024,
1067 ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
1068 architecture: Some(Architecture::X86_64),
1069 gpu: None,
1070 }),
1071 min_size: 2,
1072 max_size: 2,
1073 scale_policy: Some(CapacityGroupScalePolicy::Fixed {
1074 machines: ComputeChoiceRange {
1075 min: 1,
1076 max: 5,
1077 default: 2,
1078 },
1079 }),
1080 nested_virtualization: Some(true),
1081 })
1082 .build();
1083 let stack = Stack {
1084 id: "vm-runtime".to_string(),
1085 resources: [
1086 (
1087 "vm-runtime-loader".to_string(),
1088 ResourceEntry {
1089 config: Resource::new(daemon),
1090 lifecycle: ResourceLifecycle::Live,
1091 dependencies: Vec::new(),
1092 remote_access: false,
1093 enabled_when: None,
1094 },
1095 ),
1096 (
1097 "vm-runtime".to_string(),
1098 ResourceEntry {
1099 config: Resource::new(cluster),
1100 lifecycle: ResourceLifecycle::Frozen,
1101 dependencies: Vec::new(),
1102 remote_access: false,
1103 enabled_when: None,
1104 },
1105 ),
1106 ]
1107 .into_iter()
1108 .collect(),
1109 permissions: crate::permissions::PermissionsConfig::default(),
1110 supported_platforms: None,
1111 inputs: vec![],
1112 dynamic_container_repositories: Vec::new(),
1113 dynamic_container_image_resources: Vec::new(),
1114 };
1115
1116 let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
1117 let pool = plan.pools.first().expect("general pool should exist");
1118 assert_eq!(pool.recommended.machine(), Some("m8i.2xlarge"));
1119 assert_eq!(pool.recommended.min_size(), 2);
1120 assert_eq!(pool.recommended.max_size(), 2);
1121 assert_eq!(
1122 pool.scale,
1123 CapacityGroupScalePolicy::Fixed {
1124 machines: ComputeChoiceRange {
1125 min: 1,
1126 max: 5,
1127 default: 2,
1128 },
1129 }
1130 );
1131 assert!(!pool
1132 .machines
1133 .iter()
1134 .any(|option| option.machine == "m7g.2xlarge"));
1135
1136 let invalid_settings = ComputeSettings {
1137 pools: [(
1138 "general".to_string(),
1139 ComputePoolSelection::Fixed {
1140 machines: 2,
1141 machine: Some("m7g.2xlarge".to_string()),
1142 failure_domains: None,
1143 },
1144 )]
1145 .into_iter()
1146 .collect(),
1147 };
1148 let invalid_plan = plan_compute(&stack, Platform::Aws, Some(&invalid_settings))
1149 .expect("plan should build");
1150 assert!(!invalid_plan.pools[0].errors.is_empty());
1151 }
1152
1153 #[test]
1154 fn local_plan_has_no_provider_machine_choices() {
1155 let stack = stack_with_container();
1156
1157 let plan = plan_compute(&stack, Platform::Local, None).expect("plan should build");
1158
1159 let pool = plan.pools.first().expect("general pool should exist");
1160 assert_eq!(pool.selected.machine(), None);
1161 assert!(pool.machines.is_empty());
1162 assert!(pool.errors.is_empty());
1163 }
1164}