|
|
@@ -4,6 +4,7 @@ pub use crate::pre_serve::{
|
|
|
use anyhow::{Context, Result};
|
|
|
use serde::Deserialize;
|
|
|
use std::collections::HashMap;
|
|
|
+use std::collections::HashSet;
|
|
|
use tracing::info;
|
|
|
#[derive(Clone, Deserialize)]
|
|
|
pub struct Node {
|
|
|
@@ -160,11 +161,12 @@ fn default_tensor_parallel_size() -> u32 {
|
|
|
1
|
|
|
}
|
|
|
/// Selects a cluster + NCCL config by name or auto-picks the smallest
|
|
|
-/// cluster with NCCL config satisfying `required_nodes`.
|
|
|
+/// cluster with NCCL config and enough free nodes.
|
|
|
pub fn select_cluster<'c>(
|
|
|
config: &'c Config,
|
|
|
cluster_name: &str,
|
|
|
required_nodes: Option<u32>,
|
|
|
+ occupied: &HashSet<String>,
|
|
|
) -> Result<(Cluster, &'c NcclConfig)> {
|
|
|
if cluster_name != "auto" {
|
|
|
let cluster = resolve_cluster_name(&config.clusters, cluster_name)?;
|
|
|
@@ -174,19 +176,26 @@ pub fn select_cluster<'c>(
|
|
|
.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| (c.nodes.len() as u32) >= r))
|
|
|
+ .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" => cluster.nodes.len(),
|
|
|
+ "nodes" => avail,
|
|
|
})
|
|
|
);
|
|
|
return Ok((cluster.clone(), nccl));
|
|
|
@@ -202,6 +211,7 @@ pub fn resolve_cluster<'c>(
|
|
|
config: &'c Config,
|
|
|
recipe: &Recipe,
|
|
|
cluster_name: &str,
|
|
|
+ occupied: &HashSet<String>,
|
|
|
) -> Result<(Cluster, &'c NcclConfig)> {
|
|
|
let choice = if cluster_name != "auto" {
|
|
|
cluster_name.to_string()
|
|
|
@@ -210,7 +220,7 @@ pub fn resolve_cluster<'c>(
|
|
|
} else {
|
|
|
"auto".to_string()
|
|
|
};
|
|
|
- select_cluster(config, &choice, recipe.required_nodes)
|
|
|
+ select_cluster(config, &choice, recipe.required_nodes, occupied)
|
|
|
}
|
|
|
/// Resolve the path to a recipe's `mods/` directory, if it exists.
|
|
|
/// Checks cwd first, then the default config dir.
|
|
|
@@ -299,16 +309,14 @@ mod tests {
|
|
|
}
|
|
|
#[test]
|
|
|
fn test_recipe_deserialize_minimal() {
|
|
|
- let toml_str = "\nname = \"Minimal\"\nmodel = \"test/model\"\n";
|
|
|
- let recipe: Recipe = toml::from_str(toml_str).unwrap();
|
|
|
+ let recipe: Recipe =
|
|
|
+ toml::from_str("\nname = \"Minimal\"\nmodel = \"test/model\"\n").unwrap();
|
|
|
assert_eq!(recipe.name, "Minimal");
|
|
|
assert_eq!(recipe.model, "test/model");
|
|
|
assert_eq!(recipe.engine, "vllm");
|
|
|
assert_eq!(recipe.port, 30000);
|
|
|
- assert!(recipe.args.to_vec().is_empty());
|
|
|
- assert!(recipe.env.is_empty());
|
|
|
+ assert!(recipe.args.to_vec().is_empty() && recipe.env.is_empty());
|
|
|
assert_eq!(recipe.distributed_backend, None);
|
|
|
- assert_eq!(recipe.tensor_parallel_size, 1);
|
|
|
}
|
|
|
#[test]
|
|
|
fn test_recipe_deserialize_full() {
|
|
|
@@ -400,7 +408,8 @@ ib_disable = 0
|
|
|
global_mods_dir: None,
|
|
|
};
|
|
|
config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
- let (cluster, nccl) = select_cluster(&config, "test-cluster", None).unwrap();
|
|
|
+ let (cluster, nccl) =
|
|
|
+ select_cluster(&config, "test-cluster", None, &HashSet::new()).unwrap();
|
|
|
assert_eq!(cluster.name, "test-cluster");
|
|
|
assert_eq!(nccl.ib_hca, "rocep1s0f0");
|
|
|
}
|
|
|
@@ -419,21 +428,21 @@ ib_disable = 0
|
|
|
global_mods_dir: None,
|
|
|
};
|
|
|
config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
- let (cluster, nccl) = select_cluster(&config, "auto", None).unwrap();
|
|
|
+ let (cluster, nccl) = select_cluster(&config, "auto", None, &HashSet::new()).unwrap();
|
|
|
assert_eq!(cluster.name, "test-cluster");
|
|
|
assert_eq!(nccl.ib_hca, "rocep1s0f0");
|
|
|
}
|
|
|
#[test]
|
|
|
fn test_select_cluster_auto_no_match() {
|
|
|
let config = test_config();
|
|
|
- let result = select_cluster(&config, "auto", None);
|
|
|
+ let result = select_cluster(&config, "auto", None, &HashSet::new());
|
|
|
assert!(result.is_err());
|
|
|
}
|
|
|
#[test]
|
|
|
fn test_select_cluster_not_found() {
|
|
|
let mut config = test_config();
|
|
|
config.clusters = vec![];
|
|
|
- let result = select_cluster(&config, "nonexistent", None);
|
|
|
+ let result = select_cluster(&config, "nonexistent", None, &HashSet::new());
|
|
|
assert!(result.is_err());
|
|
|
}
|
|
|
#[test]
|
|
|
@@ -456,12 +465,10 @@ ib_disable = 0
|
|
|
fn test_select_cluster_by_prefix() {
|
|
|
let mut config = test_config();
|
|
|
config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
- // Unique prefix works
|
|
|
let cluster = select_cluster_by_name(&config, "test-c").unwrap();
|
|
|
assert_eq!(cluster.name, "test-cluster");
|
|
|
- // No match errors
|
|
|
- let result = select_cluster_by_name(&config, "no-such");
|
|
|
- assert!(result.is_err());
|
|
|
+ let result = select_cluster_by_name(&config, "no-such").is_err();
|
|
|
+ assert!(result);
|
|
|
}
|
|
|
#[test]
|
|
|
fn test_russian_pluralization() {
|