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                    enabled_when: None,
751                },
752            )]
753            .into_iter()
754            .collect(),
755            permissions: crate::permissions::PermissionsConfig::default(),
756            supported_platforms: None,
757            inputs: vec![],
758        }
759    }
760
761    #[test]
762    fn cloud_plan_recommends_provider_machine_without_mutating_selection() {
763        let stack = stack_with_container();
764
765        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
766
767        let pool = plan.pools.first().expect("general pool should exist");
768        assert_eq!(pool.pool_id, "general");
769        assert_eq!(pool.workloads, vec!["api"]);
770        assert!(pool.selected.machine().is_some());
771        assert!(pool.machines.iter().any(|machine| machine.recommended));
772    }
773
774    #[test]
775    fn cloud_plan_only_offers_machines_matching_the_default_image_target() {
776        let stack = stack_with_container();
777
778        for platform in [Platform::Aws, Platform::Gcp, Platform::Azure] {
779            let plan = plan_compute(&stack, platform, None).expect("plan should build");
780            let pool = plan.pools.first().expect("general pool should exist");
781            let expected = instance_catalog::default_architecture(platform)
782                .expect("managed cloud should have a default architecture");
783
784            assert!(pool.machines.iter().all(|machine| {
785                instance_catalog::find_instance_type(platform, &machine.machine)
786                    .is_some_and(|spec| spec.architecture == expected)
787            }));
788        }
789    }
790
791    #[test]
792    fn selected_machine_is_preserved_as_static_deployment_choice() {
793        let stack = stack_with_container();
794        let settings = ComputeSettings {
795            pools: [(
796                "general".to_string(),
797                ComputePoolSelection::Fixed {
798                    machines: 1,
799                    machine: Some("m7g.xlarge".to_string()),
800                    failure_domains: None,
801                },
802            )]
803            .into_iter()
804            .collect(),
805        };
806
807        let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
808
809        let pool = plan.pools.first().expect("general pool should exist");
810        assert_eq!(pool.selected.machine(), Some("m7g.xlarge"));
811        assert!(pool.errors.is_empty());
812    }
813
814    #[test]
815    fn selected_machine_defines_architecture_when_workloads_do_not() {
816        let stack = stack_with_container();
817        let settings = ComputeSettings {
818            pools: [(
819                "general".to_string(),
820                ComputePoolSelection::Fixed {
821                    machines: 1,
822                    machine: Some("m7i.xlarge".to_string()),
823                    failure_domains: None,
824                },
825            )]
826            .into_iter()
827            .collect(),
828        };
829
830        let plan = plan_compute(&stack, Platform::Aws, Some(&settings)).expect("plan should build");
831
832        let pool = plan.pools.first().expect("general pool should exist");
833        assert!(pool.errors.is_empty());
834        assert!(pool.machines.iter().all(|machine| {
835            instance_catalog::find_instance_type(Platform::Aws, &machine.machine)
836                .is_some_and(|spec| spec.architecture == Architecture::X86_64)
837        }));
838    }
839
840    #[test]
841    fn explicit_capacity_group_requirements_are_merged_with_workloads() {
842        let mut stack = stack_with_container();
843        let cluster = ComputeCluster::new("compute".to_string())
844            .capacity_group(CapacityGroup {
845                group_id: "general".to_string(),
846                instance_type: None,
847                profile: Some(MachineProfile {
848                    cpu: "4".to_string(),
849                    memory_bytes: 16 * 1024 * 1024 * 1024,
850                    ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
851                    architecture: Some(Architecture::X86_64),
852                    gpu: None,
853                }),
854                min_size: 2,
855                max_size: 5,
856                scale_policy: None,
857                nested_virtualization: Some(true),
858            })
859            .build();
860        stack.resources.insert(
861            "compute".to_string(),
862            ResourceEntry {
863                config: Resource::new(cluster),
864                lifecycle: ResourceLifecycle::Frozen,
865                dependencies: Vec::new(),
866                remote_access: false,
867                enabled_when: None,
868            },
869        );
870
871        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
872
873        let pool = plan.pools.first().expect("general pool should exist");
874        let machine = pool
875            .selected
876            .machine()
877            .expect("AWS selection should include a machine");
878        let spec = instance_catalog::find_instance_type(Platform::Aws, machine)
879            .expect("selected machine should exist in the catalog");
880        assert!(spec.is_nested_virt_capable());
881        assert_eq!(pool.selected.min_size(), 2);
882        assert_eq!(pool.selected.max_size(), 5);
883        assert!(pool.errors.is_empty());
884    }
885
886    #[test]
887    fn materialized_machine_profile_is_not_counted_as_workload_demand() {
888        let mut stack = stack_with_container();
889        let selected = instance_catalog::find_instance_type(Platform::Aws, "m7g.xlarge")
890            .expect("machine should exist in catalog");
891        let cluster = ComputeCluster::new("compute".to_string())
892            .capacity_group(CapacityGroup {
893                group_id: "general".to_string(),
894                instance_type: Some(selected.name.to_string()),
895                profile: Some(selected.to_machine_profile()),
896                min_size: 1,
897                max_size: 1,
898                scale_policy: None,
899                nested_virtualization: None,
900            })
901            .build();
902        stack.resources.insert(
903            "compute".to_string(),
904            ResourceEntry {
905                config: Resource::new(cluster),
906                lifecycle: ResourceLifecycle::Frozen,
907                dependencies: Vec::new(),
908                remote_access: false,
909                enabled_when: None,
910            },
911        );
912
913        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
914        let pool = plan.pools.first().expect("general pool should exist");
915
916        assert_eq!(pool.requirements.cpu, "2");
917        assert_eq!(pool.requirements.memory_bytes, 4 * 1024 * 1024 * 1024);
918        assert_eq!(pool.recommended.machine(), Some("m7g.xlarge"));
919        assert!(pool
920            .machines
921            .iter()
922            .any(|option| option.machine == "m7g.xlarge"));
923    }
924
925    #[test]
926    fn nested_x86_fixed_range_pool_preserves_bounds_and_rejects_graviton() {
927        let daemon = Daemon::new("vm-runtime-loader".to_string())
928            .code(DaemonCode::Image {
929                image: "example.com/vm-runtime:latest".to_string(),
930            })
931            .cluster("vm-runtime".to_string())
932            .cpu(ResourceSpec {
933                min: "2".to_string(),
934                desired: "2".to_string(),
935            })
936            .memory(ResourceSpec {
937                min: "4Gi".to_string(),
938                desired: "4Gi".to_string(),
939            })
940            .permissions("loader".to_string())
941            .build();
942        let cluster = ComputeCluster::new("vm-runtime".to_string())
943            .capacity_group(CapacityGroup {
944                group_id: "general".to_string(),
945                instance_type: None,
946                profile: Some(MachineProfile {
947                    cpu: "4".to_string(),
948                    memory_bytes: 16 * 1024 * 1024 * 1024,
949                    ephemeral_storage_bytes: 20 * 1024 * 1024 * 1024,
950                    architecture: Some(Architecture::X86_64),
951                    gpu: None,
952                }),
953                min_size: 2,
954                max_size: 2,
955                scale_policy: Some(CapacityGroupScalePolicy::Fixed {
956                    machines: ComputeChoiceRange {
957                        min: 1,
958                        max: 5,
959                        default: 2,
960                    },
961                }),
962                nested_virtualization: Some(true),
963            })
964            .build();
965        let stack = Stack {
966            id: "vm-runtime".to_string(),
967            resources: [
968                (
969                    "vm-runtime-loader".to_string(),
970                    ResourceEntry {
971                        config: Resource::new(daemon),
972                        lifecycle: ResourceLifecycle::Live,
973                        dependencies: Vec::new(),
974                        remote_access: false,
975                        enabled_when: None,
976                    },
977                ),
978                (
979                    "vm-runtime".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            .into_iter()
990            .collect(),
991            permissions: crate::permissions::PermissionsConfig::default(),
992            supported_platforms: None,
993            inputs: vec![],
994        };
995
996        let plan = plan_compute(&stack, Platform::Aws, None).expect("plan should build");
997        let pool = plan.pools.first().expect("general pool should exist");
998        assert_eq!(pool.recommended.machine(), Some("m8i.2xlarge"));
999        assert_eq!(pool.recommended.min_size(), 2);
1000        assert_eq!(pool.recommended.max_size(), 2);
1001        assert_eq!(
1002            pool.scale,
1003            CapacityGroupScalePolicy::Fixed {
1004                machines: ComputeChoiceRange {
1005                    min: 1,
1006                    max: 5,
1007                    default: 2,
1008                },
1009            }
1010        );
1011        assert!(!pool
1012            .machines
1013            .iter()
1014            .any(|option| option.machine == "m7g.2xlarge"));
1015
1016        let invalid_settings = ComputeSettings {
1017            pools: [(
1018                "general".to_string(),
1019                ComputePoolSelection::Fixed {
1020                    machines: 2,
1021                    machine: Some("m7g.2xlarge".to_string()),
1022                    failure_domains: None,
1023                },
1024            )]
1025            .into_iter()
1026            .collect(),
1027        };
1028        let invalid_plan = plan_compute(&stack, Platform::Aws, Some(&invalid_settings))
1029            .expect("plan should build");
1030        assert!(!invalid_plan.pools[0].errors.is_empty());
1031    }
1032
1033    #[test]
1034    fn local_plan_has_no_provider_machine_choices() {
1035        let stack = stack_with_container();
1036
1037        let plan = plan_compute(&stack, Platform::Local, None).expect("plan should build");
1038
1039        let pool = plan.pools.first().expect("general pool should exist");
1040        assert_eq!(pool.selected.machine(), None);
1041        assert!(pool.machines.is_empty());
1042        assert!(pool.errors.is_empty());
1043    }
1044}