Jelajahi Sumber

flare: translate our cluster config into FLARE's, per cluster

One flat [flare] block only fits the machine it was written for: running the
same recipe on dgx1-4 needed a different orbit cluster, different peer names, a
different controller and a different local package dir. [flare.clusters.<our
cluster>] now supplies {orbit_cluster, orbit_nodes, controller_node, local_dir}
and wins over the flat fields, so 'serve --cluster 4nodes-switch' launches on
orbit 'prod4' with gx10-dgx1..4 and dgx1 as controller, while dgx5-8 keeps using
orbit 'sparks'.

Verified by dry-run: dgx58-switch → orbit run 'gx10-dgx5:...' --cluster
'sparks'; 4nodes-switch → 'gx10-dgx1:...' --cluster 'prod4'. Test covers the
mapping precedence.
Alexander Korolev 19 jam lalu
induk
melakukan
de8b73b131
6 mengubah file dengan 125 tambahan dan 40 penghapusan
  1. 1 0
      src/config.rs
  2. 21 0
      src/config/flare_cluster.rs
  3. 9 5
      src/config/types.rs
  4. 2 1
      src/handlers.rs
  5. 84 29
      src/native/flare.rs
  6. 8 5
      src/web/flare_nodes.rs

+ 1 - 0
src/config.rs

@@ -56,6 +56,7 @@ mod staging;
 pub use staging::{Staging, StagingFormat, StagingKind, StagingVerify};
 mod types;
 pub use types::*;
