fix kueue management checking

Author: memalhot

batchtools/bd.py

@@ -1,23 +1,26 @@
 import sys
 import argparse
 from typing import cast
-
 import openshift_client as oc
-
 from .basecommand import Command, override
 from .basecommand import SubParserFactory
-from .helpers import oc_delete
+from .helpers import oc_delete, is_kueue_managed_job
 
 
 class DeleteJobsCommand(Command):
     """
     batchtools bd [job-name [job-name ...]]
 
-    Delete specified Jobs, or all Jobs if none are specified.
+    Delete specified Kueue-managed GPU jobs, or all such jobs if none are specified.
+
+    Description:
+        Deletes only those Jobs that are both:
+          - named like your GPU jobs (name starts with 'job-'), and
+          - detected as Kueue-managed (via labels/Workload linkage).
     """
 
     name: str = "bd"
-    help: str = "Delete specified Jobs, or all if none are specified"
+    help: str = "Delete specified Kueue-managed GPU jobs, or all if none are specified"
 
     @classmethod
     @override
@@ -26,34 +29,41 @@ def build_parser(cls, subparsers: SubParserFactory):
         p.add_argument(
             "job_names",
             nargs="*",
-            help="Optional list of job names to delete",
+            help="Optional list of job names to delete (must be Kueue-managed)",
         )
         return p
 
     @staticmethod
     @override
     def run(args: argparse.Namespace):
         args = cast(DeleteJobsCommand, args)
-
         try:
             jobs = oc.selector("jobs").objects()
             if not jobs:
                 print("No jobs found.")
                 return
 
+            gpu_jobs = [job for job in jobs]
+
+            # only want to delete kueue jobs so filter for kueue jobs
+            kueue_gpu_jobs = [job for job in gpu_jobs if is_kueue_managed_job(job)]
+
+            if not kueue_gpu_jobs:
+                print("No Kueue-managed GPU jobs to delete.")
+                return
+
             if args.job_names:
-                # delete only specified jobs
-                existing = {job.model.metadata.name for job in jobs}
+                # if jobs are specified, only delete specified jobs
+                allowed = {job.model.metadata.name for job in kueue_gpu_jobs}
                 for name in args.job_names:
-                    if name not in existing:
-                        print(f"{name} does not exist; skipping.")
+                    if name not in allowed:
+                        print(f"{name} is not a Kueue-managed GPU job; skipping.")
                         continue
                     oc_delete("job", name)
                     print(f"Deleted job: {name}")
             else:
-                # delete all jobs
-                print("No job names provided -> deleting ALL jobs:\n")
-                for job in jobs:
+                print("No job names provided -> deleting all Kueue-managed GPU jobs:\n")
+                for job in kueue_gpu_jobs:
                     name = job.model.metadata.name
                     oc_delete("job", name)
                     print(f"Deleted job: {name}")

batchtools/bj.py

@@ -4,6 +4,7 @@
 import sys
 import openshift_client as oc
 
+from .helpers import is_kueue_managed_job
 from .basecommand import Command
 
 
@@ -35,8 +36,12 @@ def run(args: argparse.Namespace):
                 print("No jobs found.")
                 return
 
-            print(f"Found {len(jobs)} job(s):\n")
-            for job in jobs:
+            # filter only Kueue-managed jobs
+            managed = [job for job in jobs if is_kueue_managed_job(job)]
+
+            print(f"Found {len(managed)} job(s):\n")
+
+            for job in managed:
                 print(f"- {job.model.metadata.name}")
 
         except oc.OpenShiftPythonException as e:

batchtools/helpers.py

@@ -29,40 +29,11 @@ def oc_delete(obj_type: str, obj_name: str) -> None:
         print(f"Error occurred while deleting {obj_type}/{obj_name}: {e}")
 
 
