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/// Compute a deterministic deployment-time plan.
64pub 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                        // A materialized capacity group profile describes the
200                        // selected machine, not additional workload demand.
201                        // Merging its full CPU and memory into requirements
202                        // would cause deployment package planning to size an
203                        // already-sized machine a second time. Keep only the
204                        // constraints that affect artifact compatibility.
205                        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
266/// Validate one selected compute pool against platform machine requirements and
267/// source-declared scale bounds.
268pub 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
308/// Convert a capacity group declaration into planner requirements.
309pub 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_for_storage(requirements.max_ephemeral_storage_bytes),
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.has_configurable_ephemeral_storage() {
377        if crate::instance_catalog::max_configurable_ephemeral_storage_bytes(spec.platform)
378            .is_none_or(|max| requirements.max_ephemeral_storage_bytes > max)
379        {
380            return false;
381        }
382    } else if spec.ephemeral_storage_bytes < requirements.max_ephemeral_storage_bytes {
383        return false;
384    }
385    match (&requirements.gpu, spec.gpu) {
386        (Some(required), Some(actual)) => {
387            (required.gpu_type == "any" || required.gpu_type == actual.gpu_type)
388                && actual.count >= required.count
389        }
390        (Some(_), None) => false,
391        (None, _) => true,
392    }
393}
394
395#[derive(Debug, Clone)]
396struct Workload {
397    id: String,
398    cpu: f64,
399    memory_bytes: u64,
400    desired_replicas: f64,
401    max_replicas: f64,
402    ephemeral_storage_bytes: u64,
403    gpu: Option<GpuSpec>,
404    requires_failure_domain: bool,
405}
406
407impl Workload {
408    fn from_container(container: &Container) -> Result<Self, ErrorData> {
409        let cpu = parse_cpu(&container.id, &container.cpu)?;
410        let memory_bytes = parse_memory(&container.id, &container.memory)?;
411        let desired_replicas = container
412            .autoscaling
413            .as_ref()
414            .map(|a| a.desired)
415            .or(container.replicas)
416            .unwrap_or(1) as f64;
417        let max_replicas = container
418            .autoscaling
419            .as_ref()
420            .map(|a| a.max)
421            .or(container.replicas)
422            .unwrap_or(1) as f64;
423        let ephemeral_storage_bytes = container
424            .ephemeral_storage
425            .as_deref()
426            .map(instance_catalog::parse_memory_bytes)
427            .transpose()
428            .map_err(|message| {
429                AlienError::new(ErrorData::GenericError {
430                    message: format!(
431                        "Failed to parse ephemeral storage for '{}': {message}",
432                        container.id
433                    ),
434                })
435            })?
436            .unwrap_or(0);
437
438        Ok(Self {
439            id: container.id.clone(),
440            cpu,
441            memory_bytes,
442            desired_replicas,
443            max_replicas,
444            ephemeral_storage_bytes,
445            gpu: container.gpu.as_ref().map(|gpu| GpuSpec {
446                gpu_type: gpu.gpu_type.clone(),
447                count: gpu.count,
448            }),
449            requires_failure_domain: container.stateful && container.persistent_storage.is_some(),
450        })
451    }
452
453    fn from_daemon(daemon: &Daemon) -> Result<Self, ErrorData> {
454        Ok(Self {
455            id: daemon.id.clone(),
456            cpu: parse_cpu(&daemon.id, &daemon.cpu)?,
457            memory_bytes: parse_memory(&daemon.id, &daemon.memory)?,
458            desired_replicas: 1.0,
459            max_replicas: 1.0,
460            ephemeral_storage_bytes: 0,
461            gpu: None,
462            requires_failure_domain: false,
463        })
464    }
465}
466
467fn parse_cpu(resource_id: &str, spec: &ResourceSpec) -> Result<f64, ErrorData> {
468    instance_catalog::parse_cpu(&spec.desired).map_err(|message| {
469        AlienError::new(ErrorData::GenericError {
470            message: format!(
471                "Failed to parse CPU requirement '{}' for '{}': {message}",
472                spec.desired, resource_id
473            ),
474        })
475    })
476}
477
478fn parse_memory(resource_id: &str, spec: &ResourceSpec) -> Result<u64, ErrorData> {
479    instance_catalog::parse_memory_bytes(&spec.desired).map_err(|message| {
480        AlienError::new(ErrorData::GenericError {
481            message: format!(
482                "Failed to parse memory requirement '{}' for '{}': {message}",
483                spec.desired, resource_id
484            ),
485        })
486    })
487}
488
489fn aggregate_workloads(workloads: &[Workload]) -> WorkloadRequirements {
490    let mut requirements = default_requirements();
491    requirements.total_cpu_at_desired = 0.0;
492    requirements.total_memory_bytes_at_desired = 0;
493    requirements.total_cpu_at_max = 0.0;
494    requirements.total_memory_bytes_at_max = 0;
495    requirements.max_cpu_per_container = 0.0;
496    requirements.max_memory_per_container = 0;
497    requirements.max_ephemeral_storage_bytes = 0;
498    requirements.gpu = None;
499
500    for workload in workloads {
501        requirements.total_cpu_at_desired += workload.cpu * workload.desired_replicas;
502        requirements.total_cpu_at_max += workload.cpu * workload.max_replicas;
503        requirements.total_memory_bytes_at_desired +=
504            (workload.memory_bytes as f64 * workload.desired_replicas) as u64;
505        requirements.total_memory_bytes_at_max +=
506            (workload.memory_bytes as f64 * workload.max_replicas) as u64;
507        requirements.max_cpu_per_container = requirements.max_cpu_per_container.max(workload.cpu);
508        requirements.max_memory_per_container = requirements
509            .max_memory_per_container
510            .max(workload.memory_bytes);
511        requirements.max_ephemeral_storage_bytes = requirements
512            .max_ephemeral_storage_bytes
513            .max(workload.ephemeral_storage_bytes);
514        if requirements.gpu.is_none() {
515            requirements.gpu = workload.gpu.clone();
516        }
517    }
518    requirements
519}
520
521fn default_requirements() -> WorkloadRequirements {
522    WorkloadRequirements {
523        total_cpu_at_desired: 1.0,
524        total_memory_bytes_at_desired: 2 * 1024 * 1024 * 1024,
525        total_cpu_at_max: 1.0,
526        total_memory_bytes_at_max: 2 * 1024 * 1024 * 1024,
527        max_cpu_per_container: 1.0,
528        max_memory_per_container: 2 * 1024 * 1024 * 1024,
529        max_ephemeral_storage_bytes: 0,
530        gpu: None,
531        architecture: None,
532        nested_virt: false,
533    }
534}
535
536fn profile_to_requirements(
537    profile: Option<&MachineProfile>,
538    nested_virt: bool,
539) -> WorkloadRequirements {
540    let Some(profile) = profile else {
541        return WorkloadRequirements {
542            nested_virt,
543            ..default_requirements()
544        };
545    };
546    let cpu = instance_catalog::parse_cpu(&profile.cpu).unwrap_or(1.0);
547    WorkloadRequirements {
548        total_cpu_at_desired: cpu,
549        total_memory_bytes_at_desired: profile.memory_bytes,
550        total_cpu_at_max: cpu,
551        total_memory_bytes_at_max: profile.memory_bytes,
552        max_cpu_per_container: cpu,
553        max_memory_per_container: profile.memory_bytes,
554        max_ephemeral_storage_bytes: profile.ephemeral_storage_bytes,
555        gpu: profile.gpu.clone(),
556        architecture: profile.architecture,
557        nested_virt,
558    }
559}
560
561fn merge_requirements(existing: &mut WorkloadRequirements, declared: &WorkloadRequirements) {
562    existing.total_cpu_at_desired = existing
563        .total_cpu_at_desired
564        .max(declared.total_cpu_at_desired);
565    existing.total_memory_bytes_at_desired = existing
566        .total_memory_bytes_at_desired
567        .max(declared.total_memory_bytes_at_desired);
568    existing.total_cpu_at_max = existing.total_cpu_at_max.max(declared.total_cpu_at_max);
569    existing.total_memory_bytes_at_max = existing
570        .total_memory_bytes_at_max
571        .max(declared.total_memory_bytes_at_max);
572    existing.max_cpu_per_container = existing
573        .max_cpu_per_container
574        .max(declared.max_cpu_per_container);
575    existing.max_memory_per_container = existing
576        .max_memory_per_container
577        .max(declared.max_memory_per_container);
578    existing.max_ephemeral_storage_bytes = existing
579        .max_ephemeral_storage_bytes
580        .max(declared.max_ephemeral_storage_bytes);
581    if existing.gpu.is_none() {
582        existing.gpu = declared.gpu.clone();
583    }
584    existing.nested_virt |= declared.nested_virt;
585    if existing.architecture.is_none() {
586        existing.architecture = declared.architecture;
587    }
588}
589
590fn requirements_to_profile(requirements: &WorkloadRequirements) -> MachineProfile {
591    MachineProfile {
592        cpu: requirements.max_cpu_per_container.to_string(),
593        memory_bytes: requirements.max_memory_per_container,
594        ephemeral_storage_bytes: requirements.max_ephemeral_storage_bytes,
595        architecture: requirements.architecture,
596        gpu: requirements.gpu.clone(),
597    }
598}
599
600fn merge_scale_policy(
601    existing: &CapacityGroupScalePolicy,
602    declared: &CapacityGroupScalePolicy,
603) -> CapacityGroupScalePolicy {
604    match (existing, declared) {
605        (
606            CapacityGroupScalePolicy::Fixed {
607                machines: existing_machines,
608            },
609            CapacityGroupScalePolicy::Fixed {
610                machines: declared_machines,
611            },
612        ) => CapacityGroupScalePolicy::Fixed {
613            machines: merge_choice_range(existing_machines, declared_machines),
614        },
615        (_, declared) => declared.clone(),
616    }
617}
618
619fn merge_choice_range(
620    existing: &ComputeChoiceRange,
621    declared: &ComputeChoiceRange,
622) -> ComputeChoiceRange {
623    ComputeChoiceRange {
624        min: existing.min.max(declared.min),
625        max: existing.max.max(declared.max),
626        default: declared.default,
627    }
628}
629
630fn validate_selection_against_scale(
631    selection: &ComputePoolSelection,
632    scale: &CapacityGroupScalePolicy,
633) -> std::result::Result<(), String> {
634    match (selection, scale) {
635        (
636            ComputePoolSelection::Fixed { machines, .. },
637            CapacityGroupScalePolicy::Fixed { machines: allowed },
638        ) => {
639            if allowed.contains(*machines) {
640                Ok(())
641            } else {
642                Err(format!(
643                    "fixed machine count {machines} is outside the allowed range {}-{}",
644                    allowed.min, allowed.max
645                ))
646            }
647        }
648        (
649            ComputePoolSelection::Autoscale { min, max, .. },
650            CapacityGroupScalePolicy::Autoscale {
651                min: allowed_min,
652                max: allowed_max,
653            },
654        ) => {
655            if !allowed_min.contains(*min) {
656                return Err(format!(
657                    "autoscale minimum {min} is outside the allowed range {}-{}",
658                    allowed_min.min, allowed_min.max
659                ));
660            }
661            if !allowed_max.contains(*max) {
662                return Err(format!(
663                    "autoscale maximum {max} is outside the allowed range {}-{}",
664                    allowed_max.min, allowed_max.max
665                ));
666            }
667            Ok(())
668        }
669        (ComputePoolSelection::Fixed { .. }, CapacityGroupScalePolicy::Autoscale { .. }) => {
670            Err("must use autoscale mode".to_string())
671        }
672        (ComputePoolSelection::Autoscale { .. }, CapacityGroupScalePolicy::Fixed { .. }) => {
673            Err("must use fixed mode".to_string())
674        }
675    }
676}
677
678fn needed_container_pool(container: &Container) -> &'static str {
679    if container.stateful && container.persistent_storage.is_some() {
680        return "stateful";
681    }
682    if container.gpu.is_some() {
683        return "gpu";
684    }
685    if let Some(storage) = &container.ephemeral_storage {
686        if instance_catalog::parse_memory_bytes(storage).unwrap_or(0) > 200 * 1024 * 1024 * 1024 {
687            return "storage";
688        }
689    }
690    "general"
691}
692
693fn default_min_machines(requirements: &WorkloadRequirements) -> u32 {
694    if requirements.total_cpu_at_desired > 0.0 || requirements.total_memory_bytes_at_desired > 0 {
695        1
696    } else {
697        0
698    }
699}
700
701fn default_max_machines(requirements: &WorkloadRequirements) -> u32 {
702    let min = default_min_machines(requirements);
703    let by_cpu =
704        (requirements.total_cpu_at_max / requirements.max_cpu_per_container.max(1.0)).ceil() as u32;
705    let by_mem = requirements
706        .total_memory_bytes_at_max
707        .div_ceil(requirements.max_memory_per_container.max(1)) as u32;
708    min.max(by_cpu).max(by_mem).max(1)
709}
710
711#[cfg(test)]
712mod tests {
713    use super::*;
714    use crate::{
715        instance_catalog::Architecture, CapacityGroup, CapacityGroupScalePolicy,
716        ComputeChoiceRange, ComputeCluster, ComputeSettings, ContainerCode, DaemonCode, Resource,
717        ResourceEntry, ResourceLifecycle, Stack, Storage, Worker, WorkerCode,
718    };
719
720    fn stack_with_container() -> Stack {
721        let container = Container::new("api".to_string())
722            .code(ContainerCode::Image {
723                image: "api:latest".to_string(),
724            })
725            .cpu(ResourceSpec {
726                min: "1".to_string(),
727                desired: "2".to_string(),
728            })
729            .memory(ResourceSpec {
730                min: "2Gi".to_string(),
731                desired: "4Gi".to_string(),
732            })
733            .permissions("api".to_string())
734            .build();
735        Stack {
736            id: "test".to_string(),
737            resources: [(
738                "api".to_string(),
739                ResourceEntry {
740                    config: Resource::new(container),
741                    lifecycle: ResourceLifecycle::Live,
742                    dependencies: Vec::new(),
743                    remote_access: false,
744                    enabled_when: None,
745                },
746            )]
747            .into_iter()
748            .collect(),
749            permissions: crate::permissions::PermissionsConfig::default(),
750            supported_platforms: None,
751            inputs: vec![],
752            dynamic_container_repositories: Vec::new(),
753            dynamic_container_image_resources: Vec::new(),
754        }
755    }
756
757    #[test]
758    fn serverless_workers_and_storage_do_not_allocate_machine_pools() {
759        let mut stack = stack_with_container();
760        stack.resources.clear();
761        for (id, resource) in [
762            (
763                "api",
764                Resource::new(
765                    Worker::new("api".to_string())
766                        .code(WorkerCode::Image {
767                            image: "example/worker:v1".to_string(),
768                        })
769                        .permissions("execution".to_string())
770                        .build(),
771                ),
772            ),
773            (
774                "data",
775                Resource::new(Storage::new("data".to_string()).build()),
776            ),
777        ] {
778            stack.resources.insert(
779                id.to_string(),
780                ResourceEntry {
781                    config: resource,
782                    lifecycle: ResourceLifecycle::Frozen,
783                    dependencies: Vec::new(),
784                    remote_access: false,
785                    enabled_when: None,
786                },
787            );
788        }
789
790        for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
791            let plan = plan_compute(&stack, platform, None).expect("plan should build");
792            assert!(plan.pools.is_empty(), "serverless plan allocated {plan:?}");
793        }
794    }
795
796    #[test]
797    fn empty_stack_does_not_allocate_machine_pools() {
798        let mut stack = stack_with_container();
799        stack.resources.clear();
800        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
801        assert!(plan.pools.is_empty());
802    }
803
804    #[test]
805    fn explicit_cluster_without_workloads_only_plans_declared_pools() {
806        let mut stack = stack_with_container();
807        stack.resources.clear();
808        let cluster = ComputeCluster::new("compute".to_string())
809            .capacity_group(CapacityGroup {
810                group_id: "batch".to_string(),
811                instance_type: Some("t4g.medium".to_string()),
812                profile: None,
813                min_size: 1,
814                max_size: 1,
815                scale_policy: None,
816                nested_virtualization: None,
817            })
818            .build();
819        stack.resources.insert(
820            "compute".to_string(),
821            ResourceEntry {
822                config: Resource::new(cluster),
823                lifecycle: ResourceLifecycle::Frozen,
824                dependencies: Vec::new(),
825                remote_access: false,
826                enabled_when: None,
827            },
828        );
829
830        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
831        assert_eq!(plan.pools.len(), 1);
832        assert_eq!(plan.pools[0].pool_id, "batch");
833        assert!(plan.pools[0].errors.is_empty());
834    }
835
836    #[test]
837    fn cloud_plan_recommends_provider_machine_without_mutating_selection() {
838        let stack = stack_with_container();
839
840        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
841
842        let pool = plan.pools.first().expect("general pool should exist");
843        assert_eq!(pool.pool_id, "general");
844        assert_eq!(pool.workloads, vec!["api"]);
845        assert!(pool.selected.machine().is_some());
846        assert!(pool.machines.iter().any(|machine| machine.recommended));
847    }
848
849    #[test]
850    fn cloud_plan_only_offers_machines_matching_the_default_image_target() {
851        let stack = stack_with_container();
852
853        for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
854            let plan = plan_compute(&stack, platform, None).expect("plan should build");
855            let pool = plan.pools.first().expect("general pool should exist");
856            let expected = instance_catalog::default_architecture(platform)
857                .expect("managed cloud should have a default architecture");
858
859            assert!(pool.machines.iter().all(|machine| {
860                instance_catalog::find_instance_type(platform, &machine.machine)
861                    .is_some_and(|spec| spec.architecture == expected)
862            }));
863        }
864    }
865
866    #[test]
867    fn selected_machine_is_preserved_as_static_deployment_choice() {
868        let stack = stack_with_container();
869        let settings = ComputeSettings {
870            pools: [(
871                "general".to_string(),
872                ComputePoolSelection::Fixed {
873                    machines: 1,
874                    machine: Some("m7g.xlarge".to_string()),
875                    failure_domains: None,
876                },
877            )]
878            .into_iter()
879            .collect(),
880        };
881
882        let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
883
884        let pool = plan.pools.first().expect("general pool should exist");
885        assert_eq!(pool.selected.machine(), Some("m7g.xlarge"));
886        assert!(pool.errors.is_empty());
887    }
888
889    #[test]
890    fn configurable_cloud_disk_satisfies_storage_above_fixed_local_catalog() {
891        let mut stack = stack_with_container();
892        let container = stack
893            .resources
894            .get_mut("api")
895            .and_then(|entry| entry.config.downcast_mut::<Container>())
896            .expect("test stack should contain a container");
897        container.ephemeral_storage = Some("8000Gi".to_string());
898        let settings = ComputeSettings {
899            pools: [(
900                "storage".to_string(),
901                ComputePoolSelection::Fixed {
902                    machines: 1,
903                    machine: Some("m8i.2xlarge".to_string()),
904                    failure_domains: None,
905                },
906            )]
907            .into_iter()
908            .collect(),
909        };
910
911        let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
912        let pool = plan.pools.first().expect("general pool should exist");
913        assert!(pool.errors.is_empty());
914        let selected = pool
915            .machines
916            .iter()
917            .find(|option| option.machine == "m8i.2xlarge")
918            .expect("selected machine should remain an option");
919        assert_eq!(
920            selected.profile.ephemeral_storage_bytes,
921            8_000 * 1024 * 1024 * 1024
922        );
923    }
924
925    #[test]
926    fn selected_machine_defines_architecture_when_workloads_do_not() {
927        let stack = stack_with_container();
928        let settings = ComputeSettings {
929            pools: [(
930                "general".to_string(),
931                ComputePoolSelection::Fixed {
932                    machines: 1,
933                    machine: Some("m7i.xlarge".to_string()),
934                    failure_domains: None,
935                },
936            )]
937            .into_iter()
938            .collect(),
939        };
940
941        let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
942
943        let pool = plan.pools.first().expect("general pool should exist");
944        assert!(pool.errors.is_empty());
945        assert!(pool.machines.iter().all(|machine| {
946            instance_catalog::find_instance_type(Platform::Aws, &machine.machine)
947                .is_some_and(|spec| spec.architecture == Architecture::X86_64)
948        }));
949    }
950
951    #[test]
952    fn explicit_capacity_group_requirements_are_merged_with_workloads() {
953        let mut stack = stack_with_container();
954        let cluster = ComputeCluster::new("compute".to_string())
955            .capacity_group(CapacityGroup {
956                group_id: "general".to_string(),
957                instance_type: None,
958                profile: Some(MachineProfile {
959                    cpu: "4".to_string(),
960                    memory_bytes: 16 * 1024 * 1024 * 1024,
961                    ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
962                    architecture: Some(Architecture::X86_64),
963                    gpu: None,
964                }),
965                min_size: 2,
966                max_size: 5,
967                scale_policy: None,
968                nested_virtualization: Some(true),
969            })
970            .build();
971        stack.resources.insert(
972            "compute".to_string(),
973            ResourceEntry {
974                config: Resource::new(cluster),
975                lifecycle: ResourceLifecycle::Frozen,
976                dependencies: Vec::new(),
977                remote_access: false,
978                enabled_when: None,
979            },
980        );
981
982        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
983
984        let pool = plan.pools.first().expect("general pool should exist");
985        let machine = pool
986            .selected
987            .machine()
988            .expect("AWS selection should include a machine");
989        let spec = instance_catalog::find_instance_type(Platform::Aws, machine)
990            .expect("selected machine should exist in the catalog");
991        assert!(spec.is_nested_virt_capable());
992        assert_eq!(pool.selected.min_size(), 2);
993        assert_eq!(pool.selected.max_size(), 5);
994        assert!(pool.errors.is_empty());
995    }
996
997    #[test]
998    fn materialized_machine_profile_is_not_counted_as_workload_demand() {
999        let mut stack = stack_with_container();
1000        let selected = instance_catalog::find_instance_type(Platform::Aws, "m7g.xlarge")
1001            .expect("machine should exist in catalog");
1002        let cluster = ComputeCluster::new("compute".to_string())
1003            .capacity_group(CapacityGroup {
1004                group_id: "general".to_string(),
1005                instance_type: Some(selected.name.to_string()),
1006                profile: Some(selected.to_machine_profile()),
1007                min_size: 1,
1008                max_size: 1,
1009                scale_policy: None,
1010                nested_virtualization: None,
1011            })
1012            .build();
1013        stack.resources.insert(
1014            "compute".to_string(),
1015            ResourceEntry {
1016                config: Resource::new(cluster),
1017                lifecycle: ResourceLifecycle::Frozen,
1018                dependencies: Vec::new(),
1019                remote_access: false,
1020                enabled_when: None,
1021            },
1022        );
1023
1024        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
1025        let pool = plan.pools.first().expect("general pool should exist");
1026
1027        assert_eq!(pool.requirements.cpu, "2");
1028        assert_eq!(pool.requirements.memory_bytes, 4 * 1024 * 1024 * 1024);
1029        assert_eq!(pool.recommended.machine(), Some("m7g.xlarge"));
1030        assert!(pool
1031            .machines
1032            .iter()
1033            .any(|option| option.machine == "m7g.xlarge"));
1034    }
1035
1036    #[test]
1037    fn nested_x86_fixed_range_pool_preserves_bounds_and_rejects_graviton() {
1038        let daemon = Daemon::new("vm-runtime-loader".to_string())
1039            .code(DaemonCode::Image {
1040                image: "example.com/vm-runtime:latest".to_string(),
1041            })
1042            .cluster("vm-runtime".to_string())
1043            .cpu(ResourceSpec {
1044                min: "2".to_string(),
1045                desired: "2".to_string(),
1046            })
1047            .memory(ResourceSpec {
1048                min: "4Gi".to_string(),
1049                desired: "4Gi".to_string(),
1050            })
1051            .permissions("loader".to_string())
1052            .build();
1053        let cluster = ComputeCluster::new("vm-runtime".to_string())
1054            .capacity_group(CapacityGroup {
1055                group_id: "general".to_string(),
1056                instance_type: None,
1057                profile: Some(MachineProfile {
1058                    cpu: "4".to_string(),
1059                    memory_bytes: 16 * 1024 * 1024 * 1024,
1060                    ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
1061                    architecture: Some(Architecture::X86_64),
1062                    gpu: None,
1063                }),
1064                min_size: 2,
1065                max_size: 2,
1066                scale_policy: Some(CapacityGroupScalePolicy::Fixed {
1067                    machines: ComputeChoiceRange {
1068                        min: 1,
1069                        max: 5,
1070                        default: 2,
1071                    },
1072                }),
1073                nested_virtualization: Some(true),
1074            })
1075            .build();
1076        let stack = Stack {
1077            id: "vm-runtime".to_string(),
1078            resources: [
1079                (
1080                    "vm-runtime-loader".to_string(),
1081                    ResourceEntry {
1082                        config: Resource::new(daemon),
1083                        lifecycle: ResourceLifecycle::Live,
1084                        dependencies: Vec::new(),
1085                        remote_access: false,
1086                        enabled_when: None,
1087                    },
1088                ),
1089                (
1090                    "vm-runtime".to_string(),
1091                    ResourceEntry {
1092                        config: Resource::new(cluster),
1093                        lifecycle: ResourceLifecycle::Frozen,
1094                        dependencies: Vec::new(),
1095                        remote_access: false,
1096                        enabled_when: None,
1097                    },
1098                ),
1099            ]
1100            .into_iter()
1101            .collect(),
1102            permissions: crate::permissions::PermissionsConfig::default(),
1103            supported_platforms: None,
1104            inputs: vec![],
1105            dynamic_container_repositories: Vec::new(),
1106            dynamic_container_image_resources: Vec::new(),
1107        };
1108
1109        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
1110        let pool = plan.pools.first().expect("general pool should exist");
1111        assert_eq!(pool.recommended.machine(), Some("m8i.2xlarge"));
1112        assert_eq!(pool.recommended.min_size(), 2);
1113        assert_eq!(pool.recommended.max_size(), 2);
1114        assert_eq!(
1115            pool.scale,
1116            CapacityGroupScalePolicy::Fixed {
1117                machines: ComputeChoiceRange {
1118                    min: 1,
1119                    max: 5,
1120                    default: 2,
1121                },
1122            }
1123        );
1124        assert!(!pool
1125            .machines
1126            .iter()
1127            .any(|option| option.machine == "m7g.2xlarge"));
1128
1129        let invalid_settings = ComputeSettings {
1130            pools: [(
1131                "general".to_string(),
1132                ComputePoolSelection::Fixed {
1133                    machines: 2,
1134                    machine: Some("m7g.2xlarge".to_string()),
1135                    failure_domains: None,
1136                },
1137            )]
1138            .into_iter()
1139            .collect(),
1140        };
1141        let invalid_plan = plan_compute(&stack, Platform::Aws, Some(&invalid_settings))
1142            .expect("plan should build");
1143        assert!(!invalid_plan.pools[0].errors.is_empty());
1144    }
1145
1146    #[test]
1147    fn local_plan_has_no_provider_machine_choices() {
1148        let stack = stack_with_container();
1149
1150        let plan = plan_compute(&stack, Platform::Local, None).expect("plan should build");
1151
1152        let pool = plan.pools.first().expect("general pool should exist");
1153        assert_eq!(pool.selected.machine(), None);
1154        assert!(pool.machines.is_empty());
1155        assert!(pool.errors.is_empty());
1156    }
1157}