Refactor CLI to use client

Author: XkunW

vec_inf/cli/_cli.py

@@ -7,12 +7,12 @@
 from rich.console import Console
 from rich.live import Live
 
-import vec_inf.shared._utils as utils
+import vec_inf.client._utils as utils
 from vec_inf.cli._helper import (
-    CLILaunchHelper,
-    CLIListHelper,
-    CLIMetricsHelper,
-    CLIStatusHelper,
+    CLIMetricsCollector,
+    CLIModelLauncher,
+    CLIModelRegistry,
+    CLIModelStatusMonitor,
 )
 
 
@@ -131,14 +131,15 @@ def launch(
 ) -> None:
     """Launch a model on the cluster."""
     try:
-        launch_helper = CLILaunchHelper(model_name, cli_kwargs)
-
-        launch_helper.set_env_vars()
-        launch_command = launch_helper.build_launch_command()
-        command_output, stderr = utils.run_bash_command(launch_command)
-        if stderr:
-            raise click.ClickException(f"Error: {stderr}")
-        launch_helper.post_launch_processing(command_output, CONSOLE)
+        model_launcher = CLIModelLauncher(model_name, cli_kwargs)
+        # Launch model inference server
+        model_launcher.launch()
+        # Display launch information
+        if cli_kwargs.get("json_mode"):
+            click.echo(model_launcher.params)
+        else:
+            launch_info_table = model_launcher.format_table_output()
+            CONSOLE.print(launch_info_table)
 
     except click.ClickException as e:
         raise e
@@ -163,18 +164,14 @@ def status(
 ) -> None:
     """Get the status of a running model on the cluster."""
     try:
-        status_cmd = f"scontrol show job {slurm_job_id} --oneliner"
-        output, stderr = utils.run_bash_command(status_cmd)
-        if stderr:
-            raise click.ClickException(f"Error: {stderr}")
-
-        status_helper = CLIStatusHelper(slurm_job_id, output, log_dir)
-
-        status_helper.process_job_state()
+        # Get model inference server status
+        model_status_monitor = CLIModelStatusMonitor(slurm_job_id, log_dir)
+        model_status_monitor.process_model_status()
+        # Display status information
         if json_mode:
-            status_helper.output_json()
+            model_status_monitor.output_json()
         else:
-            status_helper.output_table(CONSOLE)
+            model_status_monitor.output_table(CONSOLE)
 
     except click.ClickException as e:
         raise e
@@ -200,8 +197,8 @@ def shutdown(slurm_job_id: int) -> None:
 def list_models(model_name: Optional[str] = None, json_mode: bool = False) -> None:
     """List all available models, or get default setup of a specific model."""
     try:
-        list_helper = CLIListHelper(json_mode)
-        list_helper.process_list_command(CONSOLE, model_name)
+        model_registry = CLIModelRegistry(json_mode)
+        model_registry.process_list_command(CONSOLE, model_name)
     except click.ClickException as e:
         raise e
     except Exception as e:
@@ -216,28 +213,28 @@ def list_models(model_name: Optional[str] = None, json_mode: bool = False) -> No
 def metrics(slurm_job_id: int, log_dir: Optional[str] = None) -> None:
     """Stream real-time performance metrics from the model endpoint."""
     try:
-        metrics_helper = CLIMetricsHelper(slurm_job_id, log_dir)
+        metrics_collector = CLIMetricsCollector(slurm_job_id, log_dir)
 
         # Check if metrics URL is ready
-        if not metrics_helper.metrics_url.startswith("http"):
+        if not metrics_collector.metrics_url.startswith("http"):
             table = utils.create_table("Metric", "Value")
