Repository: caikit/caikit
Pull request #552 (2023-10-30 19:47:41)

Health probe command 548
**What this PR does / why we need it**:

Addresses #548 

This PR adds `caikit.runtime.health_probe` which runs a health probe against one or both runtime servers. It uses the existing environment config to match the running server.

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [X] this PR contains unit tests
- [X] this PR has been tested for backwards compatibility


joerunde: I'd ask if we could vendor in the stock grpc health probe binary for non-mtls cases but that would of course lower our ossf scorecard

Patch:
diff --git a/caikit/runtime/dump_services.py b/caikit/runtime/dump_services.py
index 244a2afc..0bf37632 100644
--- a/caikit/runtime/dump_services.py
+++ b/caikit/runtime/dump_services.py
@@ -100,7 +100,7 @@ def dump_http_services(output_dir: str):
             handle.write(json.dumps(response.json(), indent=2))
 
 
-if __name__ == "__main__":
+def main():
     parser = argparse.ArgumentParser(
         description="Dump grpc and http services for inference and train"
     )
@@ -133,3 +133,7 @@ if __name__ == "__main__":
         dump_grpc_services(out_dir, write_modules_json)
     if get_config().runtime.http.enabled:
         dump_http_services(out_dir)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/caikit/runtime/grpc_server.py b/caikit/runtime/grpc_server.py
index bbd55686..0ba42c8e 100644
--- a/caikit/runtime/grpc_server.py
+++ b/caikit/runtime/grpc_server.py
@@ -196,11 +196,20 @@ class RuntimeGRPCServer(RuntimeServerBase):
             )
             if self.tls_config.client.cert:
                 log.info("<RUN10001806I>", "Running with mutual TLS")
+                # Combine the client cert with the server's own cert so that
+                # health probes can use the server's key/cert instead of needing
+                # one signed by a potentially-external CA.
+                root_certificates = b"\n".join(
+                    [
+                        bytes(self._load_secret(self.tls_config.client.cert), "utf-8"),
+                        tls_server_pair[1],
+                    ]
+                )
                 # Client will verify the server using server cert and the server
                 # will verify the client using client cert.
                 server_credentials = grpc.ssl_server_credentials(
                     [tls_server_pair],
-                    root_certificates=self._load_secret(self.tls_config.client.cert),
+                    root_certificates=root_certificates,
                     require_client_auth=True,
                 )
             else:
