Repository: caikit/caikit
Pull request #607 (2023-12-11 23:06:10)

:bug: ensure correct shutdown handlers called
<!--  Thanks for sending a pull request!  Here are some tips for you:
1. Make sure to read the Contributing Guide before submitting your PR: https://github.com/caikit/caikit/blob/main/CONTRIBUTING.md
2. If this PR closes another issue, add 'closes #<issue number>' somewhere in the PR summary. GitHub will automatically close that issue when this PR gets merged. Alternatively, adding 'refs #<issue number>' will not close the issue, but help provide the reviewer more context.-->

**What this PR does / why we need it**:

This was a tangent I went on today. This PR both:
- allows multiple servers to cleanly register shutdown handlers for interruption
- Updates the http server hackery to ensure that previously registered signal handlers are actually invoked when uvicorn overrides all the signal handlers

**Special notes for your reviewer**:

This is a bit of a departure from how we were explicitly setting signal handlers in `main` methods before, but this allows us to easily handle uvicorn's insistence on overriding all signal handlers.

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/runtime/__main__.py b/caikit/runtime/__main__.py
index 20b28853..2483e374 100644
--- a/caikit/runtime/__main__.py
+++ b/caikit/runtime/__main__.py
@@ -1,6 +1,16 @@
-# Standard
-import signal
-
+# Copyright The Caikit Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 # First Party
 import alog
 
@@ -14,25 +24,6 @@ def main():
     _grpc_server = None
     _http_server = None
 
-    def interrupt(signal_, _stack_frame):
-        log.info(
-            "<RUN87630120I>",
-            "Caikit Runtime received interrupt signal %s, shutting down",
-            signal_,
-        )
-        if _grpc_server:
-            _grpc_server.stop()
-        if _http_server:
-            _http_server.stop()
-
-    # NOTE: signal function can only be called from main thread of the main
-    # interpreter. If this function is called from a thread (like in tests)
-    # then signal handler cannot be used. Thus, we will only have real
-    # termination_handler when this is called from the __main__.
-
-    signal.signal(signal.SIGINT, interrupt)
-    signal.signal(signal.SIGTERM, interrupt)
-
     #####################
     # Start the servers
     #####################
diff --git a/caikit/runtime/grpc_server.py b/caikit/runtime/grpc_server.py
index f98f2503..8ed170b3 100644
--- a/caikit/runtime/grpc_server.py
+++ b/caikit/runtime/grpc_server.py
@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 # Standard
 from concurrent import futures
 from typing import Optional, Union
@@ -244,6 +243,7 @@ class RuntimeGRPCServer(RuntimeServerBase):
             grace_period_seconds (Union[float, int]): Grace period for service shutdown.
                 Defaults to application config
         """
+        log.info("Shutting down gRPC server")
         if grace_period_seconds is None:
             grace_period_seconds = (
                 self.config.runtime.grpc.server_shutdown_grace_period_seconds
@@ -296,7 +296,6 @@ class RuntimeGRPCServer(RuntimeServerBase):
 
 def main(blocking: bool = True):
     server = RuntimeGRPCServer()
-    server._intercept_interrupt_signal()
     server.start(blocking)
 
 
diff --git a/caikit/runtime/http_server/http_server.py b/caikit/runtime/http_server/http_server.py
index 359510d2..44e22f2d 100644
--- a/caikit/runtime/http_server/http_server.py
+++ b/caikit/runtime/http_server/http_server.py
@@ -28,6 +28,7 @@ import io
 import json
 import os
 import re
+import signal
 import ssl
 import tempfile
 import threading
@@ -237,18 +238,9 @@ class RuntimeHTTPServer(RuntimeServerBase):
             # deleted if they're ephemeral
             config.load()
 
-        # Start the server with the loaded config
+        # Build the server with the loaded config
         self.server = uvicorn.Server(config=config)
 
-        # Patch the exit handler to call this server's stop
-        original_handler = self.server.handle_exit
-
-        def shutdown_wrapper(*args, **kwargs):
-            original_handler(*args, **kwargs)
-            self.stop()
-
-        self.server.handle_exit = shutdown_wrapper
-
         # Placeholder for thread when running without blocking
         self._uvicorn_server_thread = None
 
@@ -270,10 +262,13 @@ class RuntimeHTTPServer(RuntimeServerBase):
             self.thread_pool._max_workers,
         )
 
+        # Patch the exit handler to retain correct signal handling behavior
+        self._patch_exit_handler()
+
         if blocking:
             self.server.run()
         else:
-            self.run_in_thread()
+            self._run_in_thread()
 
     def stop(self):
         """Stop the server, with an optional grace period.