-            metrics_helper.display_failed_metrics(
+            metrics_collector.display_failed_metrics(
                 table,
-                f"Metrics endpoint unavailable or server not ready - {metrics_helper.metrics_url}",
+                f"Metrics endpoint unavailable or server not ready - {metrics_collector.metrics_url}",
             )
             CONSOLE.print(table)
             return
 
         with Live(refresh_per_second=1, console=CONSOLE) as live:
             while True:
-                metrics = metrics_helper.fetch_metrics()
+                metrics = metrics_collector.fetch_metrics()
                 table = utils.create_table("Metric", "Value")
 
                 if isinstance(metrics, str):
                     # Show status information if metrics aren't available
-                    metrics_helper.display_failed_metrics(table, metrics)
+                    metrics_collector.display_failed_metrics(table, metrics)
                 else:
-                    metrics_helper.display_metrics(table, metrics)
+                    metrics_collector.display_metrics(table, metrics)
 
                 live.update(table)
                 time.sleep(2)

vec_inf/cli/_helper.py

@@ -11,13 +11,19 @@
 from rich.panel import Panel
 from rich.table import Table
 
-import vec_inf.shared._utils as utils
-from vec_inf.shared._config import ModelConfig
-from vec_inf.shared._helper import LaunchHelper, ListHelper, MetricsHelper, StatusHelper
+import vec_inf.client._utils as utils
+from vec_inf.cli._models import MODEL_TYPE_COLORS, MODEL_TYPE_PRIORITY
+from vec_inf.client._config import ModelConfig
+from vec_inf.client._helper import (
+    ModelLauncher,
+    ModelRegistry,
+    ModelStatusMonitor,
+    PerformanceMetricsCollector,
+)
 
 
-class CLILaunchHelper(LaunchHelper):
-    """CLI Helper class for handling launch information."""
+class CLIModelLauncher(ModelLauncher):
+    """CLI Helper class for handling inference server launch."""
 
     def __init__(self, model_name: str, kwargs: Optional[dict[str, Any]]):
         super().__init__(model_name, kwargs)
@@ -26,12 +32,12 @@ def _warn(self, message: str) -> None:
         """Warn the user about a potential issue."""
         click.echo(click.style(f"Warning: {message}", fg="yellow"), err=True)
 
-    def _format_table_output(self, job_id: str) -> Table:
+    def format_table_output(self) -> Table:
         """Format output as rich Table."""
         table = utils.create_table(key_title="Job Config", value_title="Value")
 
         # Add key information with consistent styling
-        table.add_row("Slurm Job ID", job_id, style="blue")
+        table.add_row("Slurm Job ID", self.slurm_job_id, style="blue")
         table.add_row("Job Name", self.model_name)
 
         # Add model details
@@ -71,33 +77,12 @@ def _format_table_output(self, job_id: str) -> Table:
 
         return table
 
-    def post_launch_processing(self, output: str, console: Console) -> None:
-        """Process and display launch output."""
-        json_mode = bool(self.kwargs.get("json_mode", False))
-        slurm_job_id = output.split(" ")[-1].strip().strip("\n")
-        self.params["slurm_job_id"] = slurm_job_id
-        job_json = Path(
-            self.params["log_dir"],
-            f"{self.model_name}.{slurm_job_id}",
-            f"{self.model_name}.{slurm_job_id}.json",
-        )
-        job_json.parent.mkdir(parents=True, exist_ok=True)
-        job_json.touch(exist_ok=True)
-
-        with job_json.open("w") as file:
-            json.dump(self.params, file, indent=4)
-        if json_mode:
-            click.echo(self.params)
-        else:
-            table = self._format_table_output(slurm_job_id)
-            console.print(table)
 
+class CLIModelStatusMonitor(ModelStatusMonitor):
+    """CLI Helper class for handling server status information and monitoring."""
 
-class CLIStatusHelper(StatusHelper):
-    """CLI Helper class for handling status information."""
-
-    def __init__(self, slurm_job_id: int, output: str, log_dir: Optional[str] = None):
-        super().__init__(slurm_job_id, output, log_dir)
+    def __init__(self, slurm_job_id: int, log_dir: Optional[str] = None):
+        super().__init__(slurm_job_id, log_dir)
 
     def output_json(self) -> None:
         """Format and output JSON data."""
@@ -127,7 +112,7 @@ def output_table(self, console: Console) -> None:
         console.print(table)
 
 
-class CLIMetricsHelper(MetricsHelper):
+class CLIMetricsCollector(PerformanceMetricsCollector):
     """CLI Helper class for streaming metrics information."""
 
     def __init__(self, slurm_job_id: int, log_dir: Optional[str] = None):
@@ -204,8 +189,8 @@ def display_metrics(self, table: Table, metrics: dict[str, float]) -> None:
         )
 
 
-class CLIListHelper(ListHelper):
-    """Helper class for handling model listing functionality."""
+class CLIModelRegistry(ModelRegistry):
+    """CLI Helper class for handling model listing functionality."""
 
     def __init__(self, json_mode: bool = False):
         super().__init__()
@@ -237,28 +222,15 @@ def format_all_models_output(self) -> Union[list[str], list[Panel]]:
             return [config.model_name for config in self.model_configs]
 
         # Sort by model type priority
-        type_priority = {
-            "LLM": 0,
-            "VLM": 1,
-            "Text_Embedding": 2,
-            "Reward_Modeling": 3,
-        }
         sorted_configs = sorted(
             self.model_configs,
-            key=lambda x: type_priority.get(x.model_type, 4),
+            key=lambda x: MODEL_TYPE_PRIORITY.get(x.model_type, 4),
         )
 
         # Create panels with color coding
-        model_type_colors = {
-            "LLM": "cyan",
-            "VLM": "bright_blue",
-            "Text_Embedding": "purple",
-            "Reward_Modeling": "bright_magenta",
-        }
-
         panels = []
         for config in sorted_configs:
-            color = model_type_colors.get(config.model_type, "white")
+            color = MODEL_TYPE_COLORS.get(config.model_type, "white")
             variant = config.model_variant or ""
             display_text = f"[magenta]{config.model_family}[/magenta]"
             if variant:

vec_inf/cli/_models.py

@@ -0,0 +1,15 @@
+"""Data models for CLI rendering."""
+
+MODEL_TYPE_PRIORITY = {
+    "LLM": 0,
+    "VLM": 1,
+    "Text_Embedding": 2,
+    "Reward_Modeling": 3,
+}
+
+MODEL_TYPE_COLORS = {
+    "LLM": "cyan",
+    "VLM": "bright_blue",
+    "Text_Embedding": "purple",
+    "Reward_Modeling": "bright_magenta",
+}