adding back br wait

Author: memalhot

batchtools/br.py

@@ -21,6 +21,7 @@
 from .prom_metrics import (
     PROMETHEUS_INSTANCE,
     IN_PROGRESS,
+    # record helpers (emit both histogram + counter)
     record_batch_observation,
     record_queue_observation,
     record_wall_observation,
@@ -37,7 +38,7 @@ class CreateJobCommandArgs(argparse.Namespace):
     name: str = "job"
     job_id: str = uuid.uuid5(uuid.NAMESPACE_OID, f"{os.getpid()}-{time.time()}").hex
     job_delete: bool = True
-    wait: bool = True
+    # no more wait flag exposed to users; br always waits for completion
     timeout: int = 60 * 15 * 4
     max_sec: int = 60 * 15
     gpu_numreq: int = 1
@@ -80,12 +81,7 @@ def build_parser(cls, subparsers: SubParserFactory):
             default=CreateJobCommandArgs.job_delete,
             help="Delete job on completion",
         )
-        p.add_argument(
-            "--wait",
-            action=argparse.BooleanOptionalAction,
-            default=CreateJobCommandArgs.wait,
-            help="Wait for job completion",
-        )
+        # removed --wait / --no-wait flag: br always waits for completion now
         p.add_argument(
             "--timeout",
             default=CreateJobCommandArgs.timeout,
@@ -185,11 +181,12 @@ 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
-                )
+            # always wait for the job to complete
+            result_phase, run_elapsed, queue_wait, total_wall = log_job_output(
+                job_name=job_name, wait=True, timeout=args.timeout
+            )
 