@@ -282,11 +277,14 @@ class RuntimeHTTPServer(RuntimeServerBase):
             grace_period_seconds (Union[float, int]): Grace period for service shutdown.
                 Defaults to application config
         """
-        self.server.should_exit = True
+        log.info("Shutting down http server")
+
         if (
             self._uvicorn_server_thread is not None
             and self._uvicorn_server_thread.is_alive()
         ):
+            # This is required to notify the server in the thread to exit
+            self.server.should_exit = True
             self._uvicorn_server_thread.join()
 
         # Ensure we flush out any remaining billing metrics and stop metering
@@ -296,16 +294,17 @@ class RuntimeHTTPServer(RuntimeServerBase):
         # Shut down the model manager's model polling if enabled
         self._shut_down_model_manager()
 
-    def run_in_thread(self):
+    ##########
+    ## Impl ##
+    ##########
+
+    def _run_in_thread(self):
         self._uvicorn_server_thread = threading.Thread(target=self.server.run)
         self._uvicorn_server_thread.start()
         while not self.server.started:
             time.sleep(1e-3)
         log.info("HTTP Server is running in thread")
 
-    ##########
-    ## Impl ##
-    ##########
     def _bind_routes(self, service: ServicePackage):
         """Bind all rpcs as routes to the given app"""
         for rpc in service.caikit_rpcs.values():
@@ -823,10 +822,31 @@ class RuntimeHTTPServer(RuntimeServerBase):
             )
             raise ValueError() from err
 
+    def _patch_exit_handler(self):
+        """
+        🌶️🌶️🌶️ Here there are dragons! 🌶️🌶️🌶️
+        uvicorn will explicitly set the interrupt handler to `server.handle_exit` when
+        `server.run()` is called. That will override any other signal handlers that we
+        may have tried to set.
+
+        To work around this, we:
+        1. Register `server.handle_exit` as a SIGINT/SIGTERM signal handler ourselves, so that it
+          is invoked on interrupt and terminate
+        2. Set `server.handle_exit` to the existing SIGINT signal handler, so that when the uvicorn
+          server explicitly overrides the signal handler for SIGINT and SIGTERM to this, it has no
+          effect.
+
+        Since uvicorn overrides SIGINT and SIGTERM with a single common handler, any special
+            handlers added for SIGTERM but not SIGINT will not be invoked.
+        """
+        original_exit_handler = self.server.handle_exit
+        self._add_signal_handler(signal.SIGINT, original_exit_handler)
+        self._add_signal_handler(signal.SIGTERM, original_exit_handler)
+        self.server.handle_exit = signal.getsignal(signal.SIGINT)
+
 
 def main(blocking: bool = True):
     server = RuntimeHTTPServer()
-    server._intercept_interrupt_signal()
     server.start(blocking)
 
 
diff --git a/caikit/runtime/server_base.py b/caikit/runtime/server_base.py
index 5898d329..0c00a01c 100644
--- a/caikit/runtime/server_base.py
+++ b/caikit/runtime/server_base.py
@@ -112,6 +112,11 @@ class RuntimeServerBase(abc.ABC):  # pylint: disable=too-many-instance-attribute
 
         self.thread_pool: ThreadPoolExecutor = ServerThreadPool.pool
 
+        # Handle interrupts
+        # NB: This means that stop() methods will be called even if the process is interrupted
+        # before the start() method is called
+        self._intercept_interrupt_signal()
+
     @classmethod
     def _start_metrics_server(cls) -> None:
         """Start a single instance of the metrics server based on configuration"""
@@ -124,11 +129,6 @@ class RuntimeServerBase(abc.ABC):  # pylint: disable=too-many-instance-attribute
                 start_http_server(get_config().runtime.metrics.port)
             cls._metrics_server_started = True
 
-    def _intercept_interrupt_signal(self) -> None:
-        """intercept signal handler"""
-        signal.signal(signal.SIGINT, self.interrupt)
-        signal.signal(signal.SIGTERM, self.interrupt)
-
     def interrupt(self, signal_, _stack_frame):
         log.info(
             "<RUN87630120I>",
@@ -137,6 +137,43 @@ class RuntimeServerBase(abc.ABC):  # pylint: disable=too-many-instance-attribute
         )
         self.stop()
 
+    def _intercept_interrupt_signal(self) -> None:
+        """Intercept signal handlers to allow the server to stop on interrupt.
+        Calling this on a non-main thread has no effect.
+        This does not override any existing non-default signal handlers,
+        it will call them all in the reverse order they are registered.
+        """
+        self._add_signal_handler(signal.SIGINT, self.interrupt)
+        self._add_signal_handler(signal.SIGTERM, self.interrupt)
+
+    @staticmethod
+    def _add_signal_handler(sig, handler):
+        def nested_interrupt_builder(*handlers):
+            """Build and return an interrupt handler that calls all of *handlers"""
+
+            log.debug("Building interrupt handler: %s", handlers)
+
+            def interrupt(signal_, _stack_frame):
+                for handler in handlers:
+                    # Only call the handler if it is a callable fn that is _not_ a default handler
+                    log.debug("Running interrupt handler: %s", handler)
+                    if (
+                        handler
+                        and callable(handler)
+                        and handler != signal.SIG_DFL
+                        and handler is not signal.default_int_handler
+                    ):
+                        handler(signal_, _stack_frame)
+
+            return interrupt
+
+        try:
+            signal.signal(sig, nested_interrupt_builder(handler, signal.getsignal(sig)))
+        except ValueError:
+            log.info(
+                "Unable to register signal handler. Server was started from a non-main thread."
+            )
+
     def _shut_down_model_manager(self):
         """Shared utility for shutting down the model manager"""
         ModelManager.get_instance().shut_down()


Test patch:
diff --git a/tests/runtime/conftest.py b/tests/runtime/conftest.py
index 1390ab6a..e8387868 100644
--- a/tests/runtime/conftest.py
+++ b/tests/runtime/conftest.py
@@ -69,10 +69,9 @@ def http_session_scoped_open_port():
     return _open_port()
 
 
-def _open_port():
+def _open_port(start=8888):
     # TODO: This has obvious problems where the port returned for use by a test is not immediately
     # put into use, so parallel tests could attempt to use the same port.
-    start = 8888
     end = start + 1000
     host = "localhost"
     for port in range(start, end):
@@ -405,7 +404,9 @@ class ModuleSubproc:
             self.proc.kill()
 
     def __enter__(self):
-        self.proc = subprocess.Popen(self._cmd, env=self._env)
+        self.proc = subprocess.Popen(
+            self._cmd, env=self._env, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+        )
         self._kill_timer.start()
         return self.proc
 
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 91a1825e..b0809775 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -17,14 +17,11 @@
 # Standard
 from dataclasses import dataclass
 from unittest import mock
-from unittest.mock import patch
 import json
-import os
 import signal
 import tempfile
 import threading
 import time
-import uuid
 
 # Third Party
 from google.protobuf.descriptor_pool import DescriptorPool
@@ -35,6 +32,7 @@ from grpc_reflection.v1alpha.proto_reflection_descriptor_database import (
 )
 import grpc
 import pytest
+import requests
 import tls_test_tools
 
 # First Party
@@ -42,17 +40,20 @@ import alog
 
 # Local
 from caikit import get_config
-from caikit.core import MODEL_MANAGER
 from caikit.core.data_model.producer import ProducerId
 from caikit.interfaces.runtime.data_model import (
     RuntimeInfoRequest,
     RuntimeInfoResponse,
     TrainingInfoRequest,
     TrainingJob,
-    TrainingStatus,
     TrainingStatusResponse,
 )
-from caikit.runtime import get_inference_request, get_train_params, get_train_request
+from caikit.runtime import (
+    get_inference_request,
+    get_train_params,
+    get_train_request,
+    http_server,
+)
 from caikit.runtime.grpc_server import RuntimeGRPCServer
 from caikit.runtime.model_management.model_manager import ModelManager
 from caikit.runtime.protobufs import (
@@ -77,6 +78,7 @@ from tests.core.helpers import *
 from tests.fixtures import Fixtures
 from tests.runtime.conftest import (
     ModuleSubproc,
+    _open_port,
     register_trained_model,
     runtime_grpc_test_server,
 )
@@ -1350,6 +1352,63 @@ def test_grpc_sever_shutdown_with_model_poll(open_port):
         assert not server_proc.killed
 
 
+def test_all_signal_handlers_invoked(open_port):
+    """Test that a SIGINT successfully shuts down all running servers"""
+
+    # whoops, need 2 ports. Try to find another open one that isn't the one we already have
+    other_open_port = _open_port(start=open_port + 1)
+
+    with tempfile.TemporaryDirectory() as workdir:
+        server_proc = ModuleSubproc(
+            "caikit.runtime",
+            kill_timeout=30.0,
+            RUNTIME_GRPC_PORT=str(open_port),
+            RUNTIME_HTTP_PORT=str(other_open_port),
+            RUNTIME_LOCAL_MODELS_DIR=workdir,
+            RUNTIME_LAZY_LOAD_LOCAL_MODELS="true",
+            RUNTIME_LAZY_LOAD_POLL_PERIOD_SECONDS="0.1",
+            RUNTIME_METRICS_ENABLED="false",
+            RUNTIME_GRPC_ENABLED="true",
+            RUNTIME_HTTP_ENABLED="true",
+            LOG_LEVEL="info",
+        )
+        with server_proc as proc:
+            # Wait for the grpc server to be up:
+            _assert_connection(
+                grpc.insecure_channel(f"localhost:{open_port}"), max_failures=500
+            )
+
+            # Then wait for the http server as well:
+            http_failures = 0
+            while http_failures < 500:
+                try:
+                    resp = requests.get(
+                        f"http://localhost:{other_open_port}{http_server.HEALTH_ENDPOINT}",
+                        timeout=0.1,
+                    )
+                    resp.raise_for_status()
+                    break
+                except (
+                    requests.HTTPError,
+                    requests.ConnectionError,
+                    requests.ConnectTimeout,
+                ):
+                    http_failures += 1
+                    # tiny sleep because a connection refused won't hit the full `0.1`s timeout
+                    time.sleep(0.001)
+
+            # Signal the server to shut down
+            proc.send_signal(signal.SIGINT)
+
+        # Make sure the process was not killed
+        assert not server_proc.killed
+        # Check the logs (barf) to see if both grpc and http signal handlers called
+        # communicate returns (stdout, stderr) in bytes
+        logs = server_proc.proc.communicate()[1].decode("utf-8")
+        assert "Shutting down gRPC server" in logs
+        assert "Shutting down http server" in logs
+
+
 def test_construct_with_options(open_port, sample_train_service, sample_int_file):
     """Make sure that the server can be booted with config options"""
     with temp_config(