Repository: caikit/caikit (python)

Issue:
HTTP Server Shutdown with Local Model Dirs Sync
## Describe the bug

When running `caikit.runtime.http_server` with `runtime.lazy_load_local_models: true`, the server fails to terminate on ctrl-C due to the polling timer.


Pull request:
Server shutdown with sync
## Description

Closes #364

This PR updates the shutdown logic in `grpc_server` and `http_server` to correctly shut down the runtime Model Manager's poll timer for lazy model loading. While I was working on this, I discovered several other bugs that needed fixing:

1. Runtime was using `runtime.port` instead of `runtime.grpc.port`
2. The prometheus metrics were _always_ booting up in the `grpc_server`'s `main()`, regardless of `runtime.metrics.enabled`

Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index 979d6dc6..3fd5d499 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -104,9 +104,6 @@ runtime:
         server_shutdown_grace_period_seconds: 45
         # Listen on a unix socket (model mesh feature)
         unix_socket_path: /tmp/mmesh/grpc.sock
-    # Service exposure options
-    port: 8085
-    unix_socket_path: /tmp/mmesh/grpc.sock
 
     # Configuration for the http server
     http:
diff --git a/caikit/runtime/grpc_server.py b/caikit/runtime/grpc_server.py
index ad90e823..69b2aacc 100644
--- a/caikit/runtime/grpc_server.py
+++ b/caikit/runtime/grpc_server.py
@@ -81,7 +81,6 @@ class RuntimeGRPCServer(RuntimeServerBase):
         if handle_terminations:
             signal.signal(signal.SIGINT, self.interrupt)
             signal.signal(signal.SIGTERM, self.interrupt)
-        self.port = self.config.runtime.port
 
         # Initialize basic server
         # py_grpc_prometheus.server_metrics.
@@ -246,6 +245,7 @@ class RuntimeGRPCServer(RuntimeServerBase):
             grace_period_seconds (Union[float, int]): Grace period for service shutdown.
                 Defaults to application config
         """
+        log.debug("Shutting down grpc server")
         if grace_period_seconds is None:
             grace_period_seconds = (
                 self.config.runtime.grpc.server_shutdown_grace_period_seconds
@@ -254,6 +254,8 @@ class RuntimeGRPCServer(RuntimeServerBase):
         # Ensure we flush out any remaining billing metrics and stop metering
         if self.config.runtime.metering.enabled:
             self._global_predict_servicer.stop_metering()
+        # Shut down the model manager's model polling if enabled
+        self._shut_down_model_manager()
 
     def render_protos(self, proto_out_dir: str) -> None:
         """Renders all the necessary protos for this service into a directory
@@ -297,10 +299,16 @@ class RuntimeGRPCServer(RuntimeServerBase):
 def main(blocking: bool = True):
     # Configure using the log level and formatter type specified in config.
     caikit.core.toolkit.logging.configure()
+    log.debug("Starting up caikit.runtime.grpc_server")
 
     # Start serving Prometheus metrics
     caikit_config = get_config()
-    start_http_server(caikit_config.runtime.metrics.port)
+    if caikit_config.runtime.metrics.enabled:
+        log.info(
+            "Serving prometheus metrics on port %s", caikit_config.runtime.metrics.port
+        )
+        with alog.ContextTimer(log.info, "Booted metrics server in "):
+            start_http_server(caikit_config.runtime.metrics.port)
 
     # Enable signal handling
     handle_terminations = True