diff --git a/caikit_health_probe/__init__.py b/caikit_health_probe/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/caikit_health_probe/__main__.py b/caikit_health_probe/__main__.py
new file mode 100644
index 00000000..53aa33c9
--- /dev/null
+++ b/caikit_health_probe/__main__.py
@@ -0,0 +1,285 @@
+# 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.
+"""
+This module implements a common health probe for all running runtime servers.
+"""
+# Standard
+from contextlib import contextmanager
+from typing import Optional, Tuple
+import importlib.util
+import os
+import sys
+import tempfile
+import warnings
+
+# First Party
+import alog
+
+# We play some tricks to import caikit's config here and avoid importing all of
+# caikit itself. The reason for this is that importing caikit can be very costly
+# relative to the actual cost of probing the servers and since this is intended
+# to stand alone as an executable, that import time cost is added to every call.
+caikit_spec = importlib.util.find_spec("caikit")
+sys.path = [os.path.dirname(caikit_spec.origin)] + sys.path
+# Third Party
+from config import get_config
+
+sys.path = sys.path[1:]
+
+
+log = alog.use_channel("PROBE")
+
+
+@alog.timed_function(log.debug)
+def health_probe() -> bool:
+    """Run a health probe against all running runtime servers.
+
+    This function is intended to be run from an environment where the config is
+    identical to the config that the server is running such as from inside a
+    kubernetes pod where the server is also running.
+
+    Returns:
+        healthy (bool): True if all servers are healthy, False otherwise
+    """
+
+    # Get TLS key/cert files if possible
+    config = get_config()
+    tls_key = config.runtime.tls.server.key
+    tls_cert = config.runtime.tls.server.cert
+    client_ca = config.runtime.tls.client.cert
+    http_healthy, grpc_healthy = None, None
+
+    if config.runtime.http.enabled:
+        log.debug("Checking HTTP server health")
+        http_healthy = _http_health_probe(
+            config.runtime.http.port, tls_key, tls_cert, client_ca
+        )
+
+    if config.runtime.grpc.enabled:
+        log.debug("Checking gRPC server health")
+        grpc_healthy = _grpc_health_probe(
+            config.runtime.grpc.port, tls_key, tls_cert, client_ca
+        )
+
+    if False in [http_healthy, grpc_healthy]:
+        log.info(
+            "<RUN64273066I>",
+            "Server not healthy. HTTP: %s, gRPC: %s",
+            http_healthy,
+            grpc_healthy,
+        )
+        return False
+    return True
+
+
+## Implementation ##############################################################
+
+
+def _http_health_probe(
+    port: int,
+    tls_key: Optional[str],
+    tls_cert: Optional[str],
+    client_ca: Optional[str],
+) -> bool:
+    """Probe the http server
+
+    The implementation of this utility is a bit tricky because mTLS makes this
+    quite challenging. For insecure or TLS servers, we expect a valid healthy
+    response, but for mTLS servers, we may not have a valid key/cert pair that
+    the client can present to the server that is signed by the expected CA if
+    the trusted client CA does not match the one that signed the server's
+    key/cert pair.
+
+    The workaround for this is to detect SSLError and consider that to be a
+    passing health check. If the server is healthy enough to _reject_ bad SSL
+    requests, it's healthy enough to server good ones!
+
+    Args:
+        port (int): The port that the HTTP server is serving on
+        tls_key (Optional[str]): Body or path to the TLS key file if TLS/mTLS
+            enabled
+        tls_cert (Optional[str]): Body or path to the TLS cert file if TLS/mTLS
+            enabled
+        client_ca (Optional[str]): The client ca cert that the server is using
+            for mutual client auth
+
+    Returns:
+        healthy (bool): True if all servers are healthy, False otherwise
+    """
+    # NOTE: Local imports for optional dependency
+    with alog.ContextTimer(log.debug2, "Done with local grpc imports: "):
+
+        # Third Party
+        import requests  # pylint: disable=import-outside-toplevel
+
+    # Requests requires that the TLS information be in files
+    with _tls_files(tls_key, tls_cert) as tls_files:
+        key_file, cert_file = tls_files
+        if key_file and cert_file:
+            protocol = "https"
+            kwargs = {"verify": False}
+            if client_ca:
+                log.debug("Probing mTLS HTTP Server")
+                kwargs["cert"] = (key_file, cert_file)
+            else:
+                log.debug("Probing TLS HTTP Server")
+        else:
+            log.debug("Probing INSECURE HTTP Server")
+            protocol = "http"
+            kwargs = {}
+
+        try:
+            # Suppress insecure connection warnings since we disable server
+            # verification. This is ok since this probe will be run against
+            # localhost in a pod where the server is _known_ to be authentic.
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", module="urllib3")
+                session = requests.Session()
+                retries = requests.adapters.Retry(total=0)
+                session.mount(
+                    f"{protocol}://", requests.adapters.HTTPAdapter(max_retries=retries)
+                )
+                # NOTE: Not using the constant to avoid big imports
+                resp = session.get(
+                    f"{protocol}://localhost:{port}/health",
+                    timeout=0.01,
+                    **kwargs,
+                )
+            resp.raise_for_status()
+            return True
+        except requests.exceptions.SSLError as err:
+            log.debug("Got SSLError indicating a healthy SSL server! %s", err)
+            log.debug2(err, exc_info=True)
+            return True
+        except Exception as err:  # pylint: disable=broad-exception-caught
+            log.debug2("Caught unexpected error: %s", err, exc_info=True)
+        return False
+
+
+def _grpc_health_probe(
+    port: int,
+    tls_key: Optional[str],
+    tls_cert: Optional[str],
+    client_ca: Optional[str],
+) -> bool:
+    """Probe the grpc server
+
+    Since the gRPC server trusts its own cert for client verification, we can
+    make a valid health probe against the running server regardless of (m)TLS
+    config.
+
+    Args:
+        port (int): The port that the gRPC server is serving on
+        tls_key (Optional[str]): Body or path to the TLS key file if TLS/mTLS
+            enabled
+        tls_cert (Optional[str]): Body or path to the TLS cert file if TLS/mTLS
+            enabled
+        client_ca (Optional[str]): The client ca cert that the server is using
+            for mutual client auth
+
+    Returns:
+        healthy (bool): True if all servers are healthy, False otherwise
+    """
+    # NOTE: Local imports for optional dependency
+    with alog.ContextTimer(log.debug2, "Done with local grpc imports: "):
+
+        # Third Party
+        from grpc_health.v1 import (  # pylint: disable=import-outside-toplevel
+            health_pb2,
+            health_pb2_grpc,
+        )
+        import grpc  # pylint: disable=import-outside-toplevel
+
+    hostname = f"localhost:{port}"
+    if tls_key and tls_cert:
+        tls_server_key = bytes(_load_secret(tls_key), "utf-8")
+        tls_server_cert = bytes(_load_secret(tls_cert), "utf-8")
+        if client_ca:
+            log.debug("Probing mTLS gRPC server")
+            credentials = grpc.ssl_channel_credentials(
+                root_certificates=tls_server_cert,
+                private_key=tls_server_key,
+                certificate_chain=tls_server_cert,
+            )
+        else:
+            log.debug("Probing TLS gRPC server")
+            credentials = grpc.ssl_channel_credentials(
+                root_certificates=tls_server_cert,
+            )
+        channel = grpc.secure_channel(hostname, credentials=credentials)
+    else:
+        log.debug("Probing INSECURE gRPC server")
+        channel = grpc.insecure_channel(hostname)
+
+    client = health_pb2_grpc.HealthStub(channel)
+    try:
+        client.Check(health_pb2.HealthCheckRequest())
+        return True
+    except Exception as err:  # pylint: disable=broad-exception-caught
+        log.debug2("Caught unexpected error: %s", err, exc_info=True)
+        return False
+
+
+@contextmanager
+def _tls_files(
+    tls_key: Optional[str],
+    tls_cert: Optional[str],
+) -> Tuple[Optional[str], Optional[str]]:
+    """Get files for the TLS key/cert if given"""
+    if not tls_key or not tls_cert:
+        yield None, None
+        return
+    valid_file_vals = [
+        ((os.path.exists(fname) and fname) or None) for fname in [tls_key, tls_cert]
+    ]
+    if all(valid_file_vals):
+        yield tls_key, tls_cert
+        return
+    with tempfile.TemporaryDirectory() as workdir:
+        key_file, cert_file = valid_file_vals
+        if not key_file:
+            key_file = os.path.join(workdir, "tls.key")
+            with open(key_file, "w", encoding="utf-8") as handle:
+                handle.write(tls_key)
+        if not cert_file:
+            cert_file = os.path.join(workdir, "tls.cert")
+            with open(cert_file, "w", encoding="utf-8") as handle:
+                handle.write(tls_cert)
+        yield key_file, cert_file
+        return
+
+
+def _load_secret(secret: str) -> str:
+    """NOTE: Copied from grpc_server to avoid costly imports"""
+    if os.path.exists(secret):
+        with open(secret, "r", encoding="utf-8") as secret_file:
+            return secret_file.read()
+    return secret
+
+
+## Main ########################################################################
+def main():
+    caikit_config = get_config()
+    alog.configure(
+        default_level=caikit_config.log.level,
+        filters=caikit_config.log.filters,
+        thread_id=caikit_config.log.thread_id,
+        formatter=caikit_config.log.formatter,
+    )
+    if not health_probe():
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/pyproject.toml b/pyproject.toml
index ff2ddf52..9f570b48 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -30,6 +30,12 @@ dependencies = [
     "werkzeug>=2.3.7,<4.0.0"
 ]
 
