diff --git a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py b/pathwaysutils/experimental/shared_pathways_service/gke_utils.py index ac17acf..b51c914 100644 --- a/pathwaysutils/experimental/shared_pathways_service/gke_utils.py +++ b/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,6 +293,7 @@ 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. @@ -299,6 +301,9 @@ def enable_port_forwarding( 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) diff --git a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py b/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py index bc988a4..35aa8c8 100644 --- a/pathwaysutils/experimental/shared_pathways_service/isc_pathways.py +++ b/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, ) diff --git a/pathwaysutils/experimental/shared_pathways_service/run_workload.py b/pathwaysutils/experimental/shared_pathways_service/run_workload.py index 59c0d97..2ed7bb7 100644 --- a/pathwaysutils/experimental/shared_pathways_service/run_workload.py +++ b/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.") diff --git a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py index 1b69dcb..84ff763 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/gke_utils_test.py @@ -957,6 +957,53 @@ def test_get_worker_sidecar_image_custom_namespace(self): text=True, ) + def test_enable_port_forwarding_with_process_callback(self): + self.enter_context( + mock.patch.object(socket, "create_connection", autospec=True) + ) + mock_popen = self.enter_context( + mock.patch.object(subprocess, "Popen", autospec=True) + ) + mock_pick_port = self.enter_context( + mock.patch.object(portpicker, "pick_unused_port", autospec=True) + ) + mock_pick_port.return_value = 29007 + mock_process = mock_popen.return_value + mock_process.stdout = io.StringIO( + "Forwarding from 127.0.0.1:29007 -> 8080\n" + ) + mock_callback = mock.Mock() + + port, process = gke_utils.enable_port_forwarding( + "test-pod-123", 8080, process_callback=mock_callback + ) + + mock_callback.assert_called_once_with(mock_process) + self.assertEqual(port, 29007) + self.assertEqual(process, mock_process) + + def test_enable_port_forwarding_interrupted_terminates_process(self): + mock_create_connection = self.enter_context( + mock.patch.object(socket, "create_connection", autospec=True) + ) + mock_popen = self.enter_context( + mock.patch.object(subprocess, "Popen", autospec=True) + ) + mock_pick_port = self.enter_context( + mock.patch.object(portpicker, "pick_unused_port", autospec=True) + ) + mock_pick_port.return_value = 29007 + mock_process = mock_popen.return_value + mock_process.stdout = io.StringIO( + "Forwarding from 127.0.0.1:29007 -> 8080\n" + ) + mock_create_connection.side_effect = KeyboardInterrupt() + + with self.assertRaises(KeyboardInterrupt): + gke_utils.enable_port_forwarding("test-pod-123", 8080) + + mock_process.terminate.assert_called_once() + if __name__ == "__main__": absltest.main() diff --git a/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py b/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py index a7adccd..07d9042 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/isc_pathways_test.py @@ -2,7 +2,9 @@ """ import io import os +import signal import subprocess +import sys from unittest import mock from absl import flags @@ -675,6 +677,7 @@ def test_connect_success(self): mock_manager_instance.metrics_collector, mock_manager_instance.start_time, mock_manager_instance.total_chips, + mock.ANY, ), daemon=True, ) @@ -876,6 +879,151 @@ def test_connect_with_sidecar_validation_mismatch_raises_error(self): "us-docker.pkg.dev/.../sidecar:20260423-python_3.12-jax_0.10.0" ) + def test_cleanup_idempotent(self): + manager = isc_pathways._ISCPathways( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service:1234", + expected_tpu_instances={"tpuv6e:2x2": 1}, + proxy_job_name="test-proxy", + proxy_server_image="test-image", + ) + mock_delete = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "delete_gke_resource", autospec=True + ) + ) + manager._cleanup() + mock_delete.assert_called_once_with("job", "test-proxy") + # Second cleanup should be a no-op + manager._cleanup() + mock_delete.assert_called_once_with("job", "test-proxy") + + def test_enter_base_exception_triggers_cleanup(self): + manager = isc_pathways._ISCPathways( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service:1234", + expected_tpu_instances={"tpuv6e:2x2": 1}, + proxy_job_name="test-proxy", + proxy_server_image="test-image", + ) + self.enter_context( + mock.patch.object( + isc_pathways, + "_deploy_pathways_proxy_server", + side_effect=KeyboardInterrupt(), + ) + ) + mock_cleanup = self.enter_context( + mock.patch.object(manager, "_cleanup", wraps=manager._cleanup) + ) + + with self.assertRaises(KeyboardInterrupt): + manager.__enter__() + + mock_cleanup.assert_called_once() + + def test_signal_handler_triggers_cleanup_and_exits(self): + manager = isc_pathways._ISCPathways( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service:1234", + expected_tpu_instances={"tpuv6e:2x2": 1}, + proxy_job_name="test-proxy", + proxy_server_image="test-image", + ) + mock_cleanup = self.enter_context( + mock.patch.object(manager, "_cleanup", autospec=True) + ) + + manager._register_signal_handlers() + try: + self.assertIn(signal.SIGTERM, manager._original_signal_handlers) + sigterm_handler = signal.getsignal(signal.SIGTERM) + self.assertTrue(callable(sigterm_handler)) + assert callable(sigterm_handler) + with mock.patch("sys.exit") as mock_exit: + sigterm_handler(signal.SIGTERM, None) + mock_cleanup.assert_called_once() + mock_exit.assert_called_once_with(128 + signal.SIGTERM) + finally: + manager._restore_signal_handlers() + + def test_signal_handler_sigint_raises_keyboard_interrupt(self): + manager = isc_pathways._ISCPathways( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service:1234", + expected_tpu_instances={"tpuv6e:2x2": 1}, + proxy_job_name="test-proxy", + proxy_server_image="test-image", + ) + mock_cleanup = self.enter_context( + mock.patch.object(manager, "_cleanup", autospec=True) + ) + + manager._register_signal_handlers() + try: + self.assertIn(signal.SIGINT, manager._original_signal_handlers) + sigint_handler = signal.getsignal(signal.SIGINT) + self.assertTrue(callable(sigint_handler)) + assert callable(sigint_handler) + with self.assertRaises(KeyboardInterrupt): + sigint_handler(signal.SIGINT, None) + mock_cleanup.assert_called_once() + finally: + manager._restore_signal_handlers() + + def test_enter_port_forward_interrupted_triggers_cleanup(self): + manager = isc_pathways._ISCPathways( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service:1234", + expected_tpu_instances={"tpuv6e:2x2": 1}, + proxy_job_name="test-proxy", + proxy_server_image="test-image", + ) + self.enter_context( + mock.patch.object( + isc_pathways, + "_deploy_pathways_proxy_server", + autospec=True, + ) + ) + self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, + "wait_for_pod", + return_value="test-pod-123", + autospec=True, + ) + ) + self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, + "enable_port_forwarding", + side_effect=KeyboardInterrupt(), + ) + ) + mock_delete = self.enter_context( + mock.patch.object( + isc_pathways.gke_utils, "delete_gke_resource", autospec=True + ) + ) + with self.assertRaises(KeyboardInterrupt): + manager.__enter__() + mock_delete.assert_called_once_with("job", "test-proxy") if __name__ == "__main__": absltest.main() diff --git a/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py b/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py index a7ddffb..7323000 100644 --- a/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py +++ b/pathwaysutils/test/experimental/shared_pathways_service/run_workload_test.py @@ -43,9 +43,11 @@ def fake_connect_fn(**kwargs: Any) -> Generator[FakeConnect, None, None]: with fake as f: yield f - mock_run = self.enter_context( - mock.patch.object(subprocess, "run", autospec=True) + mock_popen = self.enter_context( + mock.patch.object(subprocess, "Popen", autospec=True) ) + mock_proc = mock_popen.return_value + mock_proc.wait.return_value = 0 run_workload.run_command( cluster="test-cluster", @@ -71,9 +73,10 @@ def fake_connect_fn(**kwargs: Any) -> Generator[FakeConnect, None, None]: self.assertTrue(fake.exited) with self.subTest("Command executed"): - mock_run.assert_called_once_with( - ["echo", "hello"], check=True, env=mock.ANY + mock_popen.assert_called_once_with( + ["echo", "hello"], env=mock.ANY ) + mock_proc.wait.assert_called_once() def test_run_command_runs_command_inside_context(self): """Verifies that the command is executed while the connection is active.""" @@ -84,14 +87,16 @@ def fake_connect_fn(**kwargs: Any) -> Generator[None, None, None]: del kwargs yield None - def mock_run_side_effect(*args: Any, **kwargs: Any) -> None: + def mock_wait_side_effect(*args: Any, **kwargs: Any) -> int: nonlocal connection_active_during_run connection_active_during_run = True + return 0 - mock_run = self.enter_context( - mock.patch.object(subprocess, "run", autospec=True) + mock_popen = self.enter_context( + mock.patch.object(subprocess, "Popen", autospec=True) ) - mock_run.side_effect = mock_run_side_effect + mock_proc = mock_popen.return_value + mock_proc.wait.side_effect = mock_wait_side_effect run_workload.run_command( cluster="test-cluster", @@ -114,10 +119,11 @@ def fake_connect_fn(**kwargs: Any) -> Generator[None, None, None]: del kwargs yield None - mock_run = self.enter_context( - mock.patch.object(subprocess, "run", autospec=True) + mock_popen = self.enter_context( + mock.patch.object(subprocess, "Popen", autospec=True) ) - mock_run.side_effect = subprocess.CalledProcessError(1, "false") + mock_proc = mock_popen.return_value + mock_proc.wait.return_value = 1 with self.assertRaises(subprocess.CalledProcessError): run_workload.run_command( @@ -132,6 +138,34 @@ def fake_connect_fn(**kwargs: Any) -> Generator[None, None, None]: connect_fn=fake_connect_fn, ) + def test_run_command_interrupted_terminates_subprocess(self): + + @contextlib.contextmanager + def fake_connect_fn(**kwargs: Any) -> Generator[None, None, None]: + del kwargs + yield None + + mock_popen = self.enter_context( + mock.patch.object(subprocess, "Popen", autospec=True) + ) + mock_proc = mock_popen.return_value + mock_proc.wait.side_effect = KeyboardInterrupt() + + with self.assertRaises(KeyboardInterrupt): + run_workload.run_command( + cluster="test-cluster", + project="test-project", + region="test-region", + gcs_bucket="test-bucket", + pathways_service="test-service:1234", + tpu_type="tpuv6e:4x8", + tpu_count=1, + command="sleep 100", + connect_fn=fake_connect_fn, + ) + + mock_proc.terminate.assert_called_once() + @flagsaver.flagsaver( cluster="test-cluster", project="test-project",