have metrics pushing also work for --no-wait jobs

Author: memalhot

batchtools/br.py

@@ -1,4 +1,3 @@
-# pyright: reportUninitializedInstanceVariable=false
 from typing import cast
 from typing_extensions import override, Optional
 
@@ -185,10 +184,9 @@ def run(args: argparse.Namespace):
             queue_wait = None
             total_wall = None
 
-            if args.wait:
-                result_phase, run_elapsed, queue_wait, total_wall = log_job_output(
-                    job_name=job_name, wait=True, timeout=args.timeout
-                )
+            result_phase, run_elapsed, queue_wait, total_wall = log_job_output(
+                job_name=job_name, wait=True, timeout=args.timeout
+            )
 
             if (
                 run_elapsed is not None
@@ -270,48 +268,48 @@ def log_job_output(
     queue_wait = None
     total_wall = None
 
-    if wait:
-        start_poll = time.monotonic()
-        while True:
-            phase = get_pod_status(pod_name)
-            if phase == "Running" and run_start is None:
-                # time waiting in queue is time from entering the queue to the time it takes to start running
-                run_start = time.monotonic()
-                queue_wait = run_start - start_poll  # submit -> running
-
-            if phase in ("Succeeded", "Failed"):
-                result_phase = phase.lower()
-                total_wall = time.monotonic() - start_poll  # submit -> terminal
-                print(f"Pod {pod_name} finished with phase={phase}")
-                break
-
-            if timeout and (time.monotonic() - start_poll) > timeout:
+    start_poll = time.monotonic()
+    while True:
+        phase = get_pod_status(pod_name)
+        if phase == "Running" and run_start is None:
+            # time waiting in queue is time from entering the queue to the time it takes to start running
+            run_start = time.monotonic()
+            queue_wait = run_start - start_poll  # submit -> running
+
+        if phase in ("Succeeded", "Failed"):
+            result_phase = phase.lower()
+            total_wall = time.monotonic() - start_poll  # submit -> terminal
+            print(f"Pod {pod_name} finished with phase={phase}")
+            break
+
+        if timeout and (time.monotonic() - start_poll) > timeout:
+            if wait:
                 print(f"Timeout waiting for pod {pod_name} to complete")
                 print(f"Deleting job {job_name}")
                 oc_delete("job", job_name)
-                total_wall = time.monotonic() - start_poll
-                # timeout: no run duration (didn't finish), queue_wait may or may not be set
-                print_timing(queue_wait, None, total_wall)
-                return ("timeout", None, queue_wait, total_wall)
+            total_wall = time.monotonic() - start_poll
+            # timeout: no run duration (didn't finish), queue_wait may or may not be set
+            print_timing(queue_wait, None, total_wall)
+            return ("timeout", None, queue_wait, total_wall)
 
-            time.sleep(2)
+        time.sleep(2)
 
     print(pretty_print(pod))
 
     # compute the runtime using the total time (total_wall)- time waiting in queue (queue_wait)
-    if wait:
-        if run_start is not None:
-            # running status was reached, run_elapsed = terminal - run_start
-            if total_wall is None:
-                total_wall = time.monotonic() - start_poll  # fallback
-            run_elapsed = total_wall - (queue_wait or 0.0)
-        else:
-            # Never reached Running; keep convention for failures:
-            run_elapsed = 0.0 if result_phase == "failed" else None
-            if total_wall is None:
-                total_wall = 0.0
-            if total_wall is not None and queue_wait is None:
-                queue_wait = total_wall
+
+    if run_start is not None:
+        # running status was reached, run_elapsed = terminal - run_start
+        if total_wall is None:
+            total_wall = time.monotonic() - start_poll  # fallback
+        run_elapsed = total_wall - (queue_wait or 0.0)
+    else:
+        # Never reached Running; keep convention for failures:
+        run_elapsed = 0.0 if result_phase == "failed" else None
+        if total_wall is None:
+            total_wall = 0.0
+        if total_wall is not None and queue_wait is None:
+            queue_wait = total_wall
 
     print_timing(queue_wait, run_elapsed, total_wall)
     return (result_phase, run_elapsed, queue_wait, total_wall)

tests/test_br.py

@@ -44,6 +44,7 @@ def test_invalid_gpu(args: argparse.Namespace):
         ),
     ],
 )
+@mock.patch("batchtools.br.log_job_output", return_value=("succeeded", 1.0, 0.5, 1.5))
 @mock.patch("openshift_client.create", name="create")
 @mock.patch("openshift_client.selector", name="selector")
 @mock.patch("socket.gethostname", name="gethostname")
@@ -53,13 +54,19 @@ def test_create_job_nowait(
     mock_gethostname,
     mock_selector,
     mock_create,
+    mock_log_job_output,
     gpu,
     resources,
     args: argparse.Namespace,
     tempdir,
     parser,
     subparsers,
 ):
+    """
+    Even if args.wait is False, CreateJobCommand.run should still build the
+    correct Job object and call oc.create() with it. We stub out
+    log_job_output so this test does not depend on pod selectors or timing.
+    """
     CreateJobCommand.build_parser(subparsers)
     args = parser.parse_args(["br"])
     args.wait = False
@@ -82,7 +89,7 @@ def test_create_job_nowait(
         }
     )
 
-    mock_result = mock.Mock(spec=["object"])
+    mock_result = mock.Mock()
     mock_result.object.return_value = pod
     mock_selector.return_value = mock_result
 
@@ -121,10 +128,17 @@ def test_create_job_nowait(
         },
     }
 
+    # Make the rsync_script simple and deterministic for the test
     batchtools.build_yaml.rsync_script = "testcommand {cmdline}"
+
     CreateJobCommand.run(args)
 
+    # Verify we created the expected Job spec
     assert mock_create.call_args.args[0] == expected
+    # And that we did call log_job_output once (even with wait=False)
+    mock_log_job_output.assert_called_once()
+    called_job_name = mock_log_job_output.call_args.kwargs["job_name"]
+    assert called_job_name == f"job-{gpu}-test"
 
 
 @mock.patch("openshift_client.create", name="create")