+            # Emit metrics if we captured any timing
             if (
                 run_elapsed is not None
                 or queue_wait is not None
@@ -220,21 +217,21 @@ def run(args: argparse.Namespace):
                 IN_PROGRESS.labels(**labels, result=result_phase).dec()
 
                 group = {
-                    "instance": PROMETHEUS_INSTANCE,  # oc project
-                    "job_name": job_name,  # unique name
+                    "instance": PROMETHEUS_INSTANCE,  # e.g. "ope-test"
+                    "job_name": job_name            # e.g. "job-none-b799d46a6f995a9d8d58b6aff76abefc"
                 }
 
                 push_registry_text(grouping_key=group)
 
         except oc.OpenShiftPythonException as e:
             sys.exit(f"Error occurred while creating job: {e}")
 
-        if args.job_delete and args.wait:
+        if args.job_delete:
             print(f"RUNDIR: jobs/{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.\n"
+                f"User specified not to delete, so {job_name} must be deleted by user.\n"
                 f"You can do this by running:\n"
                 f"  batchtools bd {job_name} OR\n"
                 f"  oc delete job {job_name}"
@@ -306,7 +303,6 @@ def log_job_output(
                 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

batchtools/prom_metrics.py

@@ -1,3 +1,4 @@
+import subprocess
 from datetime import datetime, timezone
 from .helpers import is_on_project
 
@@ -17,8 +18,9 @@
 
 LONG_JOB_BUCKETS = (1, 2, 5, 10, 20, 30, 60, 120, 180, 300, 600, 900, float("inf"))
 
-PUSHGATEWAY_ADDR = os.getenv("PUSHGATEWAY_ADDR", "pushgateway.ope-test.svc:9091")
-
+PUSHGATEWAY_ADDR = os.getenv(
+    "PUSHGATEWAY_ADDR", "pushgateway.ope-test.svc:9091"
+)
 
 def detect_instance() -> str:
     if shutil.which("oc") is None:
@@ -92,7 +94,6 @@ def detect_instance() -> str:
     registry=registry,
 )
 
-
 def now_rfc3339() -> str:
     return datetime.now(timezone.utc).isoformat()
 
@@ -140,7 +141,6 @@ def generate_metrics_text() -> tuple[str, str]:
     payload = generate_latest(registry)
     return payload.decode("utf-8"), CONTENT_TYPE_LATEST
 
-
 def push_registry_text(grouping_key: dict[str, str] | None = None) -> None:
     if not PUSHGATEWAY_ADDR:
         body, _ = generate_metrics_text()

pushgateway/prom.yaml

@@ -48,4 +48,4 @@ spec:
     - name: http
       port: 9091
       targetPort: 9091
-  type: ClusterIP
+  type: ClusterIP

tests/test_prom_metrics.py

@@ -1,17 +1,19 @@
 import batchtools.prom_metrics as pm
+import types
 from datetime import datetime
 from unittest import mock
 
 
 def test_now_rfc3339_parses_with_timezone():
     s = pm.now_rfc3339()
+    # ISO 8601 parseable and includes timezone info (+00:00)
     dt = datetime.fromisoformat(s)
     assert dt.tzinfo is not None
     assert s.endswith("+00:00")
 
 
 def _labels(job="job-x", gpu="none", queue="dummy-localqueue", instance="ns"):
-    return {"job_name": job, "gpu": gpu, "queue": queue, "instance": instance}
+    return {"job": job, "gpu": gpu, "queue": queue, "instance": instance}
 
 
 def _registry_text() -> str:
@@ -76,58 +78,61 @@ def test_generate_metrics_text_returns_valid_payload_and_ctype():
 
 
 def test_push_registry_text_no_url_prints_payload(capsys):
-    # check "no PUSHGATEWAY_ADDR" branch
-    with mock.patch.object(pm, "PUSHGATEWAY_ADDR", ""):
+    # Temporarily clear the push URL so it prints instead of POSTing
+    with mock.patch.object(pm, "PROMETHEUS_PUSH_URL", "", create=True):
         pm.push_registry_text()
-
     out = capsys.readouterr().out
-    assert "PROM: PUSHGATEWAY_ADDR not set" in out
-    assert "# HELP" in out
-    assert "# TYPE" in out
+    assert "PROMETHEUS_PUSH_URL not set" in out
+    assert "# HELP" in out  # the metrics text is printed
 
 
 def test_push_registry_text_posts_success(capsys):
     with (
-        mock.patch.object(pm, "PUSHGATEWAY_ADDR", "pushgateway.example:9091"),
-        mock.patch.object(pm, "pushadd_to_gateway") as mock_push,
+        mock.patch.object(
+            pm, "PROMETHEUS_PUSH_URL", "http://example/metrics", create=True
+        ),
+        mock.patch.object(pm.subprocess, "run") as mock_run,
     ):
-        pm.push_registry_text(grouping_key={"instance": "test-ns"})
+        mock_run.return_value = types.SimpleNamespace(returncode=0)
 
-    out = capsys.readouterr().out
-    assert "PROM: metrics pushed to pushgateway=pushgateway.example:9091" in out
+        pm.push_registry_text()
+        out = capsys.readouterr().out
+        assert "metrics successfully pushed" in out
 
-    # verify pushadd_to_gateway was called with expected args
-    mock_push.assert_called_once()
-    args, kwargs = mock_push.call_args
-    assert args[0] == "pushgateway.example:9091"
-    # job and registry are passed as keyword args
-    assert kwargs.get("job") == "batchtools"
-    assert kwargs.get("registry") is pm.registry
-    assert kwargs.get("grouping_key") == {"instance": "test-ns"}
+        # Verify we invoked curl with POST and sent our payload
+        assert mock_run.called
+        args, kwargs = mock_run.call_args
+        argv = args[0]
+        assert "curl" in argv[0]
+        assert "-X" in argv and "POST" in argv
+        assert "http://example/metrics" in argv
+
+        payload = kwargs.get("input", b"").decode("utf-8")
+        assert "# HELP" in payload
+        assert "# TYPE" in payload
 
 
 def test_push_registry_text_posts_failure(capsys):
     with (
-        mock.patch.object(pm, "PUSHGATEWAY_ADDR", "pushgateway.example:9091"),
-        mock.patch.object(pm, "pushadd_to_gateway", side_effect=RuntimeError("boom")),
+        mock.patch.object(
+            pm, "PROMETHEUS_PUSH_URL", "http://example/metrics", create=True
+        ),
+        mock.patch.object(pm.subprocess, "run") as mock_run,
     ):
-        pm.push_registry_text(grouping_key={"instance": "test-ns"})
-
-    out = capsys.readouterr().out
-    assert "PROM: failed to push metrics to pushgateway pushgateway.example:9091" in out
-    assert "boom" in out
+        mock_run.return_value = types.SimpleNamespace(returncode=7)
 
+        pm.push_registry_text()
+        out = capsys.readouterr().out
+        assert "curl returned nonzero exit 7" in out
 
-def test_push_registry_text_handles_generic_exception(capsys):
-    class WeirdError(Exception):
-        pass
 
+def test_push_registry_text_handles_exception(capsys):
     with (
-        mock.patch.object(pm, "PUSHGATEWAY_ADDR", "pushgateway.example:9091"),
-        mock.patch.object(pm, "pushadd_to_gateway", side_effect=WeirdError("weird")),
+        mock.patch.object(
+            pm, "PROMETHEUS_PUSH_URL", "http://example/metrics", create=True
+        ),
+        mock.patch.object(pm.subprocess, "run", side_effect=RuntimeError("boom")),
     ):
-        pm.push_registry_text(grouping_key={"instance": "test-ns"})
-
+        pm.push_registry_text()
     out = capsys.readouterr().out
-    assert "PROM: failed to push metrics to pushgateway pushgateway.example:9091" in out
-    assert "weird" in out
+    assert "failed to push metrics via curl" in out