Skip to content

feat(spur-client): add client-side controller failover with endpoint rotation #384

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
shiv-tyagi wants to merge 1 commit into
base: main
Choose a base branch
from
+435 −79
Open

feat(spur-client): add client-side controller failover with endpoint rotation #384

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
30 changes: 29 additions & 1 deletion Cargo.lock
View file
Open in desktop

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ resolver = "2"
members = [
"crates/spur-proto",
"crates/spur-core",
"crates/spur-client",
"crates/spur-net",
"crates/spur-sched",
"crates/spur-metrics",
Expand Down
1 change: 1 addition & 0 deletions crates/spur-cli/Cargo.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ path = "src/main.rs"

[dependencies]
spur-proto = { path = "../spur-proto" }
spur-client = { path = "../spur-client" }
spur-core = { path = "../spur-core" }
spur-update = { path = "../spur-update" }
spur-net = { path = "../spur-net" }
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/exec.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,9 +38,10 @@ pub async fn main() -> Result<()> {
pub async fn main_with_args(args: Vec<String>) -> Result<()> {
let args = ExecArgs::try_parse_from(&args)?;

let mut client = SlurmControllerClient::connect(args.controller.clone())
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to controller")?;
let mut client = SlurmControllerClient::new(channel);

let resp = client
.exec_in_job(ExecInJobRequest {
Expand Down
19 changes: 4 additions & 15 deletions crates/spur-cli/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,21 +49,10 @@ fn load_controller_addr_from_config() {
}

if let Ok(config) = spur_core::config::SlurmConfig::load_from_file(&config_path) {
// Extract host and port from the config
let host = config
.controller
.hosts
.first()
.map(|h| h.as_str())
.unwrap_or("localhost");
let port = config
.controller
.listen_addr
.rsplit(':')
.next()
.unwrap_or("6817");
let addr = format!("http://{}:{}", host, port);
std::env::set_var("SPUR_CONTROLLER_ADDR", &addr);
let endpoints = config.controller.endpoints();
if !endpoints.is_empty() {
std::env::set_var("SPUR_CONTROLLER_ADDR", endpoints.join(","));
}
}
}

Expand Down
6 changes: 3 additions & 3 deletions crates/spur-cli/src/node.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ fn parse_label_args(label_args: &[String]) -> Result<(HashMap<String, String>, V
}

async fn cmd_label(controller: &str, node: String, label_args: Vec<String>) -> Result<()> {
let mut client = SlurmControllerClient::connect(controller.to_string()).await?;
let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?);
let (set_labels, remove_labels) = parse_label_args(&label_args)?;

client
Expand All @@ -128,7 +128,7 @@ async fn cmd_label(controller: &str, node: String, label_args: Vec<String>) -> R
}

async fn cmd_drain(controller: &str, node: String, reason: Option<String>) -> Result<()> {
let mut client = SlurmControllerClient::connect(controller.to_string()).await?;
let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?);
let resp = client
.drain_node(spur_proto::proto::DrainNodeRequest {
name: node.clone(),
Expand Down Expand Up @@ -158,7 +158,7 @@ async fn cmd_remove(
force: bool,
reason: Option<String>,
) -> Result<()> {
let mut client = SlurmControllerClient::connect(controller.to_string()).await?;
let mut client = SlurmControllerClient::new(spur_client::connect_channel(controller).await?);
let resp = client
.deregister_node(spur_proto::proto::DeregisterNodeRequest {
name: node.clone(),
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/salloc.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -109,9 +109,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
let ntasks = args.ntasks;
let cpus_per_task = args.cpus_per_task;

let mut client = SlurmControllerClient::connect(controller)
let channel = spur_client::connect_channel(&controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

// Submit interactive allocation (sleep infinity holds the allocation)
let job_spec = JobSpec {
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/sattach.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,9 +47,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
.and_then(|s| s.parse().ok())
.context("sattach: invalid job ID format (expected JOB_ID or JOB_ID.STEP_ID)")?;

let mut client = SlurmControllerClient::connect(args.controller)
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

// Look up the job to find which node it is running on
let job = client
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/sbatch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -646,9 +646,10 @@ pub async fn main_with_args(cli_args: Vec<String>) -> Result<()> {
};

// Submit to controller
let mut client = SlurmControllerClient::connect(args.controller)
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

let response = client
.submit_job(SubmitJobRequest {
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/scancel.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,9 +75,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
.user
.unwrap_or_else(|| whoami::username().unwrap_or_else(|_| "unknown".into()));

let mut client = SlurmControllerClient::connect(args.controller)
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

if !args.job_ids.is_empty() {
// Cancel specific jobs
Expand Down
30 changes: 20 additions & 10 deletions crates/spur-cli/src/scontrol.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,9 +168,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
}
ScontrolCommand::Requeue { job_id } => {
// Requeue = cancel + resubmit, simplified for now
let mut client = SlurmControllerClient::connect(args.controller.to_string())
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);
client
.cancel_job(spur_proto::proto::CancelJobRequest {
job_id,
Expand All @@ -183,9 +184,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
Ok(())
}
ScontrolCommand::Suspend { job_id } => {
let mut client = SlurmControllerClient::connect(args.controller.to_string())
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);
client
.suspend_job(spur_proto::proto::SuspendJobRequest {
job_id,
Expand All @@ -197,9 +199,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
Ok(())
}
ScontrolCommand::Resume { job_id } => {
let mut client = SlurmControllerClient::connect(args.controller.to_string())
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);
client
.resume_job(spur_proto::proto::ResumeJobRequest {
job_id,
Expand Down Expand Up @@ -246,9 +249,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
.filter(|s| !s.is_empty())
.collect()
};
let mut client = SlurmControllerClient::connect(args.controller.to_string())
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);
client
.update_reservation(spur_proto::proto::UpdateReservationRequest {
name: name.clone(),
Expand All @@ -272,9 +276,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
}

async fn show(controller: &str, entity: &str, name: Option<&str>) -> Result<()> {
let mut client = SlurmControllerClient::connect(controller.to_string())
let channel = spur_client::connect_channel(controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

match entity.to_lowercase().as_str() {
"job" | "jobs" => {
Expand Down Expand Up @@ -496,9 +501,10 @@ async fn show(controller: &str, entity: &str, name: Option<&str>) -> Result<()>
}

async fn ping(controller: &str) -> Result<()> {
let mut client = SlurmControllerClient::connect(controller.to_string())
let channel = spur_client::connect_channel(controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

let resp = client.ping(()).await.context("ping failed")?;

Expand Down Expand Up @@ -537,9 +543,10 @@ fn format_ts(ts: Option<&prost_types::Timestamp>) -> String {
async fn send_job_update(controller: &str, req: spur_proto::proto::UpdateJobRequest) -> Result<()> {
let hold = req.hold;
let job_id = req.job_id;
let mut client = SlurmControllerClient::connect(controller.to_string())
let channel = spur_client::connect_channel(controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

client.update_job(req).await.context("update failed")?;

Expand Down Expand Up @@ -624,9 +631,10 @@ async fn update_node(
state: Option<&str>,
reason: Option<String>,
) -> Result<()> {
let mut client = SlurmControllerClient::connect(controller.to_string())
let channel = spur_client::connect_channel(controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

let proto_state = state.map(|s| match s.to_lowercase().as_str() {
"idle" | "resume" => spur_proto::proto::NodeState::NodeIdle as i32,
Expand Down Expand Up @@ -666,9 +674,10 @@ async fn create_reservation(
accounts: &str,
users: &str,
) -> Result<()> {
let mut client = SlurmControllerClient::connect(controller.to_string())
let channel = spur_client::connect_channel(controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

let node_list: Vec<String> = nodes
.split(',')
Expand Down Expand Up @@ -704,9 +713,10 @@ async fn create_reservation(

/// Delete a reservation via the controller.
async fn delete_reservation(controller: &str, name: &str) -> Result<()> {
let mut client = SlurmControllerClient::connect(controller.to_string())
let channel = spur_client::connect_channel(controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

client
.delete_reservation(spur_proto::proto::DeleteReservationRequest {
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/sdiag.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -36,9 +36,10 @@ pub async fn main() -> Result<()> {
pub async fn main_with_args(args: Vec<String>) -> Result<()> {
let args = SdiagArgs::try_parse_from(&args)?;

let mut client = SlurmControllerClient::connect(args.controller.clone())
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

if args.reset {
client
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/sinfo.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,9 +70,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {

let fields = format_engine::parse_format(&fmt, &format_engine::sinfo_header);

let mut client = SlurmControllerClient::connect(args.controller)
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

// Get partitions
let partitions_resp = client
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/smd.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -42,9 +42,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
let args = SmdArgs::try_parse_from(&args)?;

loop {
let mut client = SlurmControllerClient::connect(args.controller.clone())
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

let nodes = client
.get_nodes(GetNodesRequest {
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/sprio.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,9 +53,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
})
.unwrap_or_default();

let mut client = SlurmControllerClient::connect(args.controller.clone())
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

// Get pending jobs only (priority is relevant for pending jobs)
let response = client
Expand Down
3 changes: 2 additions & 1 deletion crates/spur-cli/src/squeue.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,9 +93,10 @@ pub async fn main_with_args(args: Vec<String>) -> Result<()> {
.unwrap_or_default();

// Connect and fetch
let mut client = SlurmControllerClient::connect(args.controller)
let channel = spur_client::connect_channel(&args.controller)
.await
.context("failed to connect to spurctld")?;
let mut client = SlurmControllerClient::new(channel);

let response = client
.get_jobs(GetJobsRequest {
Expand Down
Loading