-# def is_kueue_managed_job(job_obj) -> bool:
-# """
-# Returns True if the Job is managed by Kueue.
-# Checks:
-#   1) Job has label 'kueue.x-k8s.io/queue-name'
-#   2) A Workload exists that either:
-#      - has an ownerReference pointing to this Job, or
-#      - has label job-name=<job-name>
-# """
-# try:
-#     md = job_obj.model.metadata
-#     labels = getattr(md, "labels", {}) or {}
-#     if "kueue.x-k8s.io/queue-name" in labels:
-#         return True
-
-#     job_name = md.name
-#     try:
-#         workloads = oc.selector("workloads").objects()
-#     except oc.OpenShiftPythonException:
-#         workloads = []
-
-#     for wl in workloads:
-#         wl_md = wl.model.metadata
-#         owners = getattr(wl_md, "ownerReferences", []) or []
-#         for o in owners:
-#             if (
-#                 getattr(o, "kind", "") == "Job"
-#                 and getattr(o, "name", "") == job_name
-#             ):
-#                 return True
-#         wl_labels = getattr(wl_md, "labels", {}) or {}
-#         if wl_labels.get("job-name") == job_name:
-#             return True
-# except Exception:
-#     return False
-
-# return False
+def is_kueue_managed_job(job_obj) -> bool:
+    try:
+        md = job_obj.model.metadata
+        labels = getattr(md, "labels", {}) or {}
+        if "kueue.x-k8s.io/queue-name" in labels:
+            return True
+    except Exception:
+        return False

tests/test_bd.py

@@ -9,11 +9,18 @@
 
 
 @pytest.fixture
-def jobs():
+def kueue_jobs():
     return [
-        DictToObject({"model": {"metadata": {"name": "job-1"}}}),
-        DictToObject({"model": {"metadata": {"name": "job-2"}}}),
-        DictToObject({"model": {"metadata": {"name": "other-1"}}}),
+        DictToObject({"model": {"metadata": {"name": "job-job-1"}}}),
+        DictToObject({"model": {"metadata": {"name": "job-job-2"}}}),
+    ]
+
+
+@pytest.fixture
+def mixed_jobs():
+    return [
+        DictToObject({"model": {"metadata": {"name": "job-job-1"}}}),
+        DictToObject({"model": {"metadata": {"name": "ignored-1"}}}),
     ]
 
 
@@ -36,85 +43,107 @@ def patch_selector_with(job_list):
         yield mock_selector
 
 
+def patch_kueue_managed(*names_that_are_kueue):
+    """
+    Patch batchtools.bd.is_kueue_managed_job so ONLY the provided job names return True.
+
+    bd.py calls is_kueue_managed_job with a STRING job name, not an APIObject.
+    This predicate accepts either a str or an object and normalizes to the name.
+    """
+
+    def _predicate(arg):
+        if isinstance(arg, str):
+            name = arg
+        else:
+            # tolerate APIObject/DictToObject
+            name = getattr(
+                getattr(getattr(arg, "model", None), "metadata", None), "name", None
+            ) or getattr(arg, "name", None)
+        return name in names_that_are_kueue
+
+    return mock.patch("batchtools.bd.is_kueue_managed_job", side_effect=_predicate)
+
+
 def test_no_jobs_found(args, capsys):
     with patch_selector_with([]):
         DeleteJobsCommand.run(args)
         out = capsys.readouterr().out
         assert "No jobs found." in out
 
 
-def test_delete_all_when_no_names_given(args, jobs, capsys):
+def test_no_kueue_managed_gpu_jobs(args, kueue_jobs, capsys):
+    with patch_selector_with(kueue_jobs), patch_kueue_managed():
+        DeleteJobsCommand.run(args)
+        out = capsys.readouterr().out
+        assert "No Kueue-managed GPU jobs to delete." in out
+
+
+def test_ignores_non_gpu_named_jobs(args, mixed_jobs, capsys):
+    with patch_selector_with(mixed_jobs), patch_kueue_managed("job-job-1"):
+        DeleteJobsCommand.run(args)
+        out = capsys.readouterr().out
+        # delete only job-job-1
+        assert "Deleting job/job-job-1" in out
+        assert "Deleted job: job-job-1" in out
+        assert "ignored-1" not in out
+
+
+def test_delete_all_when_no_names_given(args, kueue_jobs, capsys):
     args.job_names = []  # explicit