+pub(crate) mod flare_cluster;
 mod hf_cache;
 
 impl Node {

+ 21 - 0
src/config/flare_cluster.rs

@@ -0,0 +1,21 @@
+//! `[flare.clusters.<our-cluster>]` — the translation layer between our
+//! cluster definitions and FLARE's own orbit naming, so one launcher config
+//! can serve FLARE on any cluster (dgx58-switch → orbit `sparks`,
+//! prod4 → orbit `prod4`) with the right node list and package mirror dir.
+
+/// One cluster's FLARE mapping: our cluster name → orbit cluster + nodes.
+#[derive(Clone, Debug, Default, serde::Deserialize)]
+pub struct FlareCluster {
+    /// `[cluster.<name>]` in orbit.toml for this cluster.
+    #[serde(default)]
+    pub orbit_cluster: Option<String>,
+    /// Orbit peer names in our rank order (`gx10-dgx1`…).
+    #[serde(default)]
+    pub orbit_nodes: Vec<String>,
+    /// Controller override for this cluster (our node name).
+    #[serde(default)]
+    pub controller_node: Option<String>,
+    /// Local package mirror dir for this cluster's controller.
+    #[serde(default)]
+    pub local_dir: Option<String>,
+}

+ 9 - 5
src/config/types.rs

@@ -476,11 +476,15 @@ pub struct FlareSettings {
     /// multi-cluster config (the recipes panel then sees no FLARE groups).
     #[serde(default)]
     pub controller_node: Option<String>,
-    /// Local mirror dir for the package on the controller (e.g.
-    /// `~/flare-pkgs`). When set, `serve` copies the package's companion files
-    /// there (they are named by sha256, so a size match means "already there")
-    /// and launches from the local copy — runtime table reads then hit local
-    /// NVMe instead of the shared mount.
+    /// Per-cluster FLARE settings, keyed by OUR cluster name (`dgx58-switch`,
+    /// `prod4`). This is the translation layer: the same launcher config then
+    /// serves FLARE on any cluster, with the right orbit cluster name, node
+    /// list and local package dir — instead of one flat block that only fits
+    /// the machine it was written for.
+    #[serde(default)]
+    pub clusters: std::collections::HashMap<String, super::flare_cluster::FlareCluster>,
+    /// Local package mirror dir on the controller (see `native::flare`):
+    /// `serve` copies the package there and launches from the local copy.
     #[serde(default)]
     pub local_dir: Option<String>,
     /// `[cluster.<name>]` in `~/.config/flare/orbit.toml` (default `sparks`).

+ 2 - 1
src/handlers.rs

@@ -33,7 +33,8 @@ pub(crate) async fn handle_logs(
             .iter()
             .find(|n| n.rank == 0)
             .ok_or_else(|| anyhow::anyhow!("no head node (rank 0) in cluster"))?;
-        let cmd = crate::native::flare::attach_cmd(config, &recipe_obj)?;
+        let orbit = crate::native::flare::orbit_cluster_for(config, &cluster.name);
+        let cmd = crate::native::flare::attach_cmd(config, &recipe_obj, &orbit)?;
         crate::ssh::run_ssh(config, head, cmd, dry_run).await?;
         return Ok(());
     }

+ 84 - 29
src/native/flare.rs

@@ -38,21 +38,37 @@ pub fn http_port(recipe: &Recipe) -> u16 {
         .unwrap_or(30000)
 }
 
-fn orbit_cluster(config: &Config) -> String {
-    config
-        .flare
-        .as_ref()
-        .and_then(|f| f.orbit_cluster.clone())
+/// Orbit cluster name for one of OUR clusters: the per-cluster mapping wins,
+/// then the flat `[flare].orbit_cluster`, then the default. This is what makes
+/// one launcher config work on `dgx58-switch` and `prod4` alike.
+pub fn orbit_cluster_for(config: &Config, cluster: &str) -> String {
+    let f = config.flare.as_ref();
+    f.and_then(|f| f.clusters.get(cluster))
+        .and_then(|c| c.orbit_cluster.clone())
+        .or_else(|| f.and_then(|f| f.orbit_cluster.clone()))
         .unwrap_or_else(|| "sparks".to_string())
 }
 
+/// Local package mirror dir for a cluster (per-cluster map first).
+fn local_dir_for(config: &Config, cluster: &str) -> Option<String> {
+    let f = config.flare.as_ref()?;
+    f.clusters
+        .get(cluster)
+        .and_then(|c| c.local_dir.clone())
+        .or_else(|| f.local_dir.clone())
+}
+
 /// The orbit-side name for one of our nodes (rank order → `orbit_nodes`,
 /// falling back to the cluster's own name).
-fn orbit_node_name(config: &Config, node: &Node) -> String {
-    config
-        .flare
-        .as_ref()
-        .and_then(|f| f.orbit_nodes.get(node.rank as usize).cloned())
+fn orbit_node_name(config: &Config, cluster: &str, node: &Node) -> String {
+    let f = config.flare.as_ref();
+    let per_cluster = f
+        .and_then(|f| f.clusters.get(cluster))
+        .map(|c| c.orbit_nodes.as_slice())
+        .filter(|v| !v.is_empty());
+    per_cluster
+        .or_else(|| f.map(|f| f.orbit_nodes.as_slice()))
+        .and_then(|v| v.get(node.rank as usize).cloned())
         .unwrap_or_else(|| node.name.clone())
 }
 
@@ -84,6 +100,7 @@ pub fn download_cmd(config: &Config, recipe: &Recipe) -> Result<String> {
 pub fn run_cmd(
     config: &Config,
     recipe: &Recipe,
+    orbit_cluster: &str,
     controller: &str,
     package: &str,
 ) -> Result<String> {
@@ -96,7 +113,7 @@ pub fn run_cmd(
         "{} orbit run {} --cluster {} --name {} --host 0.0.0.0 --port {} --detach",
         binary(config),
         shq(&format!("{controller}:{package}")),
-        shq(&orbit_cluster(config)),
+        shq(orbit_cluster),
         shq(&model_name),
         http_port(recipe),
     );
@@ -117,7 +134,7 @@ pub fn run_cmd(
 }
 
 /// `flare orbit stop <name> --cluster <c>`.
-pub fn stop_cmd(config: &Config, recipe: &Recipe) -> Result<String> {
+pub fn stop_cmd(config: &Config, recipe: &Recipe, orbit_cluster: &str) -> Result<String> {
     let f = recipe
         .flare
         .as_ref()
@@ -127,12 +144,12 @@ pub fn stop_cmd(config: &Config, recipe: &Recipe) -> Result<String> {
         "{} orbit stop {} --cluster {}",
         binary(config),
         shq(&model_name),
-        shq(&orbit_cluster(config)),
+        shq(orbit_cluster),
     ))
 }
 
 /// `flare orbit attach <name> --cluster <c>` (log tail; the CLI's `logs`).
-pub fn attach_cmd(config: &Config, recipe: &Recipe) -> Result<String> {
+pub fn attach_cmd(config: &Config, recipe: &Recipe, orbit_cluster: &str) -> Result<String> {
     let f = recipe
         .flare
         .as_ref()
@@ -142,7 +159,7 @@ pub fn attach_cmd(config: &Config, recipe: &Recipe) -> Result<String> {
         "{} orbit attach {} --cluster {}",
         binary(config),
         shq(&model_name),
-        shq(&orbit_cluster(config)),
+        shq(orbit_cluster),
     ))
 }
 
@@ -153,7 +170,7 @@ pub fn attach_cmd(config: &Config, recipe: &Recipe) -> Result<String> {
 /// Only the controller has the orbit file with `[cluster.<name>]` — a peer
 /// just needs the binary and a live agent, so the cluster check is added
 /// there and nowhere else.
-pub fn preflight_cmd(config: &Config, controller: bool) -> String {
+pub fn preflight_cmd(config: &Config, orbit_cluster: &str, controller: bool) -> String {
     let base = format!(
         "{} --version >/dev/null && systemctl --user is-active flare-agent >/dev/null",
         binary(config)
@@ -162,7 +179,7 @@ pub fn preflight_cmd(config: &Config, controller: bool) -> String {
         format!(
             "{base} && {} orbit status --cluster {} >/dev/null",
             binary(config),
-            shq(&orbit_cluster(config)),
+            shq(orbit_cluster),
         )
     } else {
         base
@@ -221,15 +238,16 @@ pub async fn serve(
     dry_run: bool,
 ) -> Result<()> {
     let controller = head(cluster)?;
+    let orbit = orbit_cluster_for(config, &cluster.name);
     if dry_run {
         println!(
             "  [dry-run] preflight on {} nodes (cluster check on the controller only): {}",
             cluster.nodes.len(),
-            preflight_cmd(config, true)
+            preflight_cmd(config, &orbit, true)
         );
     } else {
         for node in &cluster.nodes {
-            let pre = preflight_cmd(config, node.rank == 0);
+            let pre = preflight_cmd(config, &orbit, node.rank == 0);
             if let Err(e) = crate::ssh::run_ssh(config, node, pre, false).await {
                 bail!(
                     "FLARE preflight failed on {}: {e} — install `flare`, then `flare agent install` and pair the node with the controller",
@@ -252,7 +270,7 @@ pub async fn serve(
     // engram rows are read at request time, so a shared mount is the wrong
     // place for the package even with a RAM row cache.
     let mut package = package;
-    if let Some(dir) = config.flare.as_ref().and_then(|f| f.local_dir.clone()) {
+    if let Some(dir) = local_dir_for(config, &cluster.name) {
         let dir = dir.replace('~', "$HOME");
         let cmd = mirror_cmd(&dir, &package);
         let out = crate::ssh::run_ssh(config, controller, cmd, dry_run).await?;
@@ -264,8 +282,8 @@ pub async fn serve(
             package = local;
         }
     }
-    let orbit_name = orbit_node_name(config, controller);
-    let run = run_cmd(config, recipe, &orbit_name, &package)?;
+    let orbit_name = orbit_node_name(config, &cluster.name, controller);
+    let run = run_cmd(config, recipe, &orbit, &orbit_name, &package)?;
     crate::ssh::run_ssh(config, controller, run.clone(), dry_run).await?;
     if dry_run {
         println!("  [dry-run] {}: {run}", controller.name);
@@ -311,7 +329,7 @@ pub async fn stop(
     dry_run: bool,
 ) -> Result<()> {
     let controller = head(cluster)?;
-    let cmd = stop_cmd(config, recipe)?;
+    let cmd = stop_cmd(config, recipe, &orbit_cluster_for(config, &cluster.name))?;
     crate::ssh::run_ssh(config, controller, cmd, dry_run).await?;
     Ok(())
 }
@@ -358,7 +376,14 @@ mod tests {
 
     #[test]
     fn run_passes_flags_verbatim_and_the_controller_package() {
-        let cmd = run_cmd(&cfg(), &recipe(), "dgx5", "/pkgs/dsv41.safetensors").unwrap();
+        let cmd = run_cmd(
+            &cfg(),
+            &recipe(),
+            "sparks",
+            "dgx5",
+            "/pkgs/dsv41.safetensors",
+        )
+        .unwrap();
         assert!(
             cmd.contains("orbit run 'dgx5:/pkgs/dsv41.safetensors'"),
             "{cmd}"
@@ -387,25 +412,55 @@ mod tests {
             host_key: None,
             entity: None,
         };
-        assert_eq!(orbit_node_name(&c, &node), "spark1");
+        assert_eq!(orbit_node_name(&c, "dgx58-switch", &node), "spark1");
         let node1 = Node { rank: 1, ..node };
         // no entry for rank 1 → the cluster's own name is used
-        assert_eq!(orbit_node_name(&c, &node1), "dgx5.lan");
+        assert_eq!(orbit_node_name(&c, "dgx58-switch", &node1), "dgx5.lan");
+    }
+
+    #[test]
+    fn per_cluster_mapping_wins_over_the_flat_block() {
+        // The same launcher config must serve FLARE on another cluster: the
+        // per-cluster map supplies that cluster's orbit name and node list.
+        let mut c = cfg();
+        c.flare.as_mut().unwrap().clusters.insert(
+            "prod4".into(),
+            crate::config::flare_cluster::FlareCluster {
+                orbit_cluster: Some("prod4".into()),
+                orbit_nodes: vec!["gx10-dgx1".into()],
+                controller_node: None,
+                local_dir: Some("/data/flare".into()),
+            },
+        );
+        assert_eq!(orbit_cluster_for(&c, "prod4"), "prod4");
+        assert_eq!(orbit_cluster_for(&c, "dgx58-switch"), "sparks");
+        assert_eq!(local_dir_for(&c, "prod4").as_deref(), Some("/data/flare"));
+        let node = Node {
+            name: "dgx1.lan".into(),
+            ssh_ip: "10.0.0.1".into(),
+            ib_ip: None,
+            rank: 0,
+            host_key: None,
+            entity: None,
+        };
+        assert_eq!(orbit_node_name(&c, "prod4", &node), "gx10-dgx1");
+        let cmd = stop_cmd(&c, &recipe(), &orbit_cluster_for(&c, "prod4")).unwrap();
+        assert!(cmd.contains("--cluster 'prod4'"), "{cmd}");
     }
 
     #[test]
     fn stop_names_the_group_and_cluster() {
-        let cmd = stop_cmd(&cfg(), &recipe()).unwrap();
+        let cmd = stop_cmd(&cfg(), &recipe(), "sparks").unwrap();
         assert_eq!(cmd, "flare orbit stop 'deepseek1017' --cluster 'sparks'");
     }
 
     #[test]
     fn preflight_checks_agent_and_pairing() {
-        let cmd = preflight_cmd(&cfg(), true);
+        let cmd = preflight_cmd(&cfg(), "sparks", true);
         assert!(cmd.contains("is-active flare-agent"), "{cmd}");
         assert!(cmd.contains("orbit status --cluster 'sparks'"), "{cmd}");
         // a peer has no orbit file: binary + agent only
-        assert!(!preflight_cmd(&cfg(), false).contains("orbit status"));
+        assert!(!preflight_cmd(&cfg(), "sparks", false).contains("orbit status"));
     }
 
     #[test]

+ 8 - 5
src/web/flare_nodes.rs

@@ -7,12 +7,15 @@
 
 use crate::config::{Config, Node};
 
-/// `[flare].controller_node`, when the operator pinned it.
+/// The pinned controller: a per-cluster override wins (so one config can pin
+/// a different controller on each cluster), then the flat `[flare]` field.
 pub fn pinned(config: &Config) -> Option<&Node> {
-    let want = config
-        .flare
-        .as_ref()
-        .and_then(|f| f.controller_node.clone())?;
+    let f = config.flare.as_ref()?;
+    let want = f
+        .clusters
+        .values()
+        .find_map(|c| c.controller_node.clone())
+        .or_else(|| f.controller_node.clone())?;
     find(config, &want)
 }