+[project.scripts]
+
+caikit-runtime = "caikit.runtime.__main__:main"
+caikit-health-probe = "caikit_health_probe.__main__:main"
+caikit-render-interfaces = "caikit.runtime.dump_services:main"
+
 [project.optional-dependencies]
 
 ## Runtime Extra Sets ##
@@ -43,6 +49,7 @@ runtime-grpc = [
 
 runtime-http = [
     "fastapi[all]>=0.100,<1",
+    "requests>=2.28.2,<3",
     "sse-starlette>=1.6.1,<2",
 ]
 
@@ -62,7 +69,6 @@ dev-test = [
     "pytest-cov>=2.10.1,<5.0",
     "pytest-html>=3.1.1,<5.0",
     "pytest>=6.2.5,<8.0",
-    "requests>=2.28.2,<3",
     "tls_test_tools>=0.1.1",
     "wheel>=0.38.4",
     "caikit[interfaces-vision]",


Test patch:
diff --git a/tests/runtime/conftest.py b/tests/runtime/conftest.py
index 045da428..04ce61e2 100644
--- a/tests/runtime/conftest.py
+++ b/tests/runtime/conftest.py
@@ -179,9 +179,13 @@ def runtime_http_test_server(open_port, *args, **kwargs):
             },
             "merge",
         ):
