feat: added oar job id inference

3 files changed, 307 insertions, 7 deletions

diff --git a/src/context.rs b/src/context.rs
index ae25966..8032212 100644
--- a/src/context.rs
+++ b/src/context.rs
@@ -13,20 +13,43 @@ pub enum ExecutionNode {
 pub struct Context {
     pub node: ExecutionNode,
     job_id: Option<u32>,
+    infer_job_id: bool,
     frontend_hostname: Option<String>,
 }
 
 impl Context {
-    pub async fn new(job_id: Option<u32>, frontend_hostname: Option<String>) -> Result<Self> {
+    pub async fn new(
+        job_id: Option<u32>,
+        infer_job_id: bool,
+        frontend_hostname: Option<String>,
+    ) -> Result<Self> {
         Ok(Self {
             node: get_execution_node().await?,
             job_id,
+            infer_job_id,
             frontend_hostname,
         })
     }
 
-    pub fn job_id(&self) -> Result<u32> {
-        self.job_id.ok_or_else(|| eyre::eyre!("missing job id"))
+    pub async fn job_id(&self) -> Result<u32> {
+        tracing::debug!("obtaining job id");
+        if let Some(job_id) = self.job_id {
+            tracing::debug!("job id was set, using {job_id}");
+            Ok(job_id)
+        } else if self.infer_job_id {
+            tracing::debug!("job id was not set but inference is enabled, finding job id");
+            let job_ids = crate::oar::list_user_job_ids(self).await?;
+            match job_ids.len() {
+                0 => Err(eyre::eyre!("cannot infer job id, no jobs are running")),
+                1 => Ok(job_ids[0]),
+                _ => Err(eyre::eyre!(
+                    "cannot infer job id, multiple jobs are running"
+                )),
+            }
+        } else {
+            tracing::debug!("inference was disabled and job id is not set");
+            Err(eyre::eyre!("missing job id"))
+        }
     }
 
     pub fn frontend_hostname(&self) -> Result<&str> {
diff --git a/src/main.rs b/src/main.rs
index 9b9c815..cc8afd0 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -46,6 +46,13 @@ struct Common {
     #[clap(long, env = "OAR_JOB_ID")]
     job_id: Option<u32>,
 
+    /// infer the oar job id
+    ///
+    /// if you have a single job running on the cluster, that job will be used, otherwise,
+    /// inference fails. the job_id flag takes precedence of over inference.
+    #[clap(long, env = "OAR_P2P_INFER_JOB_ID")]
+    infer_job_id: bool,
+
     /// hostname used to access the frontend using ssh.
     /// i.e. `ssh <frontend-hostname>` should work.
     #[clap(long, env = "FRONTEND_HOSTNAME")]
@@ -212,7 +219,12 @@ async fn main() -> Result<()> {
 }
 
 async fn context_from_common(common: &Common) -> Result<Context> {
-    Context::new(common.job_id, common.frontend_hostname.clone()).await
+    Context::new(
+        common.job_id,
+        common.infer_job_id,
+        common.frontend_hostname.clone(),
+    )
+    .await
 }
 
 async fn cmd_net_up(args: NetUpArgs) -> Result<()> {
@@ -658,7 +670,8 @@ async fn machines_configure(ctx: &Context, configs: &[MachineConfig]) -> Result<
 async fn machine_list_addresses(ctx: &Context, machine: Machine) -> Result<Vec<Ipv4Addr>> {
     tracing::info!("listing machine addresses");
     let interface = machine.interface();
-    let script = format!("ip addr show {interface} | grep -oE '10\\.[0-9]+\\.[0-9]+\\.[0-9]+' || true");
+    let script =
+        format!("ip addr show {interface} | grep -oE '10\\.[0-9]+\\.[0-9]+\\.[0-9]+' || true");
     let output = machine_run_script(ctx, machine, &script).await?;
     let stdout = std::str::from_utf8(&output.stdout)?;
     let mut addresses = Vec::default();
diff --git a/src/oar.rs b/src/oar.rs
index 2047d56..5a702b4 100644
--- a/src/oar.rs
+++ b/src/oar.rs
@@ -12,7 +12,7 @@ use crate::{
 pub async fn job_list_machines(ctx: &Context) -> Result<Vec<Machine>> {
     match ctx.node {
         ExecutionNode::Frontend => {
-            let job_id = ctx.job_id()?;
+            let job_id = ctx.job_id().await?;
             let output = Command::new("oarstat")
                 .arg("-j")
                 .arg(job_id.to_string())
@@ -36,7 +36,7 @@ pub async fn job_list_machines(ctx: &Context) -> Result<Vec<Machine>> {
             extract_machines_from_oar_stat_json(stdout, job_id)
         }
         ExecutionNode::Unknown => {
-            let job_id = ctx.job_id()?;
+            let job_id = ctx.job_id().await?;
             let frontend_hostname = ctx.frontend_hostname()?;
 
             let output = Command::new("ssh")
@@ -76,6 +76,51 @@ pub async fn job_list_machines(ctx: &Context) -> Result<Vec<Machine>> {
     }
 }
 
+pub async fn list_user_job_ids(ctx: &Context) -> Result<Vec<u32>> {
+    let output = match ctx.node {
+        ExecutionNode::Frontend => Command::new("oarstat").arg("-u").arg("-J").output().await?,
+        ExecutionNode::Unknown => {
+            Command::new("ssh")
+                .arg(ctx.frontend_hostname()?)
+                .arg("oarstat")
+                .arg("-u")
+                .arg("-J")
+                .output()
+                .await?
+        }
+        ExecutionNode::Machine(_) => {
+            return Err(eyre::eyre!(
+                "cannot run oarstat from inside a cluster machine"
+            ));
+        }
+    };
+
+    if !output.status.success() {
+        tracing::error!(
+            "stdout: {}",
+            std::str::from_utf8(&output.stdout).unwrap_or("stderr contains invalid uft-8")
+        );
+        tracing::error!(
+            "stderr: {}",
+            std::str::from_utf8(&output.stderr).unwrap_or("stderr contains invalid uft-8")
+        );
+        return Err(eyre::eyre!("failed to run oarstat"));
+    }
+
+    let stdout = String::from_utf8(output.stdout)?;
+    // for some reason, running oarstat with the -J flag (for json output) when you have no jobs
+    // running results in this error message instead of an empty object, so we will just assume it
+    // meant an empty object
+    let json_string = if stdout
+        == "hash- or arrayref expected (not a simple scalar, use allow_nonref to allow this) at /usr/lib/oar/oarstat line 285."
+    {
+        String::from("{}")
+    } else {
+        stdout
+    };
+    extract_job_ids_from_oarstat_output(&json_string)
+}
+
 fn extract_machines_from_oar_stat_json(output: &str, job_id: u32) -> Result<Vec<Machine>> {
     #[derive(Debug, Deserialize)]
     struct JobSchema {
@@ -96,6 +141,26 @@ fn extract_machines_from_oar_stat_json(output: &str, job_id: u32) -> Result<Vec<
     Ok(machines)
 }
 
+fn extract_job_ids_from_oarstat_output(output: &str) -> Result<Vec<u32>> {
+    let value = serde_json::from_str::<serde_json::Value>(output)?;
+    let object = match value {
+        serde_json::Value::Object(map) => map,
+        _ => {
+            return Err(eyre::eyre!(
+                "expected oar stat output to produce a json object"
+            ));
+        }
+    };
+
+    let mut job_ids = Vec::default();
+    for key in object.keys() {
+        tracing::trace!("parsing key '{key}'");
+        let job_id = key.parse()?;
+        job_ids.push(job_id);
+    }
+    Ok(job_ids)
+}
+
 #[cfg(test)]
 mod test {
     use super::*;
@@ -161,4 +226,203 @@ mod test {
         assert_eq!(machines[0], Machine::Gengar1);
         assert_eq!(machines[1], Machine::Gengar2);
     }
+
+    const OAR_STAT_ALL_USER_JOBS_OUTPUT: &'static str = r#"
+{
+   "37030" : {
+      "dependencies" : [],
+      "jobType" : "PASSIVE",
+      "state" : "Running",
+      "assigned_network_address" : [
+         "moltres-02"
+      ],
+      "command" : "sleep 365d",
+      "submissionTime" : 1752824505,
+      "name" : null,
+      "Job_Id" : 37030,
+      "stderr_file" : "OAR.37030.stderr",
+      "queue" : "default",
+      "launchingDirectory" : "/home/diogo464",
+      "reservation" : "None",
+      "properties" : "(( ( dedicated='NO' OR dedicated='protocol-labs' OR dedicated='tardis' )) AND desktop_computing = 'NO') AND drain='NO'",
+      "message" : "R=64,W=1:0:0,J=B (Karma=0.106,quota_ok)",
+      "stdout_file" : "OAR.37030.stdout",
+      "resubmit_job_id" : 0,
+      "types" : [],
+      "cpuset_name" : "diogo464_37030",
+      "array_index" : 1,
+      "project" : "default",
+      "array_id" : 37030,
+      "owner" : "diogo464",
+      "startTime" : 1752824506,
+      "assigned_resources" : [
+         893,
+         894,
+         895,
+         896,
+         897,
+         898,
+         899,
+         900,
+         901,
+         902,
+         903,
+         904,
+         905,
+         906,
+         907,
+         908,
+         909,
+         910,
+         911,
+         912,
+         913,
+         914,
+         915,
+         916,
+         917,
+         918,
+         919,
+         920,
+         921,
+         922,
+         923,
+         924,
+         925,
+         926,
+         927,
+         928,
+         929,
+         930,
+         931,
+         932,
+         933,
+         934,
+         935,
+         936,
+         937,
+         938,
+         939,
+         940,
+         941,
+         942,
+         943,
+         944,
+         945,
+         946,
+         947,
+         948,
+         949,
+         950,
+         951,
+         952,
+         953,
+         954,
+         955,
+         956
+      ]
+   },
+   "37029" : {
+      "command" : "sleep 365d",
+      "submissionTime" : 1752824490,
+      "dependencies" : [],
+      "jobType" : "PASSIVE",
+      "state" : "Running",
+      "assigned_network_address" : [
+         "moltres-01"
+      ],
+      "Job_Id" : 37029,
+      "stderr_file" : "OAR.37029.stderr",
+      "name" : null,
+      "types" : [],
+      "cpuset_name" : "diogo464_37029",
+      "launchingDirectory" : "/home/diogo464",
+      "queue" : "default",
+      "reservation" : "None",
+      "message" : "R=64,W=1:0:0,J=B (Karma=0.106,quota_ok)",
+      "stdout_file" : "OAR.37029.stdout",
+      "properties" : "(( ( dedicated='NO' OR dedicated='protocol-labs' OR dedicated='tardis' )) AND desktop_computing = 'NO') AND drain='NO'",
+      "resubmit_job_id" : 0,
+      "startTime" : 1752824491,
+      "assigned_resources" : [
+         829,
+         830,
+         831,
+         832,
+         833,
+         834,
+         835,
+         836,
+         837,
+         838,
+         839,
+         840,
+         841,
+         842,
+         843,
+         844,
+         845,
+         846,
+         847,
+         848,
+         849,
+         850,
+         851,
+         852,
+         853,
+         854,
+         855,
+         856,
+         857,
+         858,
+         859,
+         860,
+         861,
+         862,
+         863,
+         864,
+         865,
+         866,
+         867,
+         868,
+         869,
+         870,
+         871,
+         872,
+         873,
+         874,
+         875,
+         876,
+         877,
+         878,
+         879,
+         880,
+         881,
+         882,
+         883,
+         884,
+         885,
+         886,
+         887,
+         888,
+         889,
+         890,
+         891,
+         892
+      ],
+      "array_index" : 1,
+      "project" : "default",
+      "array_id" : 37029,
+      "owner" : "diogo464"
+   }
+}
+"#;
+
+    #[test]
+    fn test_extract_job_ids_from_oarstat_output() {
+        let job_ids = extract_job_ids_from_oarstat_output(OAR_STAT_ALL_USER_JOBS_OUTPUT).unwrap();
+        assert_eq!(job_ids.len(), 2);
+        assert!(job_ids.contains(&37030));
+        assert!(job_ids.contains(&37029));
+    }
 }