Merge branch 'main' into dependabot/github_actions/docker/build-push-action-6.18.0

Author: XkunW

.pre-commit-config.yaml

@@ -17,7 +17,7 @@ repos:
     - id: check-toml
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: 'v0.11.11'
+    rev: 'v0.11.13'
     hooks:
     - id: ruff
       args: [--fix, --exit-non-zero-on-fix]
@@ -26,7 +26,7 @@ repos:
       types_or: [python, jupyter]
 
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.15.0
+    rev: v1.16.0
     hooks:
     - id: mypy
       entry: python3 -m mypy --config-file pyproject.toml

README.md

@@ -107,6 +107,7 @@ export VEC_INF_CONFIG=/h/<username>/my-model-config.yaml
 * `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`, `--job-id`, and/or `--before-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/)
 

tests/vec_inf/cli/test_cli.py

@@ -531,3 +531,77 @@ def test_metrics_command_request_failed(
             in result.output
         )
         assert "Connection refused" in result.output
+
+
+def test_cli_cleanup_logs_dry_run(runner, tmp_path):
+    """Test CLI cleanup command in dry-run mode."""
+    model_dir = tmp_path / "fam_a" / "model_a.123"
+    model_dir.mkdir(parents=True)
+
+    result = runner.invoke(
+        cli,
+        [
+            "cleanup",
+            "--log-dir",
+            str(tmp_path),
+            "--model-family",
+            "fam_a",
+            "--model-name",
+            "model_a",
+            "--dry-run",
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "would be deleted" in result.output
+    assert "model_a.123" in result.output
+
+
+def test_cli_cleanup_logs_delete(tmp_path):
+    """Test cleanup_logs CLI deletes matching directories when not in dry-run mode."""
+    fam_dir = tmp_path / "fam_a"
+    fam_dir.mkdir()
+    (fam_dir / "model_a.1").mkdir()
+
+    runner = CliRunner()
+    result = runner.invoke(
+        cli,
+        [
+            "cleanup",
+            "--log-dir",
+            str(tmp_path),
+            "--model-family",
+            "fam_a",
+            "--model-name",
+            "model_a",
+            "--job-id",
+            "1",
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "Deleted 1 log directory" in result.output
+    assert not (fam_dir / "model_a.1").exists()
+
+
+def test_cli_cleanup_logs_no_match(tmp_path):
+    """Test cleanup_logs CLI when no directories match the filters."""
+    fam_dir = tmp_path / "fam_a"
+    fam_dir.mkdir()
+    (fam_dir / "model_a.1").mkdir()
+
+    runner = CliRunner()
+    result = runner.invoke(
+        cli,
+        [
+            "cleanup",
+            "--log-dir",
+            str(tmp_path),
+            "--model-family",
+            "fam_b",
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "No matching log directories were deleted." in result.output
+    assert (fam_dir / "model_a.1").exists()

tests/vec_inf/client/test_api.py

@@ -131,6 +131,77 @@ def test_wait_until_ready():
             assert mock_status.call_count == 2
 
 
+def test_cleanup_logs_no_match(tmp_path):
+    """Test when cleanup_logs returns empty list."""
+    fam_a = tmp_path / "fam_a"
+    model_a = fam_a / "model_a.999"
+    model_a.mkdir(parents=True)
+
+    client = VecInfClient()
+    deleted = client.cleanup_logs(
+        log_dir=tmp_path,
+        model_family="fam_b",
+        dry_run=False,
+    )
+
+    assert deleted == []
+    assert fam_a.exists()
+    assert model_a.exists()
+
+
+def test_cleanup_logs_deletes_matching_dirs(tmp_path):
+    """Test that cleanup_logs deletes model directories matching filters."""
+    fam_a = tmp_path / "fam_a"
+    fam_a.mkdir()
+
+    model_a_1 = fam_a / "model_a.10"
+    model_a_2 = fam_a / "model_a.20"
+    model_b = fam_a / "model_b.30"
+
+    model_a_1.mkdir()
+    model_a_2.mkdir()
+    model_b.mkdir()
+
+    client = VecInfClient()
+    deleted = client.cleanup_logs(
+        log_dir=tmp_path,
+        model_family="fam_a",
+        model_name="model_a",
+        before_job_id=15,
+        dry_run=False,
+    )
+
+    assert deleted == [model_a_1]
+    assert not model_a_1.exists()
+    assert model_a_2.exists()
+    assert model_b.exists()
+
+
+def test_cleanup_logs_matching_dirs_dry_run(tmp_path):
+    """Test that cleanup_logs find model directories matching filters."""
+    fam_a = tmp_path / "fam_a"
+    fam_a.mkdir()
+
+    model_a_1 = fam_a / "model_a.10"
+    model_a_2 = fam_a / "model_a.20"
+
+    model_a_1.mkdir()
+    model_a_2.mkdir()
+
+    client = VecInfClient()
+    deleted = client.cleanup_logs(
+        log_dir=tmp_path,
+        model_family="fam_a",
+        model_name="model_a",
+        before_job_id=15,
+        dry_run=True,
+    )
+
+    assert deleted == [model_a_1]
+    assert model_a_1.exists()
+    assert model_a_2.exists()
+
+
 def test_shutdown_model_success():
     """Test model shutdown success."""
     client = VecInfClient()

tests/vec_inf/client/test_utils.py

@@ -8,6 +8,7 @@
 
 from vec_inf.client._utils import (
     MODEL_READY_SIGNATURE,
+    find_matching_dirs,
     get_base_url,
     is_server_running,
     load_config,
@@ -208,3 +209,135 @@ def test_load_config_invalid_user_model(tmp_path):
     assert "validation error" in str(excinfo.value).lower()
     assert "model_type" in str(excinfo.value)
     assert "num_gpus" in str(excinfo.value)
+
+
+def test_find_matching_dirs_only_model_family(tmp_path):
+    """Return model_family directory when only model_family is provided."""
+    fam_dir = tmp_path / "fam_a"
+    fam_dir.mkdir()
+    (fam_dir / "model_a.1").mkdir()
+    (fam_dir / "model_b.2").mkdir()
+
+    other_dir = tmp_path / "fam_b"
+    other_dir.mkdir()
+    (other_dir / "model_c.3").mkdir()
+
+    matches = find_matching_dirs(log_dir=tmp_path, model_family="fam_a")
+    assert len(matches) == 1
+    assert matches[0].name == "fam_a"
+
+
+def test_find_matching_dirs_only_model_name(tmp_path):
+    """Return directories matching when only model_name is provided."""
+    fam_a = tmp_path / "fam_a"
+    fam_a.mkdir()
+    (fam_a / "target.1").mkdir()
+    (fam_a / "other.2").mkdir()
+
+    fam_b = tmp_path / "fam_b"
+    fam_b.mkdir()
+    (fam_b / "different.3").mkdir()
+
+    matches = find_matching_dirs(log_dir=tmp_path, model_name="target")
+    result_names = [p.name for p in matches]
+
+    assert "target.1" in result_names
+    assert "other.2" not in result_names
+    assert "different.3" not in result_names
+
+
+def test_find_matching_dirs_only_job_id(tmp_path):
+    """Return directories matching exact job_id."""
+    fam_dir = tmp_path / "fam"
+    fam_dir.mkdir()
+    (fam_dir / "model_a.10").mkdir()
+    (fam_dir / "model_b.20").mkdir()
+    (fam_dir / "model_c.30").mkdir()
+
+    matches = find_matching_dirs(log_dir=tmp_path, job_id=10)
+    result_names = [p.name for p in matches]
+
+    assert "model_a.10" in result_names
+    assert "model_b.20" not in result_names
+    assert "model_c.30" not in result_names
+
+
+def test_find_matching_dirs_only_before_job_id(tmp_path):
+    """Return directories with job_id < before_job_id."""
+    fam_dir = tmp_path / "fam_a"
+    fam_dir.mkdir()
+    (fam_dir / "model_a.1").mkdir()
+    (fam_dir / "model_a.5").mkdir()
+    (fam_dir / "model_a.100").mkdir()
+
+    fam_dir = tmp_path / "fam_b"
+    fam_dir.mkdir()
+    (fam_dir / "model_b.30").mkdir()
+
+    matches = find_matching_dirs(log_dir=tmp_path, before_job_id=50)
+    result_names = [p.name for p in matches]
+
+    assert "model_a.1" in result_names
+    assert "model_a.5" in result_names
+    assert "model_a.100" not in result_names
+    assert "model_b.30" in result_names
+
+
+def test_find_matching_dirs_family_and_before_job_id(tmp_path):
+    """Return directories under a given family with job IDs less than before_job_id."""
+    fam_dir = tmp_path / "targetfam"
+    fam_dir.mkdir()
+    (fam_dir / "model_a.10").mkdir()
+    (fam_dir / "model_a.20").mkdir()
+    (fam_dir / "model_a.99").mkdir()
+    (fam_dir / "model_a.150").mkdir()
+
+    other_fam = tmp_path / "otherfam"
+    other_fam.mkdir()
+    (other_fam / "model_b.5").mkdir()
+    (other_fam / "model_b.10").mkdir()
+    (other_fam / "model_b.100").mkdir()
+
+    matches = find_matching_dirs(
+        log_dir=tmp_path,
+        model_family="targetfam",
+        before_job_id=100,
+    )
+
+    result_names = [p.name for p in matches]
+
+    assert "model_a.10" in result_names
+    assert "model_a.20" in result_names
+    assert "model_a.99" in result_names
+    assert "model_a.150" not in result_names
+    assert all("otherfam" not in str(p) for p in matches)
+
+
+def test_find_matching_dirs_with_family_model_name_and_before_job_id(tmp_path):
+    """Return matching dirs with model_family, model_name, and before_job_id filters."""
+    fam_dir = tmp_path / "targetfam"
+    fam_dir.mkdir()
+    (fam_dir / "model_a.1").mkdir()
+    (fam_dir / "model_a.50").mkdir()
+    (fam_dir / "model_a.150").mkdir()
+    (fam_dir / "model_b.40").mkdir()
+
+    other_fam = tmp_path / "otherfam"
+    other_fam.mkdir()
+    (other_fam / "model_c.20").mkdir()
+
+    matches = find_matching_dirs(
+        log_dir=tmp_path,
+        model_family="targetfam",
+        model_name="model_a",
+        before_job_id=100,
+    )
+
+    result_names = [p.name for p in matches]
+
+    assert "model_a.1" in result_names
+    assert "model_a.50" in result_names
+    assert "model_a.150" not in result_names
+    assert "model_b.40" not in result_names
+    assert all("model_b" not in p for p in result_names)
+    assert all("otherfam" not in str(p) for p in matches)

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`, `--job-id`, and/or `--before-job-id`. Use `--dry-run` to preview what would be deleted.
 
 Use `--help` to see all available options