-            config_overrides = {}
-            if "tls_config_override" in kwargs:
-                config_overrides = kwargs["tls_config_override"]
+            # Forward the special "tls_config_override" to "tls_config_override"
+            # IFF the configs contain actual TLS (indicated by the presence of
+            # the special "use_in_test" element).
+            config_overrides = kwargs.pop("tls_config_override", {})
+            if "use_in_test" not in config_overrides:
+                config_overrides = {}
+            else:
                 kwargs["tls_config_override"] = config_overrides["runtime"]["tls"]
             with http_server.RuntimeHTTPServer(*args, **kwargs) as server:
                 _check_http_server_readiness(server, config_overrides)
diff --git a/tests/runtime/http_server/test_http_server.py b/tests/runtime/http_server/test_http_server.py
index f6fe625c..13287c8d 100644
--- a/tests/runtime/http_server/test_http_server.py
+++ b/tests/runtime/http_server/test_http_server.py
@@ -89,17 +89,17 @@ def generate_tls_configs(
     tls: bool = False,
     mtls: bool = False,
     inline: bool = False,
+    separate_client_ca: bool = False,
     **http_config_overrides,
 ) -> Dict[str, Dict]:
     """Helper to generate tls configs"""
     with tempfile.TemporaryDirectory() as workdir:
         config_overrides = {}
-        client_keyfile, client_certfile, ca_certfile = None, None, None
+        client_keyfile, client_certfile = None, None
         ca_cert, server_cert, server_key = None, None, None
         if mtls or tls:
             ca_key = tls_test_tools.generate_key()[0]
             ca_cert = tls_test_tools.generate_ca_cert(ca_key)
-            ca_certfile, _ = save_key_cert_pair("ca", workdir, cert=ca_cert)
             server_key, server_cert = tls_test_tools.generate_derived_key_cert_pair(
                 ca_key=ca_key
             )
