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