use anyhow::{Context, Result}; use std::collections::HashSet; use tracing::info; mod types; pub use types::*; pub fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&'a Cluster> { if let Some(c) = clusters.iter().find(|c| c.name == name) { return Ok(c); } let matches: Vec<&Cluster> = clusters .iter() .filter(|c| c.name.starts_with(name)) .collect(); match matches.len() { 0 => anyhow::bail!(fluent_i18n::t!("cluster-not-found", { "name" => name.to_string() })), 1 => Ok(matches[0]), n => anyhow::bail!( "Pref '{}' matches {} clusters: {}", name, n, matches .iter() .map(|c| c.name.as_str()) .collect::>() .join(", ") ), } } pub fn select_cluster<'c>( config: &'c Config, cluster_name: &str, required_nodes: Option, occupied: &HashSet, ) -> Result<(Cluster, &'c NcclConfig)> { if cluster_name != "auto" { let cluster = resolve_cluster_name(&config.clusters, cluster_name)?; let nccl = config .nccl .get(&cluster.name) .context(fluent_i18n::t!("nccl-config-missing"))?; return Ok((cluster.clone(), nccl)); } let free = |c: &Cluster| -> u32 { c.nodes .iter() .filter(|n| !occupied.contains(&n.name)) .count() as u32 }; let mut candidates: Vec<(&Cluster, &NcclConfig)> = config .clusters .iter() .filter_map(|c| Some((c, config.nccl.get(&c.name)?))) .filter(|(c, _)| required_nodes.is_none_or(|r| free(c) >= r)) .collect(); candidates.sort_by_key(|(c, _)| c.nodes.len()); if let Some((cluster, nccl)) = candidates.into_iter().next() { let avail = free(cluster); info!( "{}", fluent_i18n::t!("cluster-selected", { "name" => cluster.name.to_string(), "nodes" => avail }) ); return Ok((cluster.clone(), nccl)); } anyhow::bail!(fluent_i18n::t!("cluster-auto-fail")); } pub fn select_cluster_by_name(config: &Config, cluster_name: &str) -> Result { Ok(resolve_cluster_name(&config.clusters, cluster_name)?.clone()) } pub fn resolve_cluster<'c>( config: &'c Config, recipe: &Recipe, cluster_name: &str, occupied: &HashSet, ) -> Result<(Cluster, &'c NcclConfig)> { let choice = if cluster_name != "auto" { cluster_name.to_string() } else if let Some(ref rc) = recipe.cluster { rc.clone() } else { "auto".to_string() }; select_cluster(config, &choice, recipe.required_nodes, occupied) } pub async fn resolve_mods_dir(name: &str) -> Option { let cwd = std::path::PathBuf::from(format!("recipes/{name}/mods")); if cwd.is_dir() { return Some(cwd); } let cfg = default_config_dir()? .join("recipes") .join(name) .join("mods"); cfg.is_dir().then_some(cfg) } pub async fn load_recipe(recipe_name: &str) -> Result { let cwd_path = format!("recipes/{}/recipe.toml", recipe_name); if let Ok(content) = tokio::fs::read_to_string(&cwd_path).await { return Ok(toml::from_str(&content)?); } if let Some(dir) = default_config_dir() { let in_default = dir.join("recipes").join(recipe_name).join("recipe.toml"); if let Ok(content) = tokio::fs::read_to_string(&in_default).await { return Ok(toml::from_str(&content)?); } } let content = tokio::fs::read_to_string(&cwd_path).await?; Ok(toml::from_str(&content)?) } pub fn default_config_dir() -> Option { dirs::config_dir().map(|p| p.join("sm121")) } #[cfg(test)] mod tests;