scripts: make metrics gate() thread-safe (#14301)

Guard signal.signal registration to the main thread (it raises ValueError on
worker threads), and track live kubectl port-forward processes in a lock-guarded
registry exposing terminate_all_port_forwards for best-effort cleanup from a
main-thread signal handler. Enables gate() to run under parallel waits.

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: matanl-starkware

scripts/prod/metrics_lib.py

@@ -4,6 +4,7 @@
 import socket
 import subprocess
 import sys
+import threading
 import urllib.error
 import urllib.request
 from time import sleep
@@ -12,6 +13,29 @@
 from common_lib import Colors, get_namespace_args, print_colored, print_error
 from prometheus_client.parser import text_string_to_metric_families
 
+# Registry of live `kubectl port-forward` processes, so they can be cleaned up even when `gate()`
+# runs on a worker thread (where it cannot install its own signal handlers). Guarded by a lock since
+# parallel waits register/unregister concurrently.
+_active_port_forwards_lock = threading.Lock()
+_active_port_forwards: set[subprocess.Popen] = set()
+
+# Set when port-forwards are torn down by an interrupt handler, so a gate() running on a worker
+# thread can tell an intentional shutdown apart from a port-forward that failed to start.
+_shutdown_requested = threading.Event()
+
+
+def terminate_all_port_forwards() -> None:
+    """Terminate every port-forward process still running.
+
+    Best-effort cleanup intended to be called from a SIGINT/SIGTERM handler installed on the main
+    thread by the orchestrator when waits run in parallel (worker threads cannot install handlers).
+    """
+    _shutdown_requested.set()
+    with _active_port_forwards_lock:
+        port_forward_processes = list(_active_port_forwards)
+    for port_forward_process in port_forward_processes:
+        MetricConditionGater._terminate_port_forward_process(port_forward_process)
+
 
 class MetricConditionGater:
     """Gates progress on a metric satisfying a condition.
@@ -123,6 +147,8 @@ def _terminate_port_forward_process(pf_process: subprocess.Popen):
                 print_colored("Force killing kubectl port-forward process")
                 pf_process.kill()
                 pf_process.wait()
+        with _active_port_forwards_lock:
+            _active_port_forwards.discard(pf_process)
 
     def gate(self):
         """Wait until the nodes metrics satisfy the condition."""
@@ -142,25 +168,39 @@ def gate(self):
 
         try:
             pf_process = subprocess.Popen(cmd, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
-            print("Waiting for forwarding to start")
+            with _active_port_forwards_lock:
+                _active_port_forwards.add(pf_process)
+            print_colored("Waiting for forwarding to start")
             # Give the forwarding time to start.
             # TODO(guy.f): Consider poll until the forwarding is ready if we see any issues.
             sleep(3)
-            assert (
-                pf_process.poll() is None
-            ), f"Port forwarding process exited with code {pf_process.returncode}"
+            if pf_process.poll() is not None:
+                if _shutdown_requested.is_set():
+                    # An interrupt handler tore down the port-forward during startup; exit cleanly
+                    # rather than treating the intentional shutdown as an unexpected failure.
+                    return
+                raise RuntimeError(
+                    f"Port forwarding process for pod {self.pod} exited with code "
+                    f"{pf_process.returncode}"
+                )
 
-            print(
+            print_colored(
                 f"Forwarding started (from local port {self.local_port} to {self.pod}:{self.metrics_port})"
             )
 
-            # Set up signal handler to ensure forwarding subprocess is terminated on interruption
-            def signal_handler(signum, frame):
-                self._terminate_port_forward_process(pf_process)
-                sys.exit(0)
+            # Set up a signal handler to terminate the forwarding subprocess on interruption.
+            # signal.signal only works on the main thread, so skip it when gating runs on a worker
+            # thread (parallel waits) — the orchestrator installs a main-thread handler that calls
+            # terminate_all_port_forwards instead, and the finally block below still cleans up on
+            # normal completion.
+            if threading.current_thread() is threading.main_thread():
+
+                def signal_handler(signum, frame):
+                    self._terminate_port_forward_process(pf_process)
+                    sys.exit(0)
 
-            signal.signal(signal.SIGINT, signal_handler)
-            signal.signal(signal.SIGTERM, signal_handler)
+                signal.signal(signal.SIGINT, signal_handler)
+                signal.signal(signal.SIGTERM, signal_handler)
 
             self._poll_until_condition_met()