Skip to content
Open
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 80 additions & 0 deletions crates/spur-cli/src/k8s.rs
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,16 @@ pub enum K8sCommand {
/// Overrides --replicas.
#[arg(long = "control-plane-nodes", value_delimiter = ',')]
control_plane_nodes: Vec<String>,
/// Scope the cluster to a subset of nodes (hostlist, e.g. "gpu[01-08]"). Combined with
/// --partition/--selector as a union; empty = enroll the whole inventory.
#[arg(long)]
nodes: Option<String>,
/// Scope the cluster to a partition's nodes.
#[arg(long)]
partition: Option<String>,
/// Scope the cluster to nodes matching every key=value label (repeatable).
#[arg(long = "selector", value_parser = parse_key_val)]
selector: Vec<(String, String)>,
},
/// Tear the k0s cluster down.
Down {
Expand Down Expand Up @@ -88,12 +98,18 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
control_plane_node,
replicas,
control_plane_nodes,
nodes,
partition,
selector,
} => {
cmd_up(
&controller,
control_plane_node,
replicas,
control_plane_nodes,
nodes,
partition,
selector,
)
.await
}
Expand All @@ -112,6 +128,16 @@ fn effective_user() -> String {
whoami::username().unwrap_or_else(|_| "unknown".into())
}

fn parse_key_val(s: &str) -> Result<(String, String), String> {
let (k, v) = s
.split_once('=')
.ok_or_else(|| format!("expected key=value, got {s}"))?;
if k.is_empty() {
return Err(format!("empty selector key in {s}"));
}
Ok((k.to_string(), v.to_string()))
}