-    with patch_selector_with(jobs):
+    with patch_selector_with(kueue_jobs), patch_kueue_managed("job-job-1", "job-job-2"):
         DeleteJobsCommand.run(args)
         out = capsys.readouterr().out
 
-        assert "No job names provided -> deleting ALL jobs:" in out
-        # oc_delete should print "Deleting job/<name>" and "Deleted job: <name>"
-        for obj in jobs:
+        assert "No job names provided -> deleting all Kueue-managed GPU jobs:" in out
+        for obj in kueue_jobs:
             name = obj.model.metadata.name
             assert f"Deleting job/{name}" in out
             assert f"Deleted job: {name}" in out
 
 
-def test_delete_only_specified_existing(args, jobs, capsys):
-    args.job_names = ["job-1", "job-2"]
-    with patch_selector_with(jobs):
+def test_delete_only_specified_allowed(args, kueue_jobs, capsys):
+    args.job_names = ["job-job-1", "job-job-2"]
+    with patch_selector_with(kueue_jobs), patch_kueue_managed("job-job-1"):
         DeleteJobsCommand.run(args)
         out = capsys.readouterr().out
 
-        # Only the two specified jobs should be deleted
-        assert "Deleting job/job-1" in out
-        assert "Deleted job: job-1" in out
-        assert "Deleting job/job-2" in out
-        assert "Deleted job: job-2" in out
-
-        # other-1 exists but is not listed; should not be deleted
-        assert "Deleting job/other-1" not in out
+        assert "Deleting job/job-job-1" in out
+        assert "Deleted job: job-job-1" in out
+        assert "job-job-2 is not a Kueue-managed GPU job; skipping." in out
+        assert "Deleting job/job-job-2" not in out
 
 
-def test_skips_nonexistent_names(args, jobs, capsys):
-    args.job_names = ["job-1", "does-not-exist"]
-    with patch_selector_with(jobs):
+def test_only_deletes_listed_names_even_if_more_kueue(args, kueue_jobs, capsys):
+    args.job_names = ["job-job-2"]
+    with patch_selector_with(kueue_jobs), patch_kueue_managed("job-job-1", "job-job-2"):
         DeleteJobsCommand.run(args)
         out = capsys.readouterr().out
 
-        # Existing job gets deleted
-        assert "Deleting job/job-1" in out
-        assert "Deleted job: job-1" in out
-
-        # Missing job is skipped with a message
-        assert "does-not-exist does not exist; skipping." in out
-        assert "Deleting job/does-not-exist" not in out
+        assert "Deleting job/job-job-2" in out
+        assert "Deleted job: job-job-2" in out
+        # make sure job-job-1 is not deleted implicitly
+        assert "Deleting job/job-job-1" not in out
 
 
 def test_delete_jobs_prints_error_when_delete_raises(args, capsys):
-    # We rely on oc_delete's behavior here: it should catch the exception and
-    # print "Error occurred while deleting job/<name>: <msg>"
-    jobs = [DictToObject({"model": {"metadata": {"name": "job-1"}}})]
-
-    with patch_selector_with(jobs) as mock_selector:
+    jobs = [DictToObject({"model": {"metadata": {"name": "job-job-1"}}})]
+    with patch_selector_with(jobs) as mock_selector, patch_kueue_managed("job-job-1"):
         mock_selector.return_value.delete.side_effect = OpenShiftPythonException(
             "test exception"
         )
 
         DeleteJobsCommand.run(args)
         out = capsys.readouterr().out
-
-        # From oc_delete
-        assert "Deleting job/job-1" in out
-        assert "Error occurred while deleting job/job-1: test exception" in out
+        assert "Error occurred while deleting job/job-job-1: test exception" in out
 
 
 def test_sys_exit_when_list_selector_raises(args):
-    # If listing jobs fails, DeleteJobsCommand should exit with an error
     with mock.patch(
         "openshift_client.selector",
         side_effect=OpenShiftPythonException("not successful"),
     ):
         with pytest.raises(SystemExit) as excinfo:
             DeleteJobsCommand.run(args)
-
         assert "Error occurred while deleting jobs: not successful" in str(
             excinfo.value
         )