Improve cleanup and interrupt handling

pathwaysutils/experimental/shared_pathways_service/gke_utils.py

@@ -1,5 +1,6 @@
 """GKE utils for deploying and managing the Pathways proxy."""
 
+from collections.abc import Callable
 import json
 import logging
 import re
@@ -292,13 +293,17 @@ def enable_port_forwarding(
     remote_server: str,
     server_port: int,
     namespace: str = "default",
+    process_callback: Callable[[subprocess.Popen[str]], None] | None = None,
 ) -> tuple[int, subprocess.Popen[str]]:
   """Enables port forwarding for the given pod.
 
   Args:
     remote_server: The name of the pod or service.
     server_port: The port of the server to forward to.
     namespace: The namespace of the pod.
+    process_callback: Optional callback invoked immediately after the port
+      forwarding process is started. Useful for registering the process for
+      cleanup before waiting for it to become ready.
 
   Returns:
     A tuple containing the pod port and the port forwarding process.
@@ -350,6 +355,9 @@ def enable_port_forwarding(
     _logger.exception("Error enabling port forwarding for the pod: %r", e)
     raise
 
+  if process_callback is not None:
+    process_callback(port_forward_process)
+
   # Check that the port forwarding is ready.
   if port_forward_process.stdout is None:
     _logger.error("Port-forward process stdout is None. Terminating.")
@@ -377,8 +385,24 @@ def enable_port_forwarding(
 
   try:
     _test_remote_connection(local_port)
-  except Exception:
-    port_forward_process.terminate()
+  except BaseException:
+    _logger.warning(
+        "Terminating port forwarding process due to connection error or "
+        "interrupt."
+    )
+    try:
+      port_forward_process.terminate()
+      port_forward_process.wait(timeout=5)
+    except subprocess.TimeoutExpired:
+      port_forward_process.kill()
+      port_forward_process.wait(timeout=5)
+    except Exception as term_err:  # pylint: disable=broad-exception-caught
+      _logger.exception(
+          "Failed to terminate port forward process: %r. Please terminate it "
+          "manually by killing the process with PID %d.",
+          term_err,
+          port_forward_process.pid,
+      )
     raise
 
   return (local_port, port_forward_process)

pathwaysutils/experimental/shared_pathways_service/isc_pathways.py

@@ -1,14 +1,17 @@
 """Module for connecting to a Pathways server for interactive supercomputing."""
 
-from collections.abc import Iterable, Iterator, Mapping, Sequence
+import atexit
+from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
 import contextlib
 import dataclasses
 import gc
 import logging
 import os
 import random
+import signal
 import string
 import subprocess
+import sys
 import threading
 import time
 from typing import Any
@@ -21,6 +24,10 @@
 from pathwaysutils.experimental.shared_pathways_service import validators
 
 
+_CLEANUP_SIGNALS = [signal.SIGTERM, signal.SIGINT]
+if hasattr(signal, "SIGHUP"):  # SIGHUP is not available on Windows.
+  _CLEANUP_SIGNALS.append(signal.SIGHUP)
+
 PROXY_FILEPATH = os.path.join(
     os.path.dirname(__file__), "yamls/pw-proxy.yaml"
 )
@@ -178,10 +185,13 @@ def _wait_for_placement(
     metrics_collector_inst: Any = None,
     start_time: float | None = None,
     total_chips: int = 0,
+    process_callback: Callable[[subprocess.Popen[str]], None] | None = None,
 ) -> None:
   """Waits for the placement to be complete by checking proxy logs."""
   _logger.info("Streaming proxy logs until the placement is complete...")
   with stream_logs_func(pod_name) as log_process:
+    if process_callback is not None:
+      process_callback(log_process)
     keywords = [
         "placement",
         "Signaling to RM",
@@ -280,6 +290,7 @@ def __init__(
     self._proxy_job_name = proxy_job_name
     self.proxy_pod_name: str = ""
     self._port_forward_process = None
+    self._log_process = None
     self._proxy_port = None
     self.proxy_server_image = proxy_server_image
     self.proxy_options = proxy_options or ProxyOptions()
@@ -299,6 +310,42 @@ def __init__(
     self._old_jax_platforms_config = None
     self._old_jax_backend_target_config = None
     self.total_chips = self._get_total_chips()
+    self._cleaned_up = False
+    self._cleanup_lock = threading.Lock()
+    self._original_signal_handlers: dict[signal.Signals, Any] = {}
+
+  def _register_signal_handlers(self) -> None:
+    """Registers signal handlers to ensure cleanup on termination."""
+    if threading.current_thread() is not threading.main_thread():
+      return
+
+    def _handle_signal(signum, frame):
+      del frame
+      _logger.warning(
+          "Received signal %d. Triggering Pathways proxy cleanup...", signum
+      )
+      self._cleanup()
+      if signum == signal.SIGINT:
+        raise KeyboardInterrupt()
+      sys.exit(128 + signum)
+
+    for sig in _CLEANUP_SIGNALS:
+      try:
+        self._original_signal_handlers[sig] = signal.signal(sig, _handle_signal)
+      except (ValueError, OSError) as e:
+        _logger.debug("Could not register handler for signal %s: %s", sig, e)
+
+  def _restore_signal_handlers(self) -> None:
+    """Restores original signal handlers."""
+    if threading.current_thread() is not threading.main_thread():
+      return
+
+    for sig, original_handler in self._original_signal_handlers.items():
+      try:
+        signal.signal(sig, original_handler)
+      except (ValueError, OSError) as e:
+        _logger.debug("Could not restore handler for signal %s: %s", sig, e)
+    self._original_signal_handlers.clear()
 
   def __repr__(self):
     return (
@@ -325,6 +372,9 @@ def _get_total_chips(self) -> int:
 
   def __enter__(self):
     """Enters the context manager, ensuring cluster exists."""
+    atexit.register(self._cleanup)
+    self._register_signal_handlers()
+
     self.metrics_collector.record_requested_capacity(self.total_chips)
 
     self._old_jax_platforms = os.environ.get(_JAX_PLATFORMS_KEY.upper())
@@ -357,9 +407,15 @@ def __enter__(self):
       _logger.info("View proxy logs in Cloud Logging: %s", cloud_logging_link)
 
       self.proxy_pod_name = gke_utils.wait_for_pod(self._proxy_job_name)
+
+      def _set_pf_process(proc: subprocess.Popen[str]) -> None:
+        self._port_forward_process = proc
+
       self._proxy_port, self._port_forward_process = (
           gke_utils.enable_port_forwarding(
-              f"pod/{self.proxy_pod_name}", PROXY_SERVER_PORT
+              f"pod/{self.proxy_pod_name}",
+              PROXY_SERVER_PORT,
+              process_callback=_set_pf_process,
           )
       )
 
@@ -379,7 +435,7 @@ def __enter__(self):
           self.cluster,
       )
       return self
-    except Exception as e:
+    except BaseException as e:
       _logger.exception("Error setting up Pathways proxy: %r", e)
       # If any part of setup fails after deployment, cleanup.
       self._cleanup()
@@ -392,42 +448,86 @@ def __exit__(self, exc_type, exc_value, traceback):
 
   def _cleanup(self) -> None:
     """Cleans up resources created by the ISCPathways context."""
-    # Clear JAX caches and run garbage collection.
-    _logger.info("Starting Pathways proxy cleanup.")
-    jax_backend.clear_backends()
-    jax.clear_caches()
-    gc.collect()
-    _logger.info("Cleared JAX caches and ran garbage collection.")
-
-    # Terminate the port forwarding process.
-    if self._port_forward_process:
-      _logger.info("Terminating port forwarding process...")
-      self._port_forward_process.terminate()
-      try:
-        self._port_forward_process.wait(timeout=10)
-      except subprocess.TimeoutExpired as e:
-        _logger.exception(
-            "Failed to terminate port forwarding process. Not treating as an "
-            "error: %r",
-            e,
-        )
-
-    # Delete the proxy GKE job.
-    _logger.info("Deleting Pathways proxy...")
-    gke_utils.delete_gke_resource("job", self._proxy_job_name)
-    _logger.info("Pathways proxy GKE job deletion complete.")
-
-    # Restore JAX variables.
-    _logger.info("Restoring JAX env and config variables...")
-    _restore_env_var(_JAX_PLATFORMS_KEY.upper(), self._old_jax_platforms)
-    _restore_env_var(
-        _JAX_BACKEND_TARGET_KEY.upper(), self._old_jax_backend_target
-    )
-    jax.config.update(_JAX_PLATFORMS_KEY, self._old_jax_platforms_config)
-    jax.config.update(
-        _JAX_BACKEND_TARGET_KEY, self._old_jax_backend_target_config
-    )
-    _logger.info("JAX variables restored.")
+    with self._cleanup_lock:
+      if self._cleaned_up:
+        return
+      self._cleaned_up = True
+
+      atexit.unregister(self._cleanup)
+      self._restore_signal_handlers()
+
+      # Clear JAX caches and run garbage collection.
+      _logger.info("Starting Pathways proxy cleanup.")
+      jax_backend.clear_backends()
+      jax.clear_caches()
+      gc.collect()
+      _logger.info("Cleared JAX caches and ran garbage collection.")
+
+      # Terminate the port forwarding process.
+      if self._port_forward_process:
+        _logger.info("Terminating port forwarding process...")
+        try:
+          self._port_forward_process.terminate()
+          self._port_forward_process.wait(timeout=10)
+        except subprocess.TimeoutExpired as e:
+          _logger.exception(
+              "Failed to terminate port forwarding process. Killing: %r",
+              e,
+          )
+          try:
+            self._port_forward_process.kill()
+            self._port_forward_process.wait(timeout=5)
+          except Exception as kill_err:  # pylint: disable=broad-exception-caught
+            _logger.exception(
+                "Failed to kill port forwarding process: %r", kill_err
+            )
+        except Exception as e:  # pylint: disable=broad-exception-caught
+          _logger.exception(
+              "Failed to terminate port forwarding process. Not treating as an "
+              "error: %r",
+              e,
+          )
+        finally:
+          self._port_forward_process = None
+
+      # Terminate the log streaming process.
+      if self._log_process:
+        _logger.info("Terminating log streaming process...")
+        try:
+          self._log_process.terminate()
+          self._log_process.wait(timeout=5)
+        except subprocess.TimeoutExpired:
+          self._log_process.kill()
+          self._log_process.wait(timeout=5)
+        except Exception as e:  # pylint: disable=broad-exception-caught
+          _logger.exception(
+              "Failed to terminate log streaming process: %r", e
+          )
+        finally:
+          self._log_process = None
+
+      # Delete the proxy GKE job.
+      if self._proxy_job_name:
+        _logger.info("Deleting Pathways proxy...")
+        try:
+          gke_utils.delete_gke_resource("job", self._proxy_job_name)
+          _logger.info("Pathways proxy GKE job deletion complete.")
+        except Exception as e:  # pylint: disable=broad-exception-caught
+          _logger.exception(
+              "Failed to delete Pathways proxy GKE job: %r", e
+          )
+
+      # Restore JAX variables.
+      _logger.info("Restoring JAX env and config variables...")
+      _restore_env_var(_JAX_PLATFORMS_KEY.upper(), self._old_jax_platforms)
+      _restore_env_var(
+          _JAX_BACKEND_TARGET_KEY.upper(), self._old_jax_backend_target
+      )
+      jax.config.update(_JAX_PLATFORMS_KEY, self._old_jax_platforms_config)
+      jax.config.update(
+          _JAX_BACKEND_TARGET_KEY, self._old_jax_backend_target_config
+      )
+      _logger.info("JAX variables restored.")
 
 
 @contextlib.contextmanager
@@ -505,6 +605,10 @@ def connect(
   ) as t:
     if t.proxy_pod_name:
       num_slices = sum(t.expected_tpu_instances.values())
+
+      def _set_log_process(proc: subprocess.Popen[str]) -> None:
+        t._log_process = proc
+
       placement_thread = threading.Thread(
           target=_wait_for_placement,
           args=(
@@ -514,6 +618,7 @@ def connect(
               t.metrics_collector,
               t.start_time,
               t.total_chips,
+              _set_log_process,
           ),
           daemon=True,
       )

pathwaysutils/experimental/shared_pathways_service/run_workload.py

@@ -130,15 +130,29 @@ def run_command(
       collect_service_metrics=collect_service_metrics,
   ):
     logging.info("Connection established. Running command: %r", command)
+    command_args = shlex.split(command)
+    proc = subprocess.Popen(command_args, env=os.environ.copy())
     try:
-      command_args = shlex.split(command)
-      subprocess.run(command_args, check=True, env=os.environ.copy())
+      returncode = proc.wait()
+      if returncode != 0:
+        raise subprocess.CalledProcessError(returncode, command_args)
     except subprocess.CalledProcessError:
       logging.error(
           "Command failed! Find the underlying error in the logs above, where"
           " the command is invoked."
       )
       raise
+    except BaseException:
+      logging.warning(
+          "Command interrupted or terminated. Terminating child process..."
+      )
+      proc.terminate()
+      try:
+        proc.wait(timeout=5)
+      except subprocess.TimeoutExpired:
+        proc.kill()
+        proc.wait()
+      raise
     finally:
       logging.info("Command execution finished.")