added clean up logs functionality to cli.

Author: kohankhaki

README.md

@@ -102,6 +102,7 @@ Note that there are other parameters that can also be added to the config but no
 * `metrics`: Streams performance metrics to the console.
 * `shutdown`: Shutdown a model by providing its Slurm job ID.
 * `list`: List all available model names, or view the default/cached configuration of a specific model, `--json-mode` supported.
+* `cleanup`: Remove old log directories. You can filter by `--model-family`, `--model-name`, and/or `--job-id`. Use `--dry-run` to preview what would be deleted.
 
 For more details on the usage of these commands, refer to the [User Guide](https://vectorinstitute.github.io/vector-inference/user_guide/)
 

vec_inf/README.md

@@ -5,5 +5,6 @@
 * `metrics`: Streams performance metrics to the console.
 * `shutdown`: Shutdown a model by providing its Slurm job ID.
 * `list`: List all available model names, or view the default/cached configuration of a specific model, `--json-mode` supported.
+* `cleanup`: Remove old log directories. You can filter by `--model-family`, `--model-name`, and/or `--job-id`. Use `--dry-run` to preview what would be deleted.
 
 Use `--help` to see all available options

vec_inf/cli/_cli.py

@@ -336,5 +336,60 @@ def metrics(slurm_job_id: int, log_dir: Optional[str] = None) -> None:
         raise click.ClickException(f"Metrics check failed: {str(e)}") from e
 
 
+@cli.command("cleanup")
+@click.option("--log-dir", type=str, help="Path to SLURM log directory")
+@click.option("--model-family", type=str, help="Filter by model family")
+@click.option("--model-name", type=str, help="Filter by model name")
+@click.option(
+    "--job-id", type=int, help="Only remove logs with this exact SLURM job ID"
+)
+@click.option("--dry-run", is_flag=True, help="List matching logs without deleting")
+def cleanup_logs_cli(
+    log_dir: Optional[str],
+    model_family: Optional[str],
+    model_name: Optional[str],
+    job_id: Optional[int],
+    dry_run: bool,
+) -> None:
+    """Clean up log files based on optional filters.
+
+    Parameters
+    ----------
+    log_dir : str or Path, optional
+        Root directory containing log files. Defaults to ~/.vec-inf-logs.
+    model_family : str, optional
+        Only delete logs for this model family.
+    model_name : str, optional
+        Only delete logs for this model name.
+    job_id : int, optional
+        If provided, only match directories with this exact SLURM job ID.
+    dry_run : bool
+        If True, return matching files without deleting them.
+    """
+    try:
+        client = VecInfClient()
+        matched = client.cleanup_logs(
+            log_dir=log_dir,
+            model_family=model_family,
+            model_name=model_name,
+            job_id=job_id,
+            dry_run=dry_run,
+        )
+
+        if not matched:
+            if dry_run:
+                click.echo("Dry run: no matching log directories found.")
+            else:
+                click.echo("No matching log directories were deleted.")
+        elif dry_run:
+            click.echo(f"Dry run: {len(matched)} directories would be deleted:")
+            for f in matched:
+                click.echo(f"  - {f}")
+        else:
+            click.echo(f"Deleted {len(matched)} log directory(ies).")
+    except Exception as e:
+        raise click.ClickException(f"Cleanup failed: {str(e)}") from e
+
+
 if __name__ == "__main__":
     cli()

vec_inf/client/_utils.py

@@ -285,3 +285,65 @@ def parse_launch_output(output: str) -> tuple[str, dict[str, str]]:
             config_dict[key.lower().replace(" ", "_")] = value
 
     return slurm_job_id, config_dict
+
+
+def find_matching_dirs(
+    log_dir: Path,
+    model_family: Optional[str] = None,
+    model_name: Optional[str] = None,
+    job_id: Optional[int] = None,
+) -> list[Path]:
+    """
+    Find log directories based on filtering criteria.
+
+    Parameters
+    ----------
+    log_dir : Path
+        The base directory containing model family directories.
+    model_family : str, optional
+        Filter to only search inside this family.
+    model_name : str, optional
+        Filter to only match model names.
+    job_id : int, optional
+        Filter to only match this exact SLURM job ID.
+
+    Returns
+    -------
+    list[Path]
+        List of directories that match the criteria and can be deleted.
+    """
+    matched = []
+
+    if not log_dir.exists() or not log_dir.is_dir():
+        raise FileNotFoundError(f"Log directory does not exist: {log_dir}")
+
+    if not model_family and not model_name and not job_id:
+        return [log_dir]
+
+    # Iterate over model families
+    for family_dir in log_dir.iterdir():
+        if not family_dir.is_dir():
+            continue
+        if model_family and family_dir.name != model_family:
+            continue
+
+        if model_family and not model_name and not job_id:
+            return [family_dir]
+
+        for job_dir in family_dir.iterdir():
+            if not job_dir.is_dir():
+                continue
+
+            try:
+                name_part, id_part = job_dir.name.rsplit(".", 1)
+                parsed_id = int(id_part)
+            except ValueError:
+                continue
+
+            if model_name and name_part != model_name:
+                continue
+            if job_id is not None and parsed_id != job_id:
+                continue
+
+            matched.append(job_dir)
+    return matched

vec_inf/client/api.py

@@ -10,8 +10,10 @@
 vec_inf.client.models : Data models for API responses
 """
 
+import shutil
 import time
 import warnings
+from pathlib import Path
 from typing import Any, Optional, Union
 
 from vec_inf.client._exceptions import (
@@ -24,7 +26,7 @@
     ModelStatusMonitor,
     PerformanceMetricsCollector,
 )
-from vec_inf.client._utils import run_bash_command
+from vec_inf.client._utils import find_matching_dirs, run_bash_command
 from vec_inf.client.config import ModelConfig
 from vec_inf.client.models import (
     LaunchOptions,
@@ -60,6 +62,9 @@ class VecInfClient:
     wait_until_ready(slurm_job_id, timeout_seconds, poll_interval_seconds, log_dir)
         Wait for a model to become ready
 
+    cleanup_logs(log_dir, model_name, model_family, job_id, dry_run)
+        Remove logs from the log directory.
+
     Examples
     --------
     >>> from vec_inf.api import VecInfClient
@@ -300,3 +305,47 @@ def wait_until_ready(
 
             # Wait before checking again
             time.sleep(poll_interval_seconds)
+
+    def cleanup_logs(
+        self,
+        log_dir: Optional[Union[str, Path]] = None,
+        model_family: Optional[str] = None,
+        model_name: Optional[str] = None,
+        job_id: Optional[int] = None,
+        dry_run: bool = False,
+    ) -> list[Path]:
+        """Remove logs from the log directory.
+
+        Parameters
+        ----------
+        log_dir : str or Path, optional
+            Root directory containing log files. Defaults to ~/.vec-inf-logs.
+        model_family : str, optional
+            Only delete logs for this model family.
+        model_name : str, optional
+            Only delete logs for this model name.
+        job_id : int, optional
+            If provided, only match directories with this exact SLURM job ID.
+        dry_run : bool
+            If True, return matching files without deleting them.
+
+        Returns
+        -------
+        list[Path]
+            List of deleted (or matched if dry_run) log file paths.
+        """
+        log_root = Path(log_dir) if log_dir else Path.home() / ".vec-inf-logs"
+        matched = find_matching_dirs(
+            log_dir=log_root,
+            model_family=model_family,
+            model_name=model_name,
+            job_id=job_id,
+        )
+
+        if dry_run:
+            return matched
+
+        for path in matched:
+            shutil.rmtree(path)
+
+        return matched