|
|
@@ -1,5 +1,6 @@
|
|
|
-use anyhow::Result;
|
|
|
-use std::collections::HashSet;
|
|
|
+use anyhow::{Context, Result};
|
|
|
+use serde::Deserialize;
|
|
|
+use std::collections::{HashMap, HashSet};
|
|
|
use tracing::info;
|
|
|
|
|
|
mod types;
|
|
|
@@ -53,6 +54,55 @@ pub fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&
|
|
|
}
|
|
|
}
|
|
|
|
|
|
+/// A `cluster_sources` file: extra `[[clusters]]` entries, optionally with
|
|
|
+/// per-cluster `[nccl.<name>]` / `[before_serve.<name>]` tables. Merged into
|
|
|
+/// the main [`Config`] by [`merge_cluster_sources`].
|
|
|
+#[derive(Deserialize, Default)]
|
|
|
+pub struct ClusterSource {
|
|
|
+ #[serde(default)]
|
|
|
+ pub clusters: Vec<Cluster>,
|
|
|
+ #[serde(default)]
|
|
|
+ pub nccl: HashMap<String, NcclConfig>,
|
|
|
+ #[serde(default)]
|
|
|
+ pub before_serve: HashMap<String, BeforeServeTable>,
|
|
|
+}
|
|
|
+
|
|
|
+/// Merge each `Config.cluster_sources` file into `config`: extra clusters
|
|
|
+/// (main-config entries win on name conflict; among sources, first wins),
|
|
|
+/// plus per-cluster `nccl` / `before_serve` tables.
|
|
|
+pub async fn merge_cluster_sources(config: &mut Config) -> Result<()> {
|
|
|
+ for raw in &config.cluster_sources {
|
|
|
+ let path = local_path(raw);
|
|
|
+ let content = tokio::fs::read_to_string(&path)
|
|
|
+ .await
|
|
|
+ .with_context(|| format!("cluster source '{}'", path.display()))?;
|
|
|
+ let src: ClusterSource = toml::from_str(&content)
|
|
|
+ .with_context(|| format!("invalid cluster source '{}'", path.display()))?;
|
|
|
+ for c in src.clusters {
|
|
|
+ if !config.clusters.iter().any(|x| x.name == c.name) {
|
|
|
+ config.clusters.push(c);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ for (name, v) in src.nccl {
|
|
|
+ config.nccl.entry(name).or_insert(v);
|
|
|
+ }
|
|
|
+ for (name, v) in src.before_serve {
|
|
|
+ config.before_serve.entry(name).or_insert(v);
|
|
|
+ }
|
|
|
+ }
|
|
|
+ Ok(())
|
|
|
+}
|
|
|
+
|
|
|
+/// Expand `~/...` against the local user's home (unlike `expand_tilde`,
|
|
|
+/// which targets the ssh user's home *on the node*). Absolute paths pass
|
|
|
+/// through.
|
|
|
+fn local_path(s: &str) -> std::path::PathBuf {
|
|
|
+ if let (Some(rest), Some(home)) = (s.strip_prefix("~/"), dirs::home_dir()) {
|
|
|
+ return home.join(rest);
|
|
|
+ }
|
|
|
+ std::path::PathBuf::from(s)
|
|
|
+}
|
|
|
+
|
|
|
/// Whether NCCL configuration is required for cluster selection.
|
|
|
///
|
|
|
/// NCCL config is only needed for multi-node distributed runs — for
|
|
|
@@ -172,11 +222,20 @@ pub async fn resolve_mods_dir(name: &str) -> Option<std::path::PathBuf> {
|
|
|
cfg.is_dir().then_some(cfg)
|
|
|
}
|
|
|
|
|
|
-pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
|
|
|
+/// Load a recipe by name. Search order: `./recipes/<name>/recipe.toml` (cwd),
|
|
|
+/// then each `sources` entry (a dir containing `<name>/recipe.toml`, `~`
|
|
|
+/// expanded), then `<default_config_dir>/recipes/<name>/recipe.toml`.
|
|
|
+pub async fn load_recipe(recipe_name: &str, sources: &[String]) -> Result<Recipe> {
|
|
|
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)?);
|
|
|
}
|
|
|
+ for src in sources {
|
|
|
+ let p = local_path(src).join(recipe_name).join("recipe.toml");
|
|
|
+ if let Ok(content) = tokio::fs::read_to_string(&p).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 {
|