async fn cmd_install_k0s(version: &str, path: &str, force: bool) -> Result<()> {
let dest = std::path::Path::new(path);
if dest.exists() && !force {
Expand All @@ -130,11 +156,15 @@ async fn cmd_install_k0s(version: &str, path: &str, force: bool) -> Result<()> {
Ok(())
}

#[allow(clippy::too_many_arguments)]
async fn cmd_up(
controller: &str,
control_plane_node: Option<String>,
replicas: Option<u32>,
control_plane_nodes: Vec<String>,
nodes: Option<String>,
partition: Option<String>,
selector: Vec<(String, String)>,
) -> Result<()> {
let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?);
let resp = client
Expand All @@ -143,6 +173,9 @@ async fn cmd_up(
control_plane_replicas: replicas,
control_plane_nodes,
caller: effective_user(),
nodes: nodes.unwrap_or_default(),
partition: partition.unwrap_or_default(),
selector: selector.into_iter().collect(),
Comment thread
yansun1996 marked this conversation as resolved.
Outdated
})
.await?
.into_inner();
Expand Down Expand Up @@ -186,6 +219,11 @@ async fn cmd_status(controller: &str) -> Result<()> {
} else if !resp.control_plane_node.is_empty() {
println!("control-plane: {}", resp.control_plane_node);
}
if resp.member_nodes.is_empty() {
println!("members: all nodes");
} else {
println!("members: {}", resp.member_nodes.join(", "));
}
for n in resp.nodes {
println!(
" {:<24} {:<11} {:<11} enabled={}",
Expand Down Expand Up @@ -223,6 +261,7 @@ mod tests {
control_plane_node,
replicas,
control_plane_nodes,
..
} => {
assert_eq!(control_plane_node.as_deref(), Some("head-node"));
assert_eq!(replicas, None);
Expand All @@ -232,6 +271,47 @@ mod tests {
}
}

#[test]
fn parses_up_with_node_scope_flags() {
let args = K8sArgs::try_parse_from([
"k8s",
"up",
"--nodes",
"gpu[01-08]",
"--partition",
"batch",
"--selector",
"zone=z1",
"--selector",
"gpu=mi300",
])
.unwrap();
match args.command {
K8sCommand::Up {
nodes,
partition,
selector,
..
} => {
assert_eq!(nodes.as_deref(), Some("gpu[01-08]"));
assert_eq!(partition.as_deref(), Some("batch"));
assert_eq!(
selector,
vec![
("zone".to_string(), "z1".to_string()),
("gpu".to_string(), "mi300".to_string())
]
);
}
_ => panic!("wrong command"),
}
}

#[test]
fn selector_without_equals_is_rejected() {
assert!(K8sArgs::try_parse_from(["k8s", "up", "--selector", "bogus"]).is_err());
}

#[test]
fn parses_up_with_replicas_and_node_set() {
let args = K8sArgs::try_parse_from(["k8s", "up", "--replicas", "3"]).unwrap();
Expand Down
30 changes: 27 additions & 3 deletions crates/spur-core/src/k0s.rs
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ mod cluster_state_tests {
phase: K0sPhase::Ready,
control_plane_node: Some("cp-1".into()),
control_plane_nodes: vec!["cp-1".into(), "cp-2".into(), "cp-3".into()],
reset_requested: false,
..Default::default()
};
assert_eq!(st.controllers(), vec!["cp-1", "cp-2", "cp-3"]);
assert_eq!(st.bootstrap().as_deref(), Some("cp-1"));
Expand All @@ -135,9 +135,8 @@ mod cluster_state_tests {
fn bootstrap_falls_back_to_first_of_set_when_singular_absent() {
let st = K0sClusterState {
phase: K0sPhase::Ready,
control_plane_node: None,
control_plane_nodes: vec!["cp-1".into(), "cp-2".into(), "cp-3".into()],
reset_requested: false,
..Default::default()
};
assert_eq!(st.bootstrap().as_deref(), Some("cp-1"));
}
Expand All @@ -146,6 +145,23 @@ mod cluster_state_tests {
fn controllers_empty_when_down() {
assert!(K0sClusterState::default().controllers().is_empty());
}

#[test]
fn is_member_empty_scope_matches_all() {
let st = K0sClusterState::default();
assert!(st.is_member("anything"), "empty scope = whole inventory");
}

#[test]
fn is_member_respects_recorded_scope() {
let st = K0sClusterState {
member_nodes: vec!["a".into(), "b".into()],
..Default::default()
};
assert!(st.is_member("a"));
assert!(st.is_member("b"));
assert!(!st.is_member("c"), "out-of-scope node excluded");
}
}

#[cfg(test)]
Expand Down Expand Up @@ -233,6 +249,9 @@ pub struct K0sClusterState {
/// All control-plane nodes (1/3/5). Empty on pre-multi-CP state — read via [`Self::controllers`].
#[serde(default)]
pub control_plane_nodes: Vec<String>,
/// Nodes the cluster is scoped to. Empty = enroll the whole inventory (back-compat).
#[serde(default)]
pub member_nodes: Vec<String>,
#[serde(default)]
pub reset_requested: bool,
}
Expand All @@ -254,4 +273,9 @@ impl K0sClusterState {
.clone()
.or_else(|| self.control_plane_nodes.first().cloned())
}

/// Whether `name` is in scope for enrollment. An empty `member_nodes` means the whole inventory.
pub fn is_member(&self, name: &str) -> bool {
self.member_nodes.is_empty() || self.member_nodes.iter().any(|n| n == name)
}
}
33 changes: 33 additions & 0 deletions crates/spur-core/src/wal.rs
Original file line number Diff line number Diff line change
Expand Up @@ -257,6 +257,8 @@ pub enum WalOperation {
#[serde(default)]
control_plane_nodes: Vec<String>,
#[serde(default)]
member_nodes: Vec<String>,
#[serde(default)]
reset_requested: bool,
},
NodeK0sClear {
Expand Down Expand Up @@ -765,6 +767,12 @@ mod deregistration_wal_tests {
phase: K0sPhase::Ready,
control_plane_node: Some("head-node".into()),
control_plane_nodes: vec!["head-node".into(), "cp-2".into(), "cp-3".into()],
member_nodes: vec![
"head-node".into(),
"cp-2".into(),
"cp-3".into(),
"w-4".into(),
],
reset_requested: false,
};
let back: WalOperation =
Expand All @@ -774,11 +782,13 @@ mod deregistration_wal_tests {
phase,
control_plane_node,
control_plane_nodes,
member_nodes,
reset_requested,
} => {
assert_eq!(phase, K0sPhase::Ready);
assert_eq!(control_plane_node.as_deref(), Some("head-node"));
assert_eq!(control_plane_nodes, vec!["head-node", "cp-2", "cp-3"]);
assert_eq!(member_nodes, vec!["head-node", "cp-2", "cp-3", "w-4"]);
assert!(!reset_requested);
}
_ => panic!("wrong variant"),
Expand All @@ -798,17 +808,40 @@ mod deregistration_wal_tests {
phase,
control_plane_node,
control_plane_nodes,
member_nodes,
reset_requested,
} => {
assert_eq!(phase, K0sPhase::Ready);
assert_eq!(control_plane_node.as_deref(), Some("head-node"));
assert!(control_plane_nodes.is_empty());
assert!(member_nodes.is_empty());
assert!(!reset_requested);
}
_ => panic!("wrong variant"),
}
}

// Frozen pre-member-scope K0sSetPhase entry (has control_plane_nodes, no member_nodes); must
// still deserialize with member_nodes defaulting empty (= whole inventory). Never regenerate.
#[test]
fn k0s_set_phase_pre_member_scope_payload_still_deserializes() {
const K0S_SET_PHASE_PRE_MEMBER_SCOPE: &str = r#"{"K0sSetPhase":{"phase":"ready","control_plane_node":"head-node","control_plane_nodes":["head-node","cp-2","cp-3"],"reset_requested":false}}"#;
let op: WalOperation = serde_json::from_str(K0S_SET_PHASE_PRE_MEMBER_SCOPE).expect(
"pre-member-scope K0sSetPhase must deserialize; member_nodes needs #[serde(default)]",
);
match op {
WalOperation::K0sSetPhase {
control_plane_nodes,
member_nodes,
..
} => {
assert_eq!(control_plane_nodes, vec!["head-node", "cp-2", "cp-3"]);
assert!(member_nodes.is_empty());
}
_ => panic!("wrong variant"),
}
}

#[test]
fn node_remove_none_reason_round_trips() {
let op = WalOperation::NodeRemove {
Expand Down
Loading