diff --git a/caikit/runtime/http_server.py b/caikit/runtime/http_server.py
index 4658360e..873c6f3c 100644
--- a/caikit/runtime/http_server.py
+++ b/caikit/runtime/http_server.py
@@ -92,6 +92,9 @@ GRPC_CODE_TO_HTTP = {
 REQUIRED_INPUTS_KEY = "inputs"
 OPTIONAL_INPUTS_KEY = "parameters"
 
+# Endpoint to use for health checks
+HEALTH_ENDPOINT = "/health"
+
 ## RuntimeHTTPServer ###########################################################
 
 
@@ -122,6 +125,9 @@ class RuntimeHTTPServer(RuntimeServerBase):
 
         self.package_name = inference_service.descriptor.full_name.rsplit(".", 1)[0]
 
+        # Add the health endpoint
+        self.app.get(HEALTH_ENDPOINT)(self._health_check)
+
         # Bind all routes to the server
         self._bind_routes(inference_service)
         # self._bind_routes(train_service)
@@ -157,6 +163,15 @@ class RuntimeHTTPServer(RuntimeServerBase):
         )
         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
 
@@ -194,6 +209,9 @@ class RuntimeHTTPServer(RuntimeServerBase):
         if self.config.runtime.metering.enabled:
             self.global_predict_servicer.stop_metering()
 
+        # Shut down the model manager's model polling if enabled
+        self._shut_down_model_manager()
+
     def run_in_thread(self):
         self._uvicorn_server_thread = threading.Thread(target=self.server.run)
         self._uvicorn_server_thread.start()
@@ -523,6 +541,10 @@ class RuntimeHTTPServer(RuntimeServerBase):
         PYDANTIC_REGISTRY[dm_class] = pydantic_model
         return pydantic_model
 
+    @staticmethod
+    def _health_check():
+        log.debug4("Server healthy")
+
 
 ## Main ########################################################################
 
diff --git a/caikit/runtime/model_management/model_manager.py b/caikit/runtime/model_management/model_manager.py
index 24f097f5..486b2708 100644
--- a/caikit/runtime/model_management/model_manager.py
+++ b/caikit/runtime/model_management/model_manager.py
@@ -141,9 +141,12 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
 
     def shut_down(self):
         """Shut down cache purging"""
+        self._enable_lazy_load_poll = False
         timer = getattr(self, "_lazy_sync_timer", None)
         if timer is not None:
             timer.cancel()
+            if timer.is_alive():
+                timer.join()
 
     ## Model Management ##
 
@@ -247,8 +250,15 @@ class ModelManager:  # pylint: disable=too-many-instance-attributes
                     self._lazy_load_poll_period_seconds,
                 )
             if self._lazy_sync_timer is not None and self._lazy_sync_timer.is_alive():
-                log.debug2("Canceling live timer")
+                log.debug3("Canceling live timer")
                 self._lazy_sync_timer.cancel()
+            log.debug3(
+                "Starting next poll timer for %ss", self._lazy_load_poll_period_seconds
+            )
+            log.debug4(
+                "All open threads: %s",
+                [thread.name for thread in threading.enumerate()],
+            )
             self._lazy_sync_timer = threading.Timer(
                 self._lazy_load_poll_period_seconds, self.sync_local_models
             )
diff --git a/caikit/runtime/server_base.py b/caikit/runtime/server_base.py
index ebe10034..3fce363f 100644
--- a/caikit/runtime/server_base.py
+++ b/caikit/runtime/server_base.py
@@ -24,6 +24,7 @@ import alog
 
 # Local
 from caikit.config import get_config
+from caikit.runtime.model_management.model_manager import ModelManager
 
 log = alog.use_channel("SERVR-BASE")
 
@@ -45,6 +46,7 @@ class RuntimeServerBase(abc.ABC):
         self.tls_config = (
             tls_config_override if tls_config_override else self.config.runtime.tls
         )
+        log.debug4("Full caikit config: %s", get_config())
 
     @classmethod
     def _find_port(cls, start=8888, end=None, host="127.0.0.1"):
@@ -75,6 +77,10 @@ class RuntimeServerBase(abc.ABC):
                 # port is open
                 return start
 
+    def _shut_down_model_manager(self):
+        """Shared utility for shutting down the model manager"""
+        ModelManager.get_instance().shut_down()
+
     @abc.abstractmethod
     def start(self, blocking: bool = True):
         pass