@@ -110,15 +110,22 @@ def generate_tls_configs(
             if inline:
                 tls_config = TLSConfig(
                     server=KeyPair(cert=server_cert, key=server_key),
-                    client=KeyPair(cert=ca_cert if mtls else "", key=""),
+                    client=KeyPair(cert="", key=""),
                 )
             else:
                 tls_config = TLSConfig(
                     server=KeyPair(cert=server_certfile, key=server_keyfile),
-                    client=KeyPair(cert=ca_certfile if mtls else "", key=""),
+                    client=KeyPair(cert="", key=""),
                 )
-            # need to save this ca_certfile in config_overrides so the tls tests below can access it from client side
-            config_overrides["use_in_test"] = {"ca_cert": ca_certfile}
+
+            # need to save this ca_certfile in config_overrides so the tls
+            # tests below can access it from client side
+            ca_certfile, _ = save_key_cert_pair("ca", workdir, cert=ca_cert)
+            config_overrides["use_in_test"] = {
+                "ca_cert": ca_certfile,
+                "server_key": server_keyfile,
+                "server_cert": server_certfile,
+            }
 
             # also saving a bad ca_certfile for a failure test case
             bad_ca_file = os.path.join(workdir, "bad_ca_cert.crt")
@@ -130,10 +137,34 @@ def generate_tls_configs(
             config_overrides["use_in_test"]["bad_ca_cert"] = bad_ca_file
 
             if mtls:
+                if separate_client_ca:
+                    subject_kwargs = {"common_name": "my.client"}
+                    client_ca_key = tls_test_tools.generate_key()[0]
+                    client_ca_cert = tls_test_tools.generate_ca_cert(
+                        client_ca_key, **subject_kwargs
+                    )
+                else:
+                    subject_kwargs = {}
+                    client_ca_key = ca_key
+                    client_ca_cert = ca_cert
+
+                # If inlining the client CA
+                if inline:
+                    tls_config.client.cert = client_ca_cert
+                else:
+                    client_ca_certfile, _ = save_key_cert_pair(
+                        "client_ca", workdir, cert=client_ca_cert
+                    )
+                    tls_config.client.cert = client_ca_certfile
+
+                # Set up the client key/cert pair derived from the client CA
                 client_certfile, client_keyfile = save_key_cert_pair(
                     "client",
                     workdir,
-                    *tls_test_tools.generate_derived_key_cert_pair(ca_key=ca_key),
+                    *tls_test_tools.generate_derived_key_cert_pair(
+                        ca_key=client_ca_key,
+                        **subject_kwargs,
+                    ),
                 )
                 # need to save the client cert and key in config_overrides so the mtls test below can access it
                 config_overrides["use_in_test"]["client_cert"] = client_certfile
@@ -161,9 +192,7 @@ def test_insecure_server(runtime_http_server, open_port):
 
 
 def test_basic_tls_server(open_port):
-    with generate_tls_configs(
-        open_port, tls=True, mtls=False, http_config_overrides={}
-    ) as config_overrides:
+    with generate_tls_configs(open_port, tls=True, mtls=False) as config_overrides:
         with runtime_http_test_server(
             open_port,
             tls_config_override=config_overrides,
@@ -177,9 +206,7 @@ def test_basic_tls_server(open_port):
 
 
 def test_basic_tls_server_with_wrong_cert(open_port):
-    with generate_tls_configs(
-        open_port, tls=True, mtls=False, http_config_overrides={}
-    ) as config_overrides:
+    with generate_tls_configs(open_port, tls=True, mtls=False) as config_overrides:
         with runtime_http_test_server(
             open_port,
             tls_config_override=config_overrides,
@@ -193,14 +220,36 @@ def test_basic_tls_server_with_wrong_cert(open_port):
 
 
 def test_mutual_tls_server(open_port):
+    with generate_tls_configs(open_port, tls=True, mtls=True) as config_overrides:
+        with runtime_http_test_server(
+            open_port,
+            tls_config_override=config_overrides,
+        ) as http_server_with_mtls:
+            # start a non-blocking http server with mutual tls
+            resp = requests.get(
+                f"https://localhost:{http_server_with_mtls.port}/docs",
+                verify=config_overrides["use_in_test"]["ca_cert"],
+                cert=(
+                    config_overrides["use_in_test"]["client_cert"],
+                    config_overrides["use_in_test"]["client_key"],
+                ),
+            )
+            resp.raise_for_status()
+
+
+def test_mutual_tls_server_different_client_ca(open_port):
     with generate_tls_configs(
-        open_port, tls=True, mtls=True, http_config_overrides={}
+        open_port,
+        tls=True,
+        mtls=True,
+        separate_client_ca=True,
     ) as config_overrides:
+        # start a non-blocking http server with mutual tls
         with runtime_http_test_server(
             open_port,
             tls_config_override=config_overrides,
         ) as http_server_with_mtls:
-            # start a non-blocking http server with mutual tls
+            # Make a request with the client's key/cert pair
             resp = requests.get(
                 f"https://localhost:{http_server_with_mtls.port}/docs",
                 verify=config_overrides["use_in_test"]["ca_cert"],
@@ -217,7 +266,7 @@ def test_mutual_tls_server_inline(open_port):
     than with files
     """
     with generate_tls_configs(
-        open_port, tls=True, mtls=True, inline=True, http_config_overrides={}
+        open_port, tls=True, mtls=True, inline=True
     ) as config_overrides:
         with runtime_http_test_server(
             open_port,
@@ -236,9 +285,7 @@ def test_mutual_tls_server_inline(open_port):
 
 
 def test_mutual_tls_server_with_wrong_cert(open_port):
-    with generate_tls_configs(
-        open_port, tls=True, mtls=True, http_config_overrides={}
-    ) as config_overrides:
+    with generate_tls_configs(open_port, tls=True, mtls=True) as config_overrides:
         with runtime_http_test_server(
             open_port,
             tls_config_override=config_overrides,
diff --git a/tests/runtime/test_caikit_health_probe.py b/tests/runtime/test_caikit_health_probe.py
new file mode 100644
index 00000000..ece6ea02
--- /dev/null
+++ b/tests/runtime/test_caikit_health_probe.py
@@ -0,0 +1,144 @@
+# 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.
+"""
+Tests for the uniform health probe
+
+🌶️🌶️🌶️ This test relies on test infrastructure in caikit.runtime, so the test
+needs to live inside tests/runtime even though the functionality being tested is
+not. If this is moved to the top of tests, the runtime test infra boots up too
+early causing some of the core tests to fail!
+"""
+# Standard
+from contextlib import contextmanager
+from enum import Enum
+from typing import Tuple
+from unittest import mock
+
+# Third Party
+import pytest
+import tls_test_tools
+
+# First Party
+from caikit_health_probe import __main__ as caikit_health_probe
+
+# Local
+from caikit import get_config
+from tests.conftest import temp_config
+from tests.runtime.conftest import runtime_grpc_test_server, runtime_http_test_server
+from tests.runtime.http_server.test_http_server import generate_tls_configs
+
+## Helpers #####################################################################
+
+
+class TlsMode(Enum):
+    INSECURE = 0
+    TLS = 1
+    MTLS_COMMON_CLIENT_CA = 2
+    MTLS_SEPARATE_CLIENT_CA = 3
+
+
+class ServerMode(Enum):
+    HTTP = 0
+    GRPC = 1
+    BOTH = 2
+
+
+@contextmanager
+def maybe_runtime_grpc_test_server(*args, **kwargs):
+    if get_config().runtime.grpc.enabled:
+        with runtime_grpc_test_server(*args, **kwargs) as grpc_server:
+            yield grpc_server
+    else:
+        yield
+
+
+@contextmanager
+def maybe_runtime_http_test_server(*args, **kwargs):
+    if get_config().runtime.http.enabled:
+        with runtime_http_test_server(*args, **kwargs) as http_server:
+            yield http_server
+    else:
+        yield
+
+
+@contextmanager
+def temp_probe_config(*args, **kwargs):
+    with temp_config(*args, **kwargs) as the_config:
+        get_config_mock = mock.MagicMock(return_value=the_config)
+        with mock.patch.object(caikit_health_probe, "get_config", get_config_mock):
+            yield
+
+
+## Tests #######################################################################
+
+
+@pytest.mark.parametrize(
+    "test_config",
+    [
+        (TlsMode.INSECURE, True),
+        (TlsMode.TLS, True),
+        (TlsMode.TLS, False),
+        (TlsMode.MTLS_COMMON_CLIENT_CA, True),
+        (TlsMode.MTLS_COMMON_CLIENT_CA, False),
+        (TlsMode.MTLS_SEPARATE_CLIENT_CA, True),
+        (TlsMode.MTLS_SEPARATE_CLIENT_CA, False),
+    ],
+)
+@pytest.mark.parametrize("server_mode", ServerMode.__members__.values())
+def test_health_probe(test_config: Tuple[TlsMode, bool], server_mode: ServerMode):
+    """Test all of the different ways that the servers could be running"""
+    tls_mode, inline = test_config
+
+    # Get ports for both servers
+    http_port = tls_test_tools.open_port()
+    grpc_port = tls_test_tools.open_port()
+
+    # Set up tls values if needed
+    tls = tls_mode == TlsMode.TLS
+    mtls = tls_mode in [TlsMode.MTLS_COMMON_CLIENT_CA, TlsMode.MTLS_SEPARATE_CLIENT_CA]
+    with generate_tls_configs(
+        port=http_port,
+        tls=tls,
+        mtls=mtls,
+        inline=inline,
+        separate_client_ca=tls_mode == TlsMode.MTLS_SEPARATE_CLIENT_CA,
+    ) as config_overrides:
+        with temp_probe_config(
+            {
+                "runtime": {
+                    "grpc": {
+                        "port": grpc_port,
+                        "enabled": server_mode in [ServerMode.GRPC, ServerMode.BOTH],
+                    },
+                    "http": {
+                        "enabled": server_mode in [ServerMode.HTTP, ServerMode.BOTH],
+                    },
+                }
+            },
+            "merge",
+        ):
+            # Health probe fails with no servers booted
+            assert not caikit_health_probe.health_probe()
+            # If booting the gRPC server, do so
+            with maybe_runtime_grpc_test_server(grpc_port):
+                # If only running gRPC, health probe should pass
+                assert caikit_health_probe.health_probe() == (
+                    server_mode == ServerMode.GRPC
+                )
+                # If booting the HTTP server, do so
+                with maybe_runtime_http_test_server(
+                    http_port, tls_config_override=config_overrides
+                ):
+                    # Probe should always pass with both possible servers up
+                    assert caikit_health_probe.health_probe()
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 6dfead16..432ab7ae 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -1119,6 +1119,49 @@ def test_mtls(open_port):
             stub.Check(health_check_request)
 
 
+def test_mtls_different_root(open_port):
+    """Make sure mtls communication works when the CA for the client is not the
+    same as the CA for the server (including health checks using the server's
+    CA)
+    """
+    # Server TLS Infra
+    server_ca_key = tls_test_tools.generate_key()[0]
+    server_ca_cert = tls_test_tools.generate_ca_cert(server_ca_key)
+    server_tls_key, server_tls_cert = tls_test_tools.generate_derived_key_cert_pair(
+        server_ca_key
+    )
+
+    # Client TLS Infra
+    client_ca_key = tls_test_tools.generate_key()[0]
+    client_ca_cert = tls_test_tools.generate_ca_cert(
+        client_ca_key, common_name="my.client"
+    )
+    client_tls_key, client_tls_cert = tls_test_tools.generate_derived_key_cert_pair(
+        client_ca_key, common_name="my.client"
+    )
+
+    server_tls_config = TLSConfig(
+        server=KeyPair(cert=server_tls_cert, key=server_tls_key),
+        client=KeyPair(cert=client_ca_cert, key=""),
+    )
+    with runtime_grpc_test_server(
+        open_port,
+        tls_config_override=server_tls_config,
+    ) as server:
+        # Connect using the client's creds
+        _assert_connection(
+            _make_secure_channel(
+                server, server_ca_cert, client_tls_key, client_tls_cert
+            )
+        )
+        # Connect using the server's creds
+        _assert_connection(
+            _make_secure_channel(
+                server, server_ca_cert, server_tls_key, server_tls_cert
+            )
+        )
+
+
 @pytest.mark.parametrize(
     "enabled_services",
     [(True, False), (False, True), (False, False)],
diff --git a/tests/runtime/test_service_factory.py b/tests/runtime/test_service_factory.py
index d36d12af..015f5c21 100644
--- a/tests/runtime/test_service_factory.py
+++ b/tests/runtime/test_service_factory.py
@@ -55,7 +55,7 @@ def clean_data_model(sample_inference_service, sample_train_service):
     with reset_global_protobuf_registry():
         with temp_dpool(
             inherit_global=True,
-            skip_inherit=[".*sampletask.*\.proto"],
+            skip_inherit=[r".*sampletask.*\.proto"],
         ) as dpool:
             yield dpool