|
|
@@ -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]
|