Update status to return all running vec-inf job status if no job ID is provided

Author: XkunW

vec_inf/cli/_cli.py

@@ -30,6 +30,7 @@
     BatchLaunchResponseFormatter,
     LaunchResponseFormatter,
     ListCmdDisplay,
+    ListStatusDisplay,
     MetricsResponseFormatter,
     StatusResponseFormatter,
 )
@@ -313,14 +314,14 @@ def batch_launch(
         raise click.ClickException(f"Batch launch failed: {str(e)}") from e
 
 
-@cli.command("status", help="Check the status of a running model on the cluster.")
-@click.argument("slurm_job_id", type=str, nargs=1)
+@cli.command("status", help="Check the status of running vec-inf jobs on the cluster.")
+@click.argument("slurm_job_id", required=False)
 @click.option(
     "--json-mode",
     is_flag=True,
     help="Output in JSON string",
 )
-def status(slurm_job_id: str, json_mode: bool = False) -> None:
+def status(slurm_job_id: str = None, json_mode: bool = False) -> None:
     """Get the status of a running model on the cluster.
 
     Parameters
@@ -338,14 +339,28 @@ def status(slurm_job_id: str, json_mode: bool = False) -> None:
     try:
         # Start the client and get model inference server status
         client = VecInfClient()
-        status_response = client.get_status(slurm_job_id)
+        if not slurm_job_id:
+            slurm_job_ids = client.fetch_running_jobs()
+            if not slurm_job_ids:
+                click.echo("No running jobs found.")
+                return
+        else:
+            slurm_job_ids = [slurm_job_id]
+        responses = []
+        for job_id in slurm_job_ids:
+            responses.append(client.get_status(job_id))
+
         # Display status information
-        status_formatter = StatusResponseFormatter(status_response)
-        if json_mode:
-            status_formatter.output_json()
+        if slurm_job_id:
+            status_formatter = StatusResponseFormatter(responses[0])
+            if json_mode:
+                status_formatter.output_json()
+            else:
+                status_info_table = status_formatter.output_table()
+                CONSOLE.print(status_info_table)
         else:
-            status_info_table = status_formatter.output_table()
-            CONSOLE.print(status_info_table)
+            list_status_display = ListStatusDisplay(slurm_job_ids, responses, json_mode)
+            list_status_display.display_multiple_status_output(CONSOLE)
 
     except click.ClickException as e:
         raise e

vec_inf/cli/_helper.py

@@ -251,6 +251,56 @@ def output_table(self) -> Table:
         return table
 
 
+class ListStatusDisplay:
+    """CLI Helper class for formatting a list of StatusResponse.
+
+    A formatter class that handles the presentation of multiple job statuses
+    in a table format.
+
+    Parameters
+    ----------
+    statuses : list[StatusResponse]
+        List of model status information
+    """
+
+    def __init__(self, job_ids: list[str], statuses: list[StatusResponse], json_mode: bool = False):
+        self.job_ids = job_ids
+        self.statuses = statuses
+        self.json_mode = json_mode
+        
+        self.table = Table(show_header=True,header_style="bold magenta")
+        self.table.add_column("Job ID")
+        self.table.add_column("Model Name")
+        self.table.add_column("Status", style="blue")
+        self.table.add_column("Base URL")
+
+    def display_multiple_status_output(self, console: Console) -> None:
+        """Format and display all model statuses.
+
+        Formats each model's status and adds it to the table.
+        """
+        if self.json_mode:
+            json_data = [
+                {
+                    "job_id": status.model_name,
+                    "model_name": status.model_name,
+                    "model_status": status.server_status,
+                    "base_url": status.base_url,
+                }
+                for status in self.statuses
+            ]
+            click.echo(json.dumps(json_data, indent=4))
+            return
+
+        for i, status in enumerate(self.statuses):
+            self.table.add_row(
+                self.job_ids[i],
+                status.model_name,
+                status.server_status,
+                status.base_url,
+            )
+        console.print(self.table)
+
 class MetricsResponseFormatter:
     """CLI Helper class for formatting MetricsResponse.