mpi_cluster_tools 0.1.5

A collection of cluster management commands for HPC environments
use crate::config::ClusterConfig;
use crate::cmd::condor::condor_q_for_user;
use crate::utils::serde::{deserialize_i64_lenient, deserialize_request_gpus};
use comfy_table::{
    presets::UTF8_FULL, Attribute, Cell, CellAlignment, Color, ContentArrangement, Table,
};
use crossterm::terminal;
use serde::Deserialize;
use std::time::{Duration, SystemTime, UNIX_EPOCH};

#[derive(Deserialize, Debug)]
struct JobRow {
    #[serde(rename = "ClusterId")]
    cluster_id: i64,
    #[serde(rename = "ProcId")]
    proc_id: i64,
    #[serde(rename = "Cmd")]
    cmd: Option<String>,
    #[serde(rename = "Args")]
    args: Option<String>,
    #[serde(rename = "JobPrio")]
    job_prio: i32,
    #[serde(
        rename = "RequestGPUs",
        default,
        deserialize_with = "deserialize_request_gpus"
    )]
    request_gpus: i32,
    #[serde(
        rename = "JobStartDate",
        default,
        deserialize_with = "deserialize_i64_lenient"
    )]
    start_unix: i64,
    #[allow(dead_code)]
    #[serde(
        rename = "QDate",
        default,
        deserialize_with = "deserialize_i64_lenient"
    )]
    q_unix: i64,
}

fn price_from_prio(job_prio: i32) -> f64 {
    (job_prio + 1000) as f64
}

fn human_duration_from_unix(start: i64) -> String {
    if start <= 0 {
        return "-".to_string();
    }
    let now = SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .unwrap_or(Duration::from_secs(0))
        .as_secs() as i64;
    let secs = (now - start).max(0) as u64;
    let h = secs / 3600;
    let m = (secs % 3600) / 60;
    let s = secs % 60;
    if h > 0 {
        format!("{}h {:02}m", h, m)
    } else if m > 0 {
        format!("{}m {:02}s", m, s)
    } else {
        format!("{}s", s)
    }
}

fn render_table(rows: &[JobRow]) -> Table {
    let mut table = Table::new();
    table
        .load_preset(UTF8_FULL)
        .set_content_arrangement(ContentArrangement::Dynamic)
        .set_header(vec![
            Cell::new("JobID").add_attribute(Attribute::Bold),
            Cell::new("Cmd").add_attribute(Attribute::Bold),
            Cell::new("Args").add_attribute(Attribute::Bold),
            Cell::new("Runtime").add_attribute(Attribute::Bold),
            Cell::new("GPUs")
                .add_attribute(Attribute::Bold)
                .set_alignment(CellAlignment::Right),
            Cell::new("Bid")
                .add_attribute(Attribute::Bold)
                .set_alignment(CellAlignment::Right),
        ]);

    let (cols, _) = terminal::size().unwrap();
    table.set_width(cols as u16);

    for j in rows {
        let jobid = format!("{}.{}", j.cluster_id, j.proc_id);
        let runtime = human_duration_from_unix(j.start_unix);
        let gpus = j.request_gpus;
        let bid = price_from_prio(j.job_prio);

        table.add_row(vec![
            Cell::new(jobid).fg(Color::Green),
            Cell::new(j.cmd.as_deref().unwrap_or("")),
            Cell::new(j.args.as_deref().unwrap_or("")),
            Cell::new(runtime),
            Cell::new(gpus.to_string()).set_alignment(CellAlignment::Right),
            Cell::new(format!("{:.0}", bid)).set_alignment(CellAlignment::Right),
        ]);
    }

    table
}

pub fn handle_list_jobs() -> Result<(), Box<dyn std::error::Error>> {
    let config = ClusterConfig::load();
    let login = match &config.login {
        Some(l) => l,
        None => {
            eprintln!("Error: No login configuration found. Run 'mct login' first.");
            std::process::exit(1);
        }
    };

    let username = config.get_username().ok_or_else(|| {
        std::io::Error::new(
            std::io::ErrorKind::Other,
            "No username found in login or ssh config",
        )
    })?;

    // Query current user's jobs (all states) with needed attributes
    let attrs = [
        "ClusterId",
        "ProcId",
        "Cmd",
        "Args",
        "JobPrio",
        "RequestGPUs",
        "RequestMemory",
        "MemoryProvisioned",
        "RequestCpus",
        "CpusProvisioned",
        "JobStartDate",
        "QDate",
    ]
    .join(",");
    let jobs: Vec<JobRow> = condor_q_for_user(login, &username, &attrs)?;
    if jobs.is_empty() {
        println!("No jobs found for user {}.", username);
        return Ok(());
    }

    let table = render_table(&jobs);
    println!("{}", table);
    Ok(())
}