Test patch:
diff --git a/tests/fixtures/config/config.yml b/tests/fixtures/config/config.yml
index a46b00df..648f1fa2 100644
--- a/tests/fixtures/config/config.yml
+++ b/tests/fixtures/config/config.yml
@@ -1,5 +1,7 @@
 # Config for tests
 runtime:
+    metrics:
+        enabled: false
     metering:
         # Switch on metering by default
         enabled: true
diff --git a/tests/runtime/conftest.py b/tests/runtime/conftest.py
index f672ba5e..36491763 100644
--- a/tests/runtime/conftest.py
+++ b/tests/runtime/conftest.py
@@ -3,12 +3,15 @@ This sets up global test configs when pytest starts
 """
 
 # Standard
-from concurrent.futures import Future
 from contextlib import contextmanager
 from typing import Type, Union
 import os
+import shlex
 import socket
+import subprocess
+import sys
 import tempfile
+import threading
 import time
 
 # Third Party
@@ -122,7 +125,7 @@ def runtime_grpc_test_server(open_port, *args, **kwargs):
                 "runtime": {
                     "metering": {"log_dir": temp_log_dir},
                     "training": {"output_dir": temp_save_dir},
-                    "port": open_port,
+                    "grpc": {"port": open_port},
                 }
             },
             "merge",
@@ -260,6 +263,48 @@ def register_trained_model(
     servicer._model_manager.loaded_models[model_id] = loaded_model
 
 
+class ModuleSubproc:
+    def __init__(
+        self, module_to_run: str, *args, kill_timeout: float = 10.0, **env_vars
+    ):
+        """Run the given python as a subprocess and kill it after the given timeout"""
+        # Set up the command
+        cmd = f"{sys.executable} -m {module_to_run}"
+        for arg in args:
+            cmd += f" {arg}"
+        self._cmd = shlex.split(cmd)
+
+        # Set up the environment
+        self._env = {**os.environ, **env_vars}
+        self._env["PYTHONPATH"] = ":".join(sys.path)
+
+        # Start the process
+        self.proc = None
+        self._killed = False
+
+        # Start the timer that will kill the process
+        self._kill_timer = threading.Timer(kill_timeout, self._kill_proc)
+
+    @property
+    def killed(self):
+        return self._killed
+
+    def _kill_proc(self):
+        self._killed = True
+        if self.proc is not None:
+            self.proc.kill()
+
+    def __enter__(self):
+        self.proc = subprocess.Popen(self._cmd, env=self._env)
+        self._kill_timer.start()
+        return self.proc
+
+    def __exit__(self, *_):
+        self.proc.wait()
+        self._kill_timer.cancel()
+        self._kill_timer.join()
+
+
 # IMPLEMENTATION DETAILS ############################################################
 
 
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 0b0f4778..f732aa48 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -19,6 +19,8 @@ from dataclasses import dataclass
 from unittest import mock
 import json
 import os
+import signal
+import tempfile
 import threading
 import time
 import uuid
@@ -66,7 +68,11 @@ from sample_lib.data_model import (
 from tests.conftest import random_test_id
 from tests.core.helpers import *
 from tests.fixtures import Fixtures
-from tests.runtime.conftest import register_trained_model, runtime_grpc_test_server
+from tests.runtime.conftest import (
+    ModuleSubproc,
+    register_trained_model,
+    runtime_grpc_test_server,
+)
 import caikit.interfaces.common
 import sample_lib
 
@@ -1053,6 +1059,32 @@ def test_streaming_handlers_are_built_correctly(runtime_grpc_server):
     assert new_handler.stream_stream.__name__ == "safe_rpc_call"
 
 
+def test_grpc_sever_shutdown_with_model_poll(open_port):
+    """Test that a SIGINT successfully shuts down the running server"""
+    with tempfile.TemporaryDirectory() as workdir:
+        server_proc = ModuleSubproc(
+            "caikit.runtime.grpc_server",
+            kill_timeout=30.0,
+            RUNTIME_GRPC_PORT=str(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",
+        )
+        with server_proc as proc:
+
+            # Wait for the server to be up
+            _assert_connection(
+                grpc.insecure_channel(f"localhost:{open_port}"), max_failures=500
+            )
+
+            # Signal the server to shut down
+            proc.send_signal(signal.SIGINT)
+
+        # Make sure the process was not killed
+        assert not server_proc.killed
+
+
 # Test implementation details #########################
 @dataclass
 class KeyPair:
@@ -1087,7 +1119,7 @@ def _make_secure_channel(
     return grpc.secure_channel(f"localhost:{server.port}", credentials=credentials)
 
 
-def _assert_connection(channel):
+def _assert_connection(channel, max_failures=20):
     """Check that we can ping the server on this channel.
 
     Assumes that it will come up, but it's hard to distinguish between a failure to boot and a
