make helpers more general for updating bd

Author: memalhot

.github/workflows/precommit.yaml

@@ -11,7 +11,7 @@ jobs:
       - uses: actions/checkout@v5
       - uses: actions/setup-python@v6
         with:
-          python-version: "3.10"
+          python-version: "3.11"
       - name: Set up uv
         uses: astral-sh/setup-uv@v7
       - name: Install dependencies

batchtools/batchtools.py

@@ -32,6 +32,7 @@ def build_parser(self) -> None:
             "--verbose",
             "-v",
             action="count",
+            default=0,
             help="Increase verbosity of output",
         )
 

batchtools/bd.py

@@ -70,13 +70,13 @@ def run(args: argparse.Namespace):
                     if name not in found:
                         print(f"{name} is not a GPU job and cannot be deleted.")
                         continue
-                    oc_delete(name)
+                    oc_delete("job", name)
             else:
                 # case where user does not provide jobs to delete, delete all
                 print("No job names provided -> deleting all GPU workloads:\n")
                 for job in gpu_jobs:
                     name = job.model.metadata.name
-                    oc_delete(name)
+                    oc_delete("job", name)
 
         except oc.OpenShiftPythonException as e:
             sys.exit(f"Error occurred while deleting jobs: {e}")

batchtools/bps.py

@@ -11,14 +11,16 @@
 
 
 class ListPodsCommandArgs(argparse.Namespace):
-    verbose: int = 0
+    verbose: int | None = 0
     node_names: list[str] = []
 
 
 class ListPodsCommand(Command):
     """
-    List active GPU pods per node. By default prints only BUSY nodes.
-    With -v/--verbose, prints FREE for nodes seen with Running pods but 0 GPUs.
+    batchtools [-v] bps [-h] [node-name [node-name ...]]
+
+    List active GPU pods per node. By default prints only BUSY nodes. So, it all nodes are FREE, it will return empty nodes/
+    With -v/--verbose, prints FREE for nodes that have Running pods but 0 GPUs.
     """
 
     name: str = "bps"
@@ -35,43 +37,45 @@ def build_parser(cls, subparsers: SubParserFactory):
     @override
     def run(args: argparse.Namespace):
         args = cast(ListPodsCommandArgs, args)
+        v = args.verbose or 0  # treat None as 0
+
         try:
             with oc.timeout(120):
                 all_pods = oc.selector("pods", all_namespaces=True).objects()
 
             if args.node_names:
-                # get individual nodes without repeats
+                # filter to Running pods on requested nodes
                 node_set = set(args.node_names)
-                # Filter to Running pods on requested nodes
                 pods_for_nodes = [
                     p
                     for p in all_pods
                     if getattr(p.model.status, "phase", None) == "Running"
                     and (getattr(p.model.spec, "nodeName", None) or "") in node_set
                 ]
+
                 # Group by node
-                pods_by_node = defaultdict(list)
+                pods_by_node: dict[str, list] = defaultdict(list)
                 for p in pods_for_nodes:
                     n = getattr(p.model.spec, "nodeName", None) or ""
                     pods_by_node[n].append(p)
 
+                # Emit per requested node (only those requested)
                 for node in node_set:
-                    lines = summarize_gpu_pods(
-                        pods_by_node.get(node, []), args.verbose > 0
-                    )
-                    if not lines and args.verbose:
+                    lines = summarize_gpu_pods(pods_by_node.get(node, []), v > 0)
+                    if not lines and v > 0:
                         print(f"{node}: FREE")
                     else:
                         for ln in lines:
                             print(ln)
+
             else:
-                # One global summary over all Running pods
+                # Global summary over all Running pods
                 running = [
                     p
                     for p in all_pods
                     if getattr(p.model.status, "phase", None) == "Running"
                 ]
-                for ln in summarize_gpu_pods(running, args.verbose > 0):
+                for ln in summarize_gpu_pods(running, v > 0):
                     print(ln)
 
         except oc.OpenShiftPythonException as e:
@@ -105,7 +109,7 @@ def summarize_gpu_pods(pods, verbose: bool) -> list[str]:
         except Exception:
             continue
 
-    lines = []
+    lines: list[str] = []
     nodes = sorted(seen_nodes or totals.keys())
     for node in nodes:
         total = totals.get(node, 0)

batchtools/br.py

@@ -225,7 +225,7 @@ def run(args: argparse.Namespace):
 
         if args.job_delete and args.wait:
             print(f"RUNDIR: jobs/{job_name}")
