Skip to main content

alien_core/
compute_planner.rs

1//! Deployment-time compute planner.
2//!
3//! The planner turns portable stack requirements plus a target platform into a
4//! renderable set of recommended deployment choices. It does not mutate the
5//! stack and does not require database access.
6
7use 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/// Full compute plan for one stack/platform pair.
18#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
19#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
20#[serde(rename_all = "camelCase")]
21pub struct ComputePlan {
22    /// Planned pools in stable pool-id order.
23    pub pools: Vec<ComputePoolPlan>,
24}
25
26/// Planner output for one compute pool.
27#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
28#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
29#[serde(rename_all = "camelCase")]
30pub struct ComputePoolPlan {
31    /// Pool ID from the stack or derived default.
32    pub pool_id: String,
33    /// Workloads assigned to this pool.
34    pub workloads: Vec<String>,
35    /// Aggregated requirements used for machine selection.
36    pub requirements: MachineProfile,
37    /// Allowed scale policy declared by source or derived for generated pools.
38    pub scale: CapacityGroupScalePolicy,
39    /// Recommended or user-selected deployment choice.
40    pub selected: ComputePoolSelection,
41    /// Planner-recommended default.
42    pub recommended: ComputePoolSelection,
43    /// Valid cloud machine choices. Empty for local and Kubernetes.
44    pub machines: Vec<ComputeMachineOption>,
45    /// Validation errors for supplied deployment settings.
46    #[serde(default, skip_serializing_if = "Vec::is_empty")]
47    pub errors: Vec<String>,
48}
49
50/// One concrete provider machine option.
51#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
52#[cfg_attr(feature = "openapi", derive(utoipa::ToSchema))]
53#[serde(rename_all = "camelCase")]
54pub struct ComputeMachineOption {
55    /// Provider machine name.
56    pub machine: String,
57    /// Machine hardware profile.
58    pub profile: MachineProfile,
59    /// Whether this machine is the planner's default recommendation.
60    pub recommended: bool,
61}
62
63/// Failure-domain policy for a pool that hosts persistent stateful workloads when the
64/// deployment does not choose one: one provider-selected domain, so the pool's machines
65/// and their volumes land in the same zone.
66pub fn default_persistent_failure_domains() -> FailureDomainSelection {
67    FailureDomainSelection {
68        spread: 1,
69        selected_failure_domains: Vec::new(),
70    }
71}
72
73/// Compute a deterministic deployment-time plan.
74pub 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                        // A materialized capacity group profile describes the
210                        // selected machine, not additional workload demand.
211                        // Merging its full CPU and memory into requirements
212                        // would cause deployment package planning to size an
213                        // already-sized machine a second time. Keep only the
214                        // constraints that affect artifact compatibility.
215                        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
273/// Validate one selected compute pool against platform machine requirements and
274/// source-declared scale bounds.
275pub 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
315/// Convert a capacity group declaration into planner requirements.
316pub 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}