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