-            oc_delete(job_name)
+            oc_delete("job", job_name)
         else:
             print(
                 f"User specified not to wait, or not to delete, so {job_name} must be deleted by user."
@@ -268,7 +268,7 @@ def log_job_output(job_name: str, *, wait: bool, timeout: int | None) -> None:
             if timeout and (time.monotonic() - start) > timeout:
                 print(f"Timeout waiting for pod {pod_name} to complete")
                 print(f"Deleting pod {pod_name}")
-                oc_delete(job_name)
+                oc_delete("job", job_name)
                 return
 
             # sleep to avoid hammering the server

batchtools/helpers.py

@@ -21,9 +21,9 @@ def pretty_print(pod: oc.APIObject) -> str:
     return formatted_logs
 
 
-def oc_delete(job_name: str) -> None:
+def oc_delete(obj_type: str, obj_name: str) -> None:
     try:
-        print(f"Deleting {job_name}")
-        oc.invoke("delete", ["job", job_name])
+        print(f"Deleting {obj_type}/{obj_name}")
+        oc.selector(f"{obj_type}/{obj_name}").delete()
     except oc.OpenShiftPythonException as e:
-        print(f"Error occurred while deleting job: {e}")
+        print(f"Error occurred while deleting {obj_type}/{obj_name}: {e}")

tests/test_bd.py

@@ -2,7 +2,6 @@
 import pytest
 from unittest import mock
 from contextlib import contextmanager
-
 import argparse
 
 from batchtools.bd import DeleteJobsCommand
@@ -12,84 +11,72 @@
 @pytest.fixture
 def jobs():
     return [
-        DictToObject(
-            {
-                "model": {"metadata": {"name": "job-job-1"}},
-            }
-        ),
-        DictToObject(
-            {
-                "model": {"metadata": {"name": "job-job-2"}},
-            }
-        ),
+        DictToObject({"model": {"metadata": {"name": "job-job-1"}}}),
+        DictToObject({"model": {"metadata": {"name": "job-job-2"}}}),
     ]
 
 
 @pytest.fixture
 def ignored_jobs():
-    return [
-        DictToObject(
-            {
-                "model": {"metadata": {"name": "ignored-1"}},
-            }
-        ),
-    ]
+    return [DictToObject({"model": {"metadata": {"name": "ignored-1"}}})]
 
 
 @pytest.fixture
 def failed_jobs():
-    return [
-        DictToObject(
-            {
-                "model": {"metadata": {"name": "job-job-1"}},
-            }
-        ),
-    ]
+    return [DictToObject({"model": {"metadata": {"name": "job-job-1"}}})]
 
 
 @contextmanager
-def patch_jobs_selector(jobs: list[DictToObject]):
+def patch_jobs_selector(job_list: list[DictToObject]):
+    """
+    Patches openshift_client.selector for BOTH:
+      - the list call: selector("job", labels=...).objects() -> job_list
+      - the delete calls: selector("job/<name>").delete()
+    """
     with mock.patch("openshift_client.selector") as mock_selector:
-        mock_result = mock.Mock(name="result")
-        mock_result.objects.return_value = jobs
-        mock_selector.return_value = mock_result
+        result = mock.Mock(name="selector_result")
+        result.objects.return_value = job_list
+        mock_selector.return_value = result
         yield mock_selector
 
 
 def test_no_jobs(args: argparse.Namespace, capsys):
     with patch_jobs_selector([]):
         DeleteJobsCommand.run(args)
-        captured = capsys.readouterr()
-        assert "No jobs found" in captured.out
+        out = capsys.readouterr().out
+        assert "No jobs found" in out
 
 
 def test_no_gpu_jobs(args: argparse.Namespace, ignored_jobs, capsys):
     with patch_jobs_selector(ignored_jobs):
         DeleteJobsCommand.run(args)
-        captured = capsys.readouterr()
-        assert "No GPU workloads to delete" in captured.out
+        out = capsys.readouterr().out
+        assert "No GPU workloads to delete" in out
 
 
-def test_delete_jobs(args: argparse.Namespace, jobs, capsys):
+def test_delete_obj(args: argparse.Namespace, jobs, capsys):
     args.job_names = []
-    with (
-        patch_jobs_selector(jobs),
-        mock.patch("openshift_client.invoke") as mock_invoke,
-    ):
+    with patch_jobs_selector(jobs) as mock_selector:
         DeleteJobsCommand.run(args)
-        captured = capsys.readouterr()
-        for job, ca in zip(jobs, mock_invoke.call_args_list):
-            assert f"Deleting {job.model.metadata.name}" in captured.out
-            assert ca.args == ("delete", ["job", job.model.metadata.name])
+        out = capsys.readouterr().out
+
+        for obj in jobs:
+            name = obj.model.metadata.name
+            assert f"Deleting job/{name}" in out
+
+        called_with = [c.args[0] for c in mock_selector.call_args_list if c.args]
+        for obj in jobs:
+            assert f"job/{obj.model.metadata.name}" in called_with
 
 
 def test_delete_jobs_fail(args: argparse.Namespace, failed_jobs, capsys):
     args.job_names = []
-    with (
-        patch_jobs_selector(failed_jobs),
-        mock.patch("openshift_client.invoke") as mock_invoke,
-    ):
-        mock_invoke.side_effect = OpenShiftPythonException("test exception")
+    with patch_jobs_selector(failed_jobs) as mock_selector:
+        mock_selector.return_value.delete.side_effect = OpenShiftPythonException(
+            "test exception"
+        )
+
         DeleteJobsCommand.run(args)
-        captured = capsys.readouterr()
-        assert "Error occurred while deleting job: test exception" in captured.out
+        out = capsys.readouterr().out
+
+        assert "Error occurred while deleting job/job-job-1: test exception" in out