@@ -1097,9 +1129,7 @@ def _assert_connection(channel):
     failures = 0
     while not done:
         try:
-            stub = health_pb2_grpc.HealthStub(channel)
-            health_check_request = health_pb2.HealthCheckRequest()
-            stub.Check(health_check_request)
+            health_pb2_grpc.HealthStub(channel).Check(health_pb2.HealthCheckRequest())
             done = True
         except grpc.RpcError as e:
             log.debug(
@@ -1108,5 +1138,5 @@ def _assert_connection(channel):
 
             time.sleep(0.1)
             failures += 1
-            if failures > 20:
+            if failures > max_failures:
                 raise e
diff --git a/tests/runtime/test_http_server.py b/tests/runtime/test_http_server.py
index 257e072a..54355dc4 100644
--- a/tests/runtime/test_http_server.py
+++ b/tests/runtime/test_http_server.py
@@ -20,6 +20,7 @@ from dataclasses import dataclass
 from typing import List
 import json
 import os
+import signal
 import tempfile
 
 # Third Party
@@ -30,10 +31,10 @@ import tls_test_tools
 
 # Local
 from caikit.core import DataObjectBase, dataobject
-from caikit.core.data_model import DataBase
 from caikit.interfaces.nlp.data_model import GeneratedTextStreamResult, GeneratedToken
 from caikit.runtime import http_server
 from tests.conftest import temp_config
+from tests.runtime.conftest import ModuleSubproc, open_port
 
 ## Helpers #####################################################################
 
@@ -335,6 +336,40 @@ def test_pydantic_wrapping_with_lists():
     assert foo.bars[0].baz == 1
 
 
+def test_http_server_shutdown_with_model_poll(open_port):
+    """Test that a SIGINT successfully shuts down the running server"""
+    with tempfile.TemporaryDirectory() as workdir:
+        server_proc = ModuleSubproc(
+            "caikit.runtime.http_server",
+            RUNTIME_HTTP_PORT=str(open_port),
+            RUNTIME_LOCAL_MODELS_DIR=workdir,
+            RUNTIME_LAZY_LOAD_LOCAL_MODELS="true",
+            RUNTIME_LAZY_LOAD_POLL_PERIOD_SECONDS="0.1",
+        )
+        with server_proc as proc:
+            # Wait for the server to be up
+            while True:
+                try:
+                    resp = requests.get(
+                        f"http://localhost:{open_port}{http_server.HEALTH_ENDPOINT}",
+                        timeout=0.1,
+                    )
+                    resp.raise_for_status()
+                    break
+                except (
+                    requests.HTTPError,
+                    requests.ConnectionError,
+                    requests.ConnectTimeout,
+                ):
+                    pass
+
+            # Signal the server to shut down
+            proc.send_signal(signal.SIGINT)
+
+        # Make sure the process was not killed
+        assert not server_proc.killed
+
+
 # TODO: uncomment later
 # def test_train():
 #     server = http_server.RuntimeHTTPServer()