|
@@ -13,6 +13,8 @@ pub struct Node {
|
|
|
pub rank: u32,
|
|
pub rank: u32,
|
|
|
#[serde(default)]
|
|
#[serde(default)]
|
|
|
pub host_key: Option<String>,
|
|
pub host_key: Option<String>,
|
|
|
|
|
+ #[serde(default)]
|
|
|
|
|
+ pub entity: Option<String>,
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
#[derive(Clone, Deserialize)]
|
|
#[derive(Clone, Deserialize)]
|
|
@@ -30,8 +32,7 @@ pub struct Cluster {
|
|
|
pub nodes: Vec<Node>,
|
|
pub nodes: Vec<Node>,
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
-/// Top-level config containing cluster definitions and per-cluster
|
|
|
|
|
-/// pre-serve commands.
|
|
|
|
|
|
|
+/// Top-level config (clusters + per-cluster pre-serve commands).
|
|
|
#[derive(Clone, Deserialize)]
|
|
#[derive(Clone, Deserialize)]
|
|
|
pub struct Config {
|
|
pub struct Config {
|
|
|
pub ssh_user: String,
|
|
pub ssh_user: String,
|
|
@@ -41,15 +42,29 @@ pub struct Config {
|
|
|
pub nccl: HashMap<String, NcclConfig>,
|
|
pub nccl: HashMap<String, NcclConfig>,
|
|
|
#[serde(default)]
|
|
#[serde(default)]
|
|
|
pub hf_cache_dir: Option<String>,
|
|
pub hf_cache_dir: Option<String>,
|
|
|
- /// TCP port used for cluster init (vLLM/SGLang). Default: 25000.
|
|
|
|
|
|
|
+ /// TCP port for cluster init (vLLM/SGLang).
|
|
|
#[serde(default = "default_init_port")]
|
|
#[serde(default = "default_init_port")]
|
|
|
pub init_port: u16,
|
|
pub init_port: u16,
|
|
|
- /// Per-cluster pre-serve commands. Looked up by cluster name.
|
|
|
|
|
|
|
+ /// Per-cluster pre-serve commands, command aliases, HA config.
|
|
|
#[serde(default)]
|
|
#[serde(default)]
|
|
|
pub before_serve: HashMap<String, BeforeServeTable>,
|
|
pub before_serve: HashMap<String, BeforeServeTable>,
|
|
|
- /// Command aliases for `cargo run exec --alias`.
|
|
|
|
|
#[serde(default)]
|
|
#[serde(default)]
|
|
|
pub exec_aliases: HashMap<String, String>,
|
|
pub exec_aliases: HashMap<String, String>,
|
|
|
|
|
+ #[serde(default)]
|
|
|
|
|
+ pub home_assistant: Option<HomeAssistantConfig>,
|
|
|
|
|
+}
|
|
|
|
|
+
|
|
|
|
|
+/// Home Assistant config for `power on`/`power off`.
|
|
|
|
|
+#[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 {
|
|
fn default_init_port() -> u16 {
|
|
@@ -57,7 +72,6 @@ fn default_init_port() -> u16 {
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
impl Config {
|
|
impl Config {
|
|
|
- /// HF cache directory for volume mounting.
|
|
|
|
|
pub fn resolved_hf_cache_dir(&self) -> String {
|
|
pub fn resolved_hf_cache_dir(&self) -> String {
|
|
|
if let Some(ref dir) = self.hf_cache_dir {
|
|
if let Some(ref dir) = self.hf_cache_dir {
|
|
|
return dir.clone();
|
|
return dir.clone();
|
|
@@ -98,8 +112,7 @@ pub struct Recipe {
|
|
|
pub entrypoint: Vec<String>,
|
|
pub entrypoint: Vec<String>,
|
|
|
#[serde(default = "default_distributed_backend")]
|
|
#[serde(default = "default_distributed_backend")]
|
|
|
pub distributed_backend: String,
|
|
pub distributed_backend: String,
|
|
|
- /// Recipe-level pre-serve commands. Appended after cluster-level
|
|
|
|
|
- /// commands from `[before_serve."<cluster>"]`.
|
|
|
|
|
|
|
+ /// Recipe-level pre-serve commands (appended after cluster-level).
|
|
|
#[serde(default)]
|
|
#[serde(default)]
|
|
|
pub before_serve: Vec<PreServeCommand>,
|
|
pub before_serve: Vec<PreServeCommand>,
|
|
|
}
|
|
}
|
|
@@ -116,8 +129,7 @@ fn default_tensor_parallel_size() -> u32 {
|
|
|
1
|
|
1
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
-/// Selects a cluster and its NCCL config by name, or auto-picks the first
|
|
|
|
|
-/// cluster that has NCCL config.
|
|
|
|
|
|
|
+/// Selects a cluster + NCCL config by name or auto-picks the first one with NCCL.
|
|
|
pub fn select_cluster<'c>(
|
|
pub fn select_cluster<'c>(
|
|
|
config: &'c Config,
|
|
config: &'c Config,
|
|
|
cluster_name: &str,
|
|
cluster_name: &str,
|
|
@@ -171,13 +183,8 @@ pub fn resolve_cluster<'c>(
|
|
|
select_cluster(config, &choice)
|
|
select_cluster(config, &choice)
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
-/// Loads a recipe from disk.
|
|
|
|
|
-///
|
|
|
|
|
-/// Search order:
|
|
|
|
|
-/// 1. `./recipes/<name>/recipe.toml` (cwd)
|
|
|
|
|
-/// 2. `<default_config_dir>/recipes/<name>/recipe.toml`
|
|
|
|
|
-///
|
|
|
|
|
-/// Falls back to the cwd path so the error message mentions it.
|
|
|
|
|
|
|
+/// Loads a recipe (cwd first, then `<default_config_dir>/recipes/<name>/recipe.toml`,
|
|
|
|
|
+/// falling back to cwd for the error).
|
|
|
pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
|
|
pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
|
|
|
let cwd_path = format!("recipes/{}/recipe.toml", recipe_name);
|
|
let cwd_path = format!("recipes/{}/recipe.toml", recipe_name);
|
|
|
if let Ok(content) = tokio::fs::read_to_string(&cwd_path).await {
|
|
if let Ok(content) = tokio::fs::read_to_string(&cwd_path).await {
|
|
@@ -194,15 +201,13 @@ pub async fn load_recipe(recipe_name: &str) -> Result<Recipe> {
|
|
|
Ok(toml::from_str(&content)?)
|
|
Ok(toml::from_str(&content)?)
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
-/// Returns the OS-appropriate default config directory for the launcher
|
|
|
|
|
-/// (e.g. `~/.config/sm121-launcher` on Linux), or `None` if the OS
|
|
|
|
|
-/// has no known config dir (some restricted sandboxes).
|
|
|
|
|
|
|
+/// OS-appropriate default config dir (e.g. `~/.config/sm121-launcher`),
|
|
|
|
|
+/// or `None` for sandboxes with no known config dir.
|
|
|
pub fn default_config_dir() -> Option<std::path::PathBuf> {
|
|
pub fn default_config_dir() -> Option<std::path::PathBuf> {
|
|
|
dirs::config_dir().map(|p| p.join("sm121-launcher"))
|
|
dirs::config_dir().map(|p| p.join("sm121-launcher"))
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
-/// Returns the cluster and recipe name for a recipe (for stop/logs commands
|
|
|
|
|
-/// that only receive a recipe name).
|
|
|
|
|
|
|
+/// Resolves the cluster for a recipe (stop/logs flow).
|
|
|
pub async fn resolve_cluster_for_recipe(
|
|
pub async fn resolve_cluster_for_recipe(
|
|
|
config: &Config,
|
|
config: &Config,
|
|
|
recipe_name: &str,
|
|
recipe_name: &str,
|
|
@@ -227,6 +232,7 @@ mod tests {
|
|
|
ib_ip: "192.168.1.1".to_string(),
|
|
ib_ip: "192.168.1.1".to_string(),
|
|
|
rank: 0,
|
|
rank: 0,
|
|
|
host_key: None,
|
|
host_key: None,
|
|
|
|
|
+ entity: None,
|
|
|
},
|
|
},
|
|
|
Node {
|
|
Node {
|
|
|
name: "node1".to_string(),
|
|
name: "node1".to_string(),
|
|
@@ -234,11 +240,26 @@ mod tests {
|
|
|
ib_ip: "192.168.1.2".to_string(),
|
|
ib_ip: "192.168.1.2".to_string(),
|
|
|
rank: 1,
|
|
rank: 1,
|
|
|
host_key: None,
|
|
host_key: None,
|
|
|
|
|
+ entity: None,
|
|
|
},
|
|
},
|
|
|
],
|
|
],
|
|
|
}
|
|
}
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
|
|
+ fn test_config() -> Config {
|
|
|
|
|
+ Config {
|
|
|
|
|
+ ssh_user: "u".to_string(),
|
|
|
|
|
+ ssh_key_path: "/k".to_string(),
|
|
|
|
|
+ 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,
|
|
|
|
|
+ }
|
|
|
|
|
+ }
|
|
|
|
|
+
|
|
|
fn test_nccl() -> NcclConfig {
|
|
fn test_nccl() -> NcclConfig {
|
|
|
NcclConfig {
|
|
NcclConfig {
|
|
|
ib_hca: "rocep1s0f0".to_string(),
|
|
ib_hca: "rocep1s0f0".to_string(),
|
|
@@ -375,6 +396,7 @@ nodes = [
|
|
|
init_port: 25000,
|
|
init_port: 25000,
|
|
|
before_serve: HashMap::new(),
|
|
before_serve: HashMap::new(),
|
|
|
exec_aliases: HashMap::new(),
|
|
exec_aliases: HashMap::new(),
|
|
|
|
|
+ home_assistant: None,
|
|
|
};
|
|
};
|
|
|
config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
|
|
|
|
@@ -394,6 +416,7 @@ nodes = [
|
|
|
init_port: 25000,
|
|
init_port: 25000,
|
|
|
before_serve: HashMap::new(),
|
|
before_serve: HashMap::new(),
|
|
|
exec_aliases: HashMap::new(),
|
|
exec_aliases: HashMap::new(),
|
|
|
|
|
+ home_assistant: None,
|
|
|
};
|
|
};
|
|
|
config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
config.nccl.insert("test-cluster".to_string(), test_nccl());
|
|
|
|
|
|
|
@@ -404,32 +427,15 @@ nodes = [
|
|
|
|
|
|
|
|
#[test]
|
|
#[test]
|
|
|
fn test_select_cluster_auto_no_match() {
|
|
fn test_select_cluster_auto_no_match() {
|
|
|
- let config = Config {
|
|
|
|
|
- ssh_user: "user".to_string(),
|
|
|
|
|
- ssh_key_path: "/path/to/key".to_string(),
|
|
|
|
|
- clusters: vec![test_cluster()],
|
|
|
|
|
- nccl: HashMap::new(),
|
|
|
|
|
- hf_cache_dir: None,
|
|
|
|
|
- init_port: 25000,
|
|
|
|
|
- before_serve: HashMap::new(),
|
|
|
|
|
- exec_aliases: HashMap::new(),
|
|
|
|
|
- };
|
|
|
|
|
|
|
+ let config = test_config();
|
|
|
let result = select_cluster(&config, "auto");
|
|
let result = select_cluster(&config, "auto");
|
|
|
assert!(result.is_err());
|
|
assert!(result.is_err());
|
|
|
}
|
|
}
|
|
|
|
|
|
|
|
#[test]
|
|
#[test]
|
|
|
fn test_select_cluster_not_found() {
|
|
fn test_select_cluster_not_found() {
|
|
|
- let config = Config {
|
|
|
|
|
- ssh_user: "user".to_string(),
|
|
|
|
|
- ssh_key_path: "/path/to/key".to_string(),
|
|
|
|
|
- clusters: vec![],
|
|
|
|
|
- nccl: HashMap::new(),
|
|
|
|
|
- hf_cache_dir: None,
|
|
|
|
|
- init_port: 25000,
|
|
|
|
|
- before_serve: HashMap::new(),
|
|
|
|
|
- exec_aliases: HashMap::new(),
|
|
|
|
|
- };
|
|
|
|
|
|
|
+ let mut config = test_config();
|
|
|
|
|
+ config.clusters = vec![];
|
|
|
let result = select_cluster(&config, "nonexistent");
|
|
let result = select_cluster(&config, "nonexistent");
|
|
|
assert!(result.is_err());
|
|
assert!(result.is_err());
|
|
|
}
|
|
}
|
|
@@ -438,16 +444,8 @@ nodes = [
|
|
|
fn test_select_cluster_by_name_found() {
|
|
fn test_select_cluster_by_name_found() {
|
|
|
let mut nccl_map = HashMap::new();
|
|
let mut nccl_map = HashMap::new();
|
|
|
nccl_map.insert("test-cluster".to_string(), test_nccl());
|
|
nccl_map.insert("test-cluster".to_string(), test_nccl());
|
|
|
- let config = Config {
|
|
|
|
|
- ssh_user: "u".to_string(),
|
|
|
|
|
- ssh_key_path: "/k".to_string(),
|
|
|
|
|
- clusters: vec![test_cluster()],
|
|
|
|
|
- nccl: nccl_map,
|
|
|
|
|
- hf_cache_dir: None,
|
|
|
|
|
- init_port: 25000,
|
|
|
|
|
- before_serve: HashMap::new(),
|
|
|
|
|
- exec_aliases: HashMap::new(),
|
|
|
|
|
- };
|
|
|
|
|
|
|
+ let mut config = test_config();
|
|
|
|
|
+ config.nccl = nccl_map;
|
|
|
let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
|
|
let cluster = select_cluster_by_name(&config, "test-cluster").unwrap();
|
|
|
assert_eq!(cluster.name, "test-cluster");
|
|
assert_eq!(cluster.name, "test-cluster");
|
|
|
assert_eq!(cluster.nodes.len(), 2);
|
|
assert_eq!(cluster.nodes.len(), 2);
|
|
@@ -455,16 +453,7 @@ nodes = [
|
|
|
|
|
|
|
|
#[test]
|
|
#[test]
|
|
|
fn test_select_cluster_by_name_not_found() {
|
|
fn test_select_cluster_by_name_not_found() {
|
|
|
- let config = Config {
|
|
|
|
|
- ssh_user: "u".to_string(),
|
|
|
|
|
- ssh_key_path: "/k".to_string(),
|
|
|
|
|
- clusters: vec![test_cluster()],
|
|
|
|
|
- nccl: HashMap::new(),
|
|
|
|
|
- hf_cache_dir: None,
|
|
|
|
|
- init_port: 25000,
|
|
|
|
|
- before_serve: HashMap::new(),
|
|
|
|
|
- exec_aliases: HashMap::new(),
|
|
|
|
|
- };
|
|
|
|
|
|
|
+ let config = test_config();
|
|
|
let result = select_cluster_by_name(&config, "no-such-cluster");
|
|
let result = select_cluster_by_name(&config, "no-such-cluster");
|
|
|
assert!(result.is_err());
|
|
assert!(result.is_err());
|
|
|
}
|
|
}
|