Add fetch running jobs in API to return all vec-inf jobs running for the user

Author: XkunW

vec_inf/client/_helper.py

@@ -721,7 +721,7 @@ def _get_base_status_data(self) -> StatusResponse:
             Basic status information for the job
         """
         try:
-            job_name = self.job_status["JobName"]
+            job_name = self.job_status["JobName"].removesuffix("-vec-inf")
             job_state = self.job_status["JobState"]
         except KeyError:
             job_name = "UNAVAILABLE"

vec_inf/client/_slurm_script_generator.py

@@ -89,6 +89,8 @@ def _generate_shebang(self) -> str:
         for arg, value in SLURM_JOB_CONFIG_ARGS.items():
             if self.params.get(value):
                 shebang.append(f"#SBATCH --{arg}={self.params[value]}")
+            if value == "model_name":
+                shebang[-1] += "-vec-inf"
         if self.is_multinode:
             shebang += SLURM_SCRIPT_TEMPLATE["shebang"]["multinode"]
         return "\n".join(shebang)
@@ -328,6 +330,8 @@ def _generate_batch_slurm_script_shebang(self) -> str:
                 model_params = self.params["models"][model_name]
                 if model_params.get(value) and value not in ["out_file", "err_file"]:
                     shebang.append(f"#SBATCH --{arg}={model_params[value]}")
+                if value == "model_name":
+                    shebang[-1] += "-vec-inf"
             shebang[-1] += "\n"
             shebang.append(BATCH_SLURM_SCRIPT_TEMPLATE["hetjob"])
         # Remove the last hetjob line

vec_inf/client/api.py

@@ -14,6 +14,8 @@
 import time
 import warnings
 from pathlib import Path
+import re
+import subprocess
 from typing import Any, Optional, Union
 
 from vec_inf.client._exceptions import (
@@ -181,6 +183,49 @@ def batch_launch_models(
         )
         return model_launcher.launch()
 
+    def fetch_running_jobs(self) -> list[str]:
+        """
+        Fetch the list of running vec-inf job IDs for the current user.
+
+        Returns
+        -------
+        list[str]
+            List of matching job names; empty list if squeue unavailable.
+        """
+
+        try:
+            # Run squeue for current user
+            res = subprocess.run(
+                ["squeue", "--me", "--noheader"],
+                capture_output=True, text=True, check=True
+            )
+            job_ids = [ln.strip().split()[0] for ln in res.stdout.splitlines() if ln.strip()]
+
+            if not job_ids:
+                return []
+
+            # For each job, fetch the full JobName and filter by suffix
+            matching_ids = []
+            for jid in job_ids:
+                try:
+                    sctl = subprocess.run(
+                        ["scontrol", "show", "job", "-o", jid],
+                        capture_output=True, text=True, check=True
+                    )
+                    # Example: "JobId=12345 JobName=my-long-job-name-vec-inf ..."
+                    m = re.search(r"\bJobName=([^\s]+)", sctl.stdout)
+                    if m and m.group(1).endswith("-vec-inf"):
+                        matching_ids.append(jid)
+                except subprocess.CalledProcessError:
+                    # Job might have finished between squeue and scontrol; skip
+                    continue
+
+            return matching_ids
+
+        except subprocess.CalledProcessError as e:
+            raise SlurmJobError(f"Error running slurm command: {e}") from e
+
+
     def get_status(self, slurm_job_id: str) -> StatusResponse:
         """Get the status of a running model.