merge with changes from main, updates to only manipulate kueue related jobs/pods

Author: memalhot

README.md

@@ -1,4 +1,4 @@
- python-batchtools
+# python-batchtools
 
 ## Overview
 
@@ -55,7 +55,7 @@ The br command is how to submit batchjobs. It submits code intended to run on GP
 
 Here's how to use thed br command:
 
-First write a CUDA program and compile it :D
+First write a CUDA program and compile it :D <br>
 Then to submit your CUDA program to the GPU node:
 
 ``` sh
@@ -79,9 +79,9 @@ Run without waiting for logs (for longer runs, similar to a more traditional bat
 ``` sh
 batchtools br --no-wait "./cuda_program"
 ```
-***WARNING***
+***WARNING***<br>
 If you run br with the --no-wait flag, it will not be cleaned up for you. You must delete it on your own by running `batchtools bd <job-name>` or `oc delete job <job-name>`
-But don't worry, running with --no-wait will give you a reminder to delete your jobs!
+But don't worry, running with --no-wait will give you a reminder to delete your jobs!<br>
 
 And if you need help or want to see more flas:
 

batchtools/bd.py

@@ -43,8 +43,10 @@ def run(args: argparse.Namespace):
                 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 jobs if is_kueue_managed_job(job)]
+            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.")

batchtools/bl.py

@@ -8,6 +8,7 @@
 from .basecommand import Command
 from .basecommand import SubParserFactory
 from .helpers import pretty_print
+from .helpers import is_kueue_managed_pod
 
 
 class LogsCommandArgs(argparse.Namespace):
@@ -63,10 +64,15 @@ def run(args: argparse.Namespace):
                     print(pretty_print(pod_dict[name]))
 
             else:
-                # case where user provides no args, print logs for all pods
+                printed_any = False
                 for name, pod in pod_dict.items():
-                    print(f"\nLogs for {name}:\n{'-' * 40}")
-                    print(pretty_print(pod))
+                    if is_kueue_managed_pod(pod):
+                        printed_any = True
+                        print(f"\nLogs for {name}:\n{'-' * 40}")
+                        print(pretty_print(pod))
+
+                if not printed_any:
+                    print("No Kueue-managed pods found")
 
         except oc.OpenShiftPythonException as e:
             sys.exit(f"Error occurred while retrieving logs: {e}")

batchtools/br.py

@@ -232,7 +232,7 @@ def run(args: argparse.Namespace):
             print(
                 f"User specified not to wait, or not to delete, so {job_name} must be deleted by user.\n"
                 f"You can do this by running:\n"
-                f"  bd {job_name} OR\n"
+                f"  batchtools bd {job_name} OR\n"
                 f"  oc delete job {job_name}"
             )
 

batchtools/helpers.py

@@ -49,6 +49,20 @@ def is_kueue_managed_job(job_obj) -> bool:
         return False
 
 
-# FOR PRINTING TIMES
 def fmt(x):
     return f"{x:.3f}s" if x is not None else "n/a"
+
+
+def is_kueue_managed_pod(pod) -> bool:
+    try:
+        owners = getattr(pod.model.metadata, "ownerReferences", []) or []
+        job_owner = next((o for o in owners if o.kind == "Job"), None)
+        if not job_owner:
+            return False
+
+        job_name = job_owner.name
+        job_obj = oc.selector(f"job/{job_name}").object()
+        return is_kueue_managed_job(job_obj)
+
+    except Exception:
+        return False

tests/test-bt.sh

@@ -0,0 +1,119 @@
+print_line() {
+   sleep 1
+   for (( i=0; i<50; i++ )); do
+       printf '-'
+   done
+   printf '\n'
+}
+
+echo testing bq
+batchtools bq
+print_line
+
+echo 'testing bps (verbose)'
+batchtools --verbose bps
+print_line
+
+echo 'testing bps (not verbose)'
+batchtools bps
+print_line
+
+echo 'testing br without gpu nowait'
+batchtools br --gpu none --no-wait --name mtest ‘sleep 10 && echo test’
+print_line
+
+echo 'testing bd w/ one job, no args'
+batchtools bd
+print_line
+
+echo 'testing br with gpu'
+cat <<EOF > hello_test.cu
+#include <stdio.h>
+
+
+__global__ void
+helloFromGPU(void)
+{
+   printf("Hello from GPU\\n");
+}
+
+
+int main(void)
+{
+   printf("Hello from CPU\\n");
+   helloFromGPU <<< 1, 3 >>>();
+   cudaDeviceSynchronize();
+   cudaDeviceReset();
+}
+EOF
+
+nvcc hello_test.cu -o hello
+
+batchtools br './hello' | tee log.txt
+
+print_line
+
+echo 'testing batchtools br no command, should return error'
+batchtools br
+
+print_line
+
+echo 'testing bp, bl, bd no specified jobs'
+batchtools br --gpu none --no-wait 'sleep 10'
+
+printf 'bp one job'
+printf '\n'
+
+batchtools bp
+printf '\n'
+
+printf 'bl one job'
+printf '\n'
+
+batchtools bl
+printf '\n'
+
+batchtools br --gpu none --no-wait 'sleep 10'
+
+printf 'bp two jobs'
+printf '\n'
+
+batchtools bp
+printf '\n'
+
+printf 'bl two jobs'
+printf '\n'
+
+batchtools bl
+
+printf 'bd two jobs, no args'
+batchtools bd
+
+print_line
+
+echo 'testing br with context'
+
+echo "Hello from CUDA file!" > input.txt
+
+cat <<'EOF' > readfile.cu
+#include <stdio.h>
+#include <cuda_runtime.h>
+
+__global__ void printFile(char *d) { printf("%s\n", d); }
+
+int main() {
+    FILE *f = fopen("input.txt", "r");
+    char t[128]; fgets(t, 128, f); fclose(f);
+
+    char *d; cudaMalloc(&d, 128);
+    cudaMemcpy(d, t, 128, cudaMemcpyHostToDevice);
+    printFile<<<1,1>>>(d);
+    cudaDeviceSynchronize();
+    cudaFree(d);
+}
+EOF
+
+nvcc readfile.cu -o readfile
+
+batchtools br './readfile'
+print_line

tests/test_bl.py

@@ -1,6 +1,7 @@
 import pytest
 from unittest import mock
 from contextlib import contextmanager
+import openshift_client as oc
 
 import argparse
 from typing import Any
@@ -41,7 +42,10 @@ def args() -> argparse.Namespace:
 
 @contextmanager
 def patch_pods_selector(pods: list[Any]):
-    with mock.patch("openshift_client.selector") as mock_selector:
+    with (
+        mock.patch("openshift_client.selector") as mock_selector,
+        mock.patch("batchtools.bl.is_kueue_managed_pod", return_value=True),
+    ):
         mock_result = mock.Mock(name="result")
         mock_result.objects.return_value = pods
         mock_selector.return_value = mock_result
@@ -75,3 +79,45 @@ def test_get_logs_selected(args: argparse.Namespace, pods: list[Any], capsys):
         captured = capsys.readouterr()
         assert "Logs for pod1" in captured.out
         assert "Logs for pod2" not in captured.out
+
+
+def test_no_kueue_managed_pods(args: argparse.Namespace, pods: list[Any], capsys):
+    with (
+        patch_pods_selector(pods),
+        mock.patch("batchtools.bl.is_kueue_managed_pod", return_value=False),
+    ):
+        args.pod_names = []
+        LogsCommand.run(args)
+        captured = capsys.readouterr()
+        assert "No Kueue-managed pods found" in captured.out
+
+
+def test_invalid_pod_name(args: argparse.Namespace, pods: list[Any], capsys):
+    with patch_pods_selector(pods):
+        args.pod_names = ["pod1", "does-not-exist"]
+        LogsCommand.run(args)
+        captured = capsys.readouterr()
+
+        # Valid pod still prints logs
+        assert "Logs for pod1" in captured.out
+        # Invalid pod hits the error path
+        assert (
+            "does-not-exist is not a valid pod. Logs cannot be retrieved."
+            in captured.out
+        )
+
+
+def test_selector_exception_exits_cleanly(args: argparse.Namespace):
+    # Make openshift_client.selector raise the same exception type used in bl.py
+    with mock.patch("openshift_client.selector") as mock_selector:
+        mock_selector.side_effect = oc.OpenShiftPythonException(
+            "test exception",
+            mock.Mock(),
+        )
+
+        with pytest.raises(SystemExit) as excinfo:
+            LogsCommand.run(args)
+
+        message = str(excinfo.value)
+        assert "Error occurred while retrieving logs:" in message
+        assert "test exception" in message