|
|
@@ -1,139 +1,11 @@
|
|
|
-pub use crate::pre_serve::{
|
|
|
- BeforeServeTable, DistributedBackend, ModsTable, OnFailure, PreServeCommand, RecipeArgs,
|
|
|
-};
|
|
|
use anyhow::{Context, Result};
|
|
|
-use serde::Deserialize;
|
|
|
-use std::collections::{HashMap, HashSet};
|
|
|
+use std::collections::HashSet;
|
|
|
use tracing::info;
|
|
|
-#[derive(Clone, Debug, Deserialize)]
|
|
|
-pub struct Node {
|
|
|
- pub name: String,
|
|
|
- pub ssh_ip: String,
|
|
|
- pub ib_ip: String,
|
|
|
- pub rank: u32,
|
|
|
- #[serde(default)]
|
|
|
- pub host_key: Option<String>,
|
|
|
- #[serde(default)]
|
|
|
- pub entity: Option<String>,
|
|
|
-}
|
|
|
-#[derive(Clone, Debug, Deserialize)]
|
|
|
-pub struct NcclConfig {
|
|
|
- pub ib_hca: String,
|
|
|
- pub socket_ifname: String,
|
|
|
- pub ib_gid_index: u32,
|
|
|
- pub ib_disable: u32,
|
|
|
- #[serde(default)]
|
|
|
- pub gloo_socket_ifname: Option<String>,
|
|
|
- #[serde(default)]
|
|
|
- pub ucx_net_devices: Option<String>,
|
|
|
- #[serde(default)]
|
|
|
- pub nccl_cross_nic: Option<u32>,
|
|
|
- #[serde(default)]
|
|
|
- pub nccl_proto: Option<String>,
|
|
|
- #[serde(default)]
|
|
|
- pub nccl_p2p_level: Option<String>,
|
|
|
-}
|
|
|
-#[derive(Clone, Debug, Deserialize)]
|
|
|
-pub struct Cluster {
|
|
|
- pub name: String,
|
|
|
- pub master_ib_ip: String,
|
|
|
- pub nodes: Vec<Node>,
|
|
|
- /// Extra HA switches.
|
|
|
- #[serde(default)]
|
|
|
- pub extra_entities: Vec<String>,
|
|
|
-}
|
|
|
-#[derive(Clone, Deserialize)]
|
|
|
-pub struct Config {
|
|
|
- pub ssh_user: String,
|
|
|
- pub ssh_key_path: std::path::PathBuf,
|
|
|
- pub clusters: Vec<Cluster>,
|
|
|
- #[serde(default)]
|
|
|
- pub nccl: HashMap<String, NcclConfig>,
|
|
|
- #[serde(default)]
|
|
|
- pub hf_cache_dir: Option<String>,
|
|
|
- /// Init port
|
|
|
- #[serde(default = "default_init_port")]
|
|
|
- pub init_port: u16,
|
|
|
- /// Pre-serve, aliases, HA
|
|
|
- #[serde(default)]
|
|
|
- pub before_serve: HashMap<String, BeforeServeTable>,
|
|
|
- #[serde(default)]
|
|
|
- pub exec_aliases: HashMap<String, String>,
|
|
|
- #[serde(default)]
|
|
|
- pub home_assistant: Option<HomeAssistantConfig>,
|
|
|
- /// Global mods dir.
|
|
|
- #[serde(default)]
|
|
|
- pub global_mods_dir: Option<std::path::PathBuf>,
|
|
|
-}
|
|
|
-/// Home Assistant config.
|
|
|
-#[derive(Clone, Deserialize)]
|
|
|
-pub struct HomeAssistantConfig {
|
|
|
- pub url: String,
|
|
|
- pub token: String,
|
|
|
- #[serde(default = "default_shutdown_delay")]
|
|
|
- pub shutdown_delay_seconds: u64,
|
|
|
-}
|
|
|
-fn default_shutdown_delay() -> u64 {
|
|
|
- 30
|
|
|
-}
|
|
|
-fn default_init_port() -> u16 {
|
|
|
- 25000
|
|
|
-}
|
|
|
-fn default_port() -> u16 {
|
|
|
- 30000
|
|
|
-}
|
|
|
-impl Config {
|
|
|
- pub fn resolved_hf_cache_dir(&self) -> std::path::PathBuf {
|
|
|
- if let Some(ref dir) = self.hf_cache_dir {
|
|
|
- return std::path::PathBuf::from(dir);
|
|
|
- }
|
|
|
- let home = if self.ssh_user == "root" {
|
|
|
- "/root"
|
|
|
- } else {
|
|
|
- // Safe to return a temporary String; PathBuf::from copies it
|
|
|
- return std::path::PathBuf::from(format!("/home/{}/.cache/huggingface", self.ssh_user));
|
|
|
- };
|
|
|
- std::path::PathBuf::from(format!("{}/.cache/huggingface", home))
|
|
|
- }
|
|
|
-}
|
|
|
-#[derive(Clone, Deserialize)]
|
|
|
-pub struct Recipe {
|
|
|
- #[serde(default = "default_engine")]
|
|
|
- pub engine: String,
|
|
|
- pub name: String,
|
|
|
- #[serde(default)]
|
|
|
- pub docker_image: String,
|
|
|
- #[serde(default)]
|
|
|
- pub model: String,
|
|
|
- #[serde(default = "default_port")]
|
|
|
- pub port: u16,
|
|
|
- #[serde(default)]
|
|
|
- pub required_nodes: Option<u32>,
|
|
|
- #[serde(default)]
|
|
|
- pub pipeline_parallel_size: Option<u32>,
|
|
|
- #[serde(default = "default_tensor_parallel_size")]
|
|
|
- pub tensor_parallel_size: u32,
|
|
|
- #[serde(default)]
|
|
|
- pub args: RecipeArgs,
|
|
|
- #[serde(default)]
|
|
|
- pub env: HashMap<String, String>,
|
|
|
- #[serde(default)]
|
|
|
- pub cluster: Option<String>,
|
|
|
- #[serde(default)]
|
|
|
- pub entrypoint: Vec<String>,
|
|
|
- pub distributed_backend: Option<DistributedBackend>,
|
|
|
- /// Shm size
|
|
|
- #[serde(default = "default_shm_size")]
|
|
|
- pub shm_size: String,
|
|
|
- /// Recipe pre-serve.
|
|
|
- #[serde(default)]
|
|
|
- pub before_serve: Vec<PreServeCommand>,
|
|
|
- /// Mods before engine.
|
|
|
- #[serde(default)]
|
|
|
- pub mods: ModsTable,
|
|
|
-}
|
|
|
-/// Resolve cluster.
|
|
|
-pub(crate) fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> Result<&'a Cluster> {
|
|
|
+
|
|
|
+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);
|
|
|
}
|
|
|
@@ -156,16 +28,7 @@ pub(crate) fn resolve_cluster_name<'a>(clusters: &'a [Cluster], name: &str) -> R
|
|
|
),
|
|
|
}
|
|
|
}
|
|
|
-fn default_engine() -> String {
|
|
|
- "vllm".to_string()
|
|
|
-}
|
|
|
-fn default_shm_size() -> String {
|
|
|
- String::new()
|
|
|
-}
|
|
|
-fn default_tensor_parallel_size() -> u32 {
|
|
|
- 1
|
|
|
-}
|
|
|
-/// Select cluster.
|
|
|
+
|
|
|
pub fn select_cluster<'c>(
|
|
|
config: &'c Config,
|
|
|
cluster_name: &str,
|
|
|
@@ -197,18 +60,17 @@ pub fn select_cluster<'c>(
|
|
|
let avail = free(cluster);
|
|
|
info!(
|
|
|
"{}",
|
|
|
- fluent_i18n::t!("cluster-selected", {
|
|
|
- "name" => cluster.name.to_string(),
|
|
|
- "nodes" => avail,
|
|
|
- })
|
|
|
+ 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<Cluster> {
|
|
|
Ok(resolve_cluster_name(&config.clusters, cluster_name)?.clone())
|
|
|
}
|
|
|
+
|
|
|
pub fn resolve_cluster<'c>(
|
|
|
config: &'c Config,
|
|
|
recipe: &Recipe,
|
|
|
@@ -224,6 +86,7 @@ pub fn resolve_cluster<'c>(
|
|
|
};
|
|
|
select_cluster(config, &choice, recipe.required_nodes, occupied)
|
|
|
}
|
|
|
+
|
|
|
pub async fn resolve_mods_dir(name: &str) -> Option<std::path::PathBuf> {
|
|
|
let cwd = std::path::PathBuf::from(format!("recipes/{name}/mods"));
|
|
|
if cwd.is_dir() {
|
|
|
@@ -235,6 +98,7 @@ pub async fn resolve_mods_dir(name: &str) -> Option<std::path::PathBuf> {
|
|
|
.join("mods");
|
|
|
cfg.is_dir().then_some(cfg)
|
|
|
}
|
|
|
+
|
|
|
pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
|
|
|
let cwd_path = format!("recipes/{}/recipe.toml", recipe_name);
|
|
|
if let Ok(content) = tokio::fs::read_to_string(&cwd_path).await {
|
|
|
@@ -249,252 +113,10 @@ pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
|
|
|
let content = tokio::fs::read_to_string(&cwd_path).await?;
|
|
|
Ok(toml::from_str(&content)?)
|
|
|
}
|
|
|
+
|
|
|
pub fn default_config_dir() -> Option<std::path::PathBuf> {
|
|
|
dirs::config_dir().map(|p| p.join("sm121-launcher"))
|
|
|
}
|
|
|
+
|
|
|
#[cfg(test)]
|
|
|
-mod tests {
|
|
|
- use super::*;
|
|
|
- fn test_cluster() -> Cluster {
|
|
|
- Cluster {
|
|
|
- name: "test-cluster".to_string(),
|
|
|
- master_ib_ip: "192.168.1.1".to_string(),
|
|
|
- nodes: vec![
|
|
|
- Node {
|
|
|
- name: "node0".to_string(),
|
|
|
- ssh_ip: "10.0.0.1".to_string(),
|
|
|
- ib_ip: "192.168.1.1".to_string(),
|
|
|
- rank: 0,
|
|
|
- host_key: None,
|
|
|
- entity: None,
|
|
|
- },
|
|
|
- Node {
|
|
|
- name: "node1".to_string(),
|
|
|
- ssh_ip: "10.0.0.2".to_string(),
|
|
|
- ib_ip: "192.168.1.2".to_string(),
|
|
|
- rank: 1,
|
|
|
- host_key: None,
|
|
|
- entity: None,
|
|
|
- },
|
|
|
- ],
|
|
|
- extra_entities: vec![],
|
|
|
- }
|
|
|
- }
|
|
|
- fn test_config() -> Config {
|
|
|
- Config {
|
|
|
- ssh_user: "u".to_string(),
|
|
|
- ssh_key_path: "/k".into(),
|
|
|
- clusters: vec![test_cluster()],
|
|
|
- nccl: HashMap::new(),
|
|
|
- hf_cache_dir: None,
|
|
|
- init_port: 25000,
|
|
|
- before_serve: HashMap::new(),
|
|
|
- exec_aliases: HashMap::new(),
|
|
|
- home_assistant: None,
|
|
|
- global_mods_dir: None,
|
|
|
- }
|
|
|
- }
|
|
|
- fn test_nccl() -> NcclConfig {
|
|
|
- NcclConfig {
|
|
|
- ib_hca: "rocep1s0f0".to_string(),
|
|
|
- socket_ifname: "enp1s0f0np0".to_string(),
|
|
|
- ib_gid_index: 3,
|
|
|
- ib_disable: 0,
|
|
|
- gloo_socket_ifname: None,
|
|
|
- ucx_net_devices: None,
|
|
|
- nccl_cross_nic: None,
|
|
|
- nccl_proto: None,
|
|
|
- nccl_p2p_level: None,
|
|
|
- }
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_recipe_deserialize_minimal() {
|
|
|
- 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() && recipe.env.is_empty());
|
|
|
- assert_eq!(recipe.distributed_backend, None);
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_recipe_deserialize_full() {
|
|
|
- let toml_str = r#"
|
|
|
-engine = "sglang"
|
|
|
-name = "Full Recipe"
|
|
|
-docker_image = "org/image:tag"
|
|
|
-model = "org/model"
|
|
|
-port = 30000
|
|
|
-required_nodes = 4
|
|
|
-pipeline_parallel_size = 2
|
|
|
-distributed_backend = "ray"
|
|
|
-cluster = "my-cluster"
|
|
|
-args = ["--arg1", "val1", "--arg2"]
|
|
|
-[env]
|
|
|
-KEY = "value"
|
|
|
-OTHER = "123"
|
|
|
-"#;
|
|
|
- let recipe: Recipe = toml::from_str(toml_str).unwrap();
|
|
|
- assert_eq!(recipe.engine, "sglang");
|
|
|
- assert_eq!(recipe.name, "Full Recipe");
|
|
|
- assert_eq!(recipe.docker_image, "org/image:tag");
|
|
|
- assert_eq!(recipe.model, "org/model");
|
|
|
- assert_eq!(recipe.port, 30000);
|
|
|
- assert_eq!(recipe.required_nodes, Some(4));
|
|
|
- assert_eq!(recipe.pipeline_parallel_size, Some(2));
|
|
|
- assert_eq!(recipe.distributed_backend, Some(DistributedBackend::Ray));
|
|
|
- assert_eq!(recipe.cluster, Some("my-cluster".to_string()));
|
|
|
- assert_eq!(recipe.args.to_vec(), vec!["--arg1", "val1", "--arg2"]);
|
|
|
- assert_eq!(recipe.env.get("KEY").unwrap(), "value");
|
|
|
- assert_eq!(recipe.env.get("OTHER").unwrap(), "123");
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_recipe_deserialize_no_required_nodes() {
|
|
|
- let toml_str = r#"
|
|
|
-name = "Flexible"
|
|
|
-model = "test/model"
|
|
|
-"#;
|
|
|
- let recipe: Recipe = toml::from_str(toml_str).unwrap();
|
|
|
- assert_eq!(recipe.required_nodes, None);
|
|
|
- assert_eq!(recipe.pipeline_parallel_size, None);
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_config_deserialize() {
|
|
|
- let toml_str = r#"
|
|
|
-ssh_user = "admin"
|
|
|
-ssh_key_path = "/home/admin/.ssh/id_rsa"
|
|
|
-[[clusters]]
|
|
|
-name = "prod"
|
|
|
-master_ib_ip = "10.0.0.1"
|
|
|
-nodes = [
|
|
|
- { name = "node1", ssh_ip = "10.0.0.10", ib_ip = "10.0.0.1", rank = 0 },
|
|
|
- { name = "node2", ssh_ip = "10.0.0.11", ib_ip = "10.0.0.2", rank = 1 },
|
|
|
-]
|
|
|
-[nccl."prod"]
|
|
|
-ib_hca = "mlx5_0"
|
|
|
-socket_ifname = "eth0"
|
|
|
-ib_gid_index = 3
|
|
|
-ib_disable = 0
|
|
|
-"#;
|
|
|
- let config: Config = toml::from_str(toml_str).unwrap();
|
|
|
- assert_eq!(config.ssh_user, "admin");
|
|
|
- assert_eq!(config.clusters.len(), 1);
|
|
|
- assert_eq!(config.clusters[0].nodes.len(), 2);
|
|
|
- assert_eq!(config.clusters[0].nodes[0].host_key, None);
|
|
|
- assert!(config.nccl.contains_key("prod") && config.nccl["prod"].ib_hca == "mlx5_0");
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_config_deserialize_with_host_key() {
|
|
|
- let toml_str = "\nssh_user = \"admin\"\nssh_key_path = \"/path/to/key\"\n[[clusters]]\nname = \"secure\"\nmaster_ib_ip = \"1.1.1.1\"\nnodes = [\n { name = \"n1\", ssh_ip = \"2.2.2.2\", ib_ip = \"1.1.1.1\", rank = 0, host_key = \"SHA256:abc123\" },\n]\n";
|
|
|
- let config: Config = toml::from_str(toml_str).unwrap();
|
|
|
- assert_eq!(
|
|
|
- config.clusters[0].nodes[0].host_key,
|
|
|
- Some("SHA256:abc123".into())
|
|
|
- );
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_select_cluster_by_name() {
|
|
|
- let mut config = Config {
|
|
|
- ssh_user: "user".to_string(),
|
|
|
- ssh_key_path: "/path/to/key".into(),
|
|
|
- clusters: vec![test_cluster()],
|
|
|
- nccl: HashMap::new(),
|
|
|
- hf_cache_dir: None,
|
|
|
- init_port: 25000,
|
|
|
- before_serve: HashMap::new(),
|
|
|
- exec_aliases: HashMap::new(),
|
|
|
- home_assistant: None,
|
|
|
- global_mods_dir: None,
|
|
|
- };
|
|
|
- config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
- let (cluster, nccl) =
|
|
|
- select_cluster(&config, "test-cluster", None, &HashSet::new()).unwrap();
|
|
|
- assert_eq!(cluster.name, "test-cluster");
|
|
|
- assert_eq!(nccl.ib_hca, "rocep1s0f0");
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_select_cluster_auto() {
|
|
|
- let mut config = Config {
|
|
|
- ssh_user: "user".to_string(),
|
|
|
- ssh_key_path: "/path/to/key".into(),
|
|
|
- clusters: vec![test_cluster()],
|
|
|
- nccl: HashMap::new(),
|
|
|
- hf_cache_dir: None,
|
|
|
- init_port: 25000,
|
|
|
- before_serve: HashMap::new(),
|
|
|
- exec_aliases: HashMap::new(),
|
|
|
- home_assistant: None,
|
|
|
- global_mods_dir: None,
|
|
|
- };
|
|
|
- config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
- 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, &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, &HashSet::new());
|
|
|
- assert!(result.is_err());
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_select_cluster_by_name_found() {
|
|
|
- let mut nccl_map = HashMap::new();
|
|
|
- nccl_map.insert("test-cluster".to_string(), test_nccl());
|
|
|
- let mut config = test_config();
|
|
|
- config.nccl = nccl_map;
|
|
|
- let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
|
|
|
- assert_eq!(cluster.name, "test-cluster");
|
|
|
- assert_eq!(cluster.nodes.len(), 2);
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_select_cluster_by_name_not_found() {
|
|
|
- let config = test_config();
|
|
|
- let result = select_cluster_by_name(&config, "no-such-cluster");
|
|
|
- assert!(result.is_err());
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_select_cluster_by_prefix() {
|
|
|
- let mut config = test_config();
|
|
|
- config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
- let cluster = select_cluster_by_name(&config, "test-c").unwrap();
|
|
|
- assert_eq!(cluster.name, "test-cluster");
|
|
|
- let result = select_cluster_by_name(&config, "no-such").is_err();
|
|
|
- assert!(result);
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_russian_pluralization() {
|
|
|
- let _prev = fluent_i18n::set_locale(Some("ru")).ok();
|
|
|
- let s1 = fluent_i18n::t!("cluster-selected", { "name" => "c", "nodes" => 1u64 });
|
|
|
- assert!(s1.contains("1 узел"), "got: {}", s1);
|
|
|
- let s2 = fluent_i18n::t!("cluster-selected", { "name" => "c", "nodes" => 2u64 });
|
|
|
- assert!(s2.contains("2 узла"), "got: {}", s2);
|
|
|
- let s5 = fluent_i18n::t!("cluster-selected", { "name" => "c", "nodes" => 5u64 });
|
|
|
- assert!(s5.contains("5 узлов"), "got: {}", s5);
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_default_engine_is_vllm() {
|
|
|
- let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
|
|
|
- assert_eq!(recipe.engine, "vllm");
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn test_default_distributed_backend_is_none() {
|
|
|
- let recipe: Recipe = toml::from_str("name = \"x\"\nmodel = \"m\"").unwrap();
|
|
|
- assert_eq!(recipe.distributed_backend, None);
|
|
|
- }
|
|
|
- #[test]
|
|
|
- fn default_config_dir_appends_sm121_launcher() {
|
|
|
- let Some(dir) = default_config_dir() else {
|
|
|
- return;
|
|
|
- };
|
|
|
- assert!(dir.ends_with("sm121-launcher") && dir.parent().is_some());
|
|
|
- }
|
|
|
-}
|
|
|
+mod tests;
|