Repository: caikit/caikit
Pull request #568 (2023-11-06 21:51:49)

added new dataobject for user warnings
<!--  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 PR adds a warnings field and associated reason dataobject in order to have the ability to issue a warning when returning classification results.






Patch:
diff --git a/caikit/interfaces/nlp/data_model/classification.py b/caikit/interfaces/nlp/data_model/classification.py
index 4687c412..b1e49d44 100644
--- a/caikit/interfaces/nlp/data_model/classification.py
+++ b/caikit/interfaces/nlp/data_model/classification.py
@@ -14,6 +14,7 @@
 """Data structures for classification representations"""
 
 # Standard
+from enum import Enum
 from typing import List, Optional
 
 # Third Party
@@ -31,6 +32,21 @@ from .text_generation import FinishReason
 log = alog.use_channel("DATAM")
 
 
+@dataobject(package=NLP_PACKAGE)
+class InputWarningReason(Enum):
+    UNSUITABLE_INPUT = 0
+
+
+@dataobject(package=NLP_PACKAGE)
+class InputWarning(DataObjectBase):
+    """Input Warning data object, which returns a reason and message associated with warnings
+    to issue to a user that causes errors (such as failed text generation)
+    """
+
+    id: Annotated[InputWarningReason, FieldNumber(1)]  # id of input error
+    message: Annotated[str, FieldNumber(2)]  # Error message detailing Warning
+
+
 @dataobject(package=NLP_PACKAGE)
 class ClassificationTrainRecord(DataObjectBase):
     """A classification training record consisting of a single train instance."""
@@ -129,6 +145,9 @@ class ClassifiedGeneratedTextResult(DataObjectBase):
         Optional[np.uint64], FieldNumber(5)
     ]  # The random seed used for text generation
     input_token_count: Annotated[Optional[int], FieldNumber(6)]
+    warnings: Annotated[
+        Optional[List[InputWarning]], FieldNumber(9)
+    ]  # Warning to user in the event of input errors
 
 
 @dataobject(package=NLP_PACKAGE)
diff --git a/caikit_health_probe/__main__.py b/caikit_health_probe/__main__.py
index 86c37a5f..b73cbc42 100644
--- a/caikit_health_probe/__main__.py
+++ b/caikit_health_probe/__main__.py
@@ -23,6 +23,10 @@ import sys
 import tempfile
 import warnings
 
+# Third Party
+from cryptography import x509
+from cryptography.hazmat.backends import default_backend
+
 # First Party
 import alog
 
@@ -201,17 +205,8 @@ def _grpc_health_probe(
         )
         import grpc  # pylint: disable=import-outside-toplevel
 
-    # Server hostname to use unless using socket mode
     hostname = f"localhost:{port}"
-    socket_file = get_config().runtime.grpc.unix_socket_path
-
-    # If available, use a unix socket
-    if socket_file and os.path.exists(os.path.dirname(socket_file)):
-        socket_address = f"unix://{socket_file}"
-        log.debug("Probing gRPC server over unix socket: %s", socket_file)
-        channel = grpc.insecure_channel(socket_address)
-
-    elif tls_key and tls_cert:
+    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:
@@ -227,14 +222,18 @@ def _grpc_health_probe(
                 root_certificates=tls_server_cert,
             )
 
-        # NOTE: If the server's certificate does not have 'localhost' in it,
-        #   this will cause certificate validation errors and fail. The original
-        #   workaround for this was to parse the cert's SANs and use hostname
-        #   overrides, but that requires a full cryptographic PEM parser which
-        #   is a security-sensitive dependency to pull that we want to avoid.
-        #   Instead, the workaround is to use the unix socket server option
-        #   above.
-        channel = grpc.secure_channel(hostname, credentials=credentials)
+        # Since we are not guaranteed to know the hostname that the server is
+        # using and the server may not have 'localhost' in the SANs for the cert
+        # it is using, we need to parse out the correct hostname here so that it
+        # can be passed to grpc hostname verification. There is no option to
+        # fully disable grpc hostname verification as of the writing of this
+        # code.
+        host_name = _get_host_name(tls_server_cert)
+        channel = grpc.secure_channel(
+            hostname,
+            credentials=credentials,
+            options=(("grpc.ssl_target_name_override", host_name),),
+        )
     else:
         log.debug("Probing INSECURE gRPC server")
         channel = grpc.insecure_channel(hostname)
@@ -285,6 +284,23 @@ def _load_secret(secret: str) -> str:
     return secret
 
 
+def _get_host_name(cert_bytes: bytes) -> str:
+    """Get the Common Name from the given certificate in order to provide
+    hostname verification for the gRPC health check
+    """
+    cert = x509.load_pem_x509_certificate(cert_bytes, default_backend())
+    if sans := cert.subject.get_attributes_for_oid(x509.OID_SUBJECT_ALTERNATIVE_NAME):
+        return sans[0].value
+    if ext_san := cert.extensions.get_extension_for_oid(
+        x509.OID_SUBJECT_ALTERNATIVE_NAME
+    ):
+        if hn_sans := ext_san.value.get_values_for_type(x509.DNSName):
+            return hn_sans[0]
+        if ip_sans := ext_san.value.get_values_for_type(x509.IPAddress):
+            return ip_sans[0]
+    return cert.subject.get_attributes_for_oid(x509.OID_COMMON_NAME)[0].value
+
+
 ## Main ########################################################################
 def main():
     caikit_config = get_config()
diff --git a/pyproject.toml b/pyproject.toml
index 266dbf10..99bf3ca2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -45,6 +45,9 @@ runtime-grpc = [
     "grpcio-reflection>=1.35.0,<2.0",
     "prometheus_client>=0.12.0,<1.0",
     "py-grpc-prometheus>=0.7.0,<0.8",
+    # This is needed to support health probes for TLS gRPC servers since we
+    # cannot fully disable hostname verification.
+    "cryptography>=40.0.1,<41",
 ]
 
 runtime-http = [


Test patch:
diff --git a/tests/interfaces/nlp/test_classification.py b/tests/interfaces/nlp/test_classification.py
index e5949eec..8015107f 100644
--- a/tests/interfaces/nlp/test_classification.py
+++ b/tests/interfaces/nlp/test_classification.py
@@ -22,6 +22,10 @@ from caikit.interfaces.nlp.data_model import (
     TokenClassificationResult,
     TokenClassificationResults,
 )
+from caikit.interfaces.nlp.data_model.classification import (
+    InputWarning,
+    InputWarningReason,
+)
 
 ## Setup #########################################################################
 
@@ -60,8 +64,14 @@ classification_generated_text_result = ClassifiedGeneratedTextResult(
     generated_token_count=7,
     finish_reason=FinishReason.STOP_SEQUENCE,
     seed=42,
+    warnings=[
+        InputWarning(
+            id=InputWarningReason.UNSUITABLE_INPUT, message="unsuitable input detected"
+        )
+    ],
 )
 
+
 ## Tests ########################################################################
 
 ### ClassificationResult
@@ -260,3 +270,5 @@ def _validate_classification_generated_text_result(obj):
     assert obj.generated_token_count == 7
     assert obj.finish_reason == 5
     assert obj.seed == 42
+    assert obj.warnings[0].id == InputWarningReason.UNSUITABLE_INPUT.value
+    assert obj.warnings[0].message == "unsuitable input detected"
diff --git a/tests/runtime/conftest.py b/tests/runtime/conftest.py
index 1390ab6a..bab997ce 100644
--- a/tests/runtime/conftest.py
+++ b/tests/runtime/conftest.py
@@ -14,7 +14,6 @@ import sys
 import tempfile
 import threading
 import time
-import warnings
 
 # Third Party
 from grpc_health.v1 import health_pb2, health_pb2_grpc
@@ -152,7 +151,9 @@ def runtime_grpc_test_server(open_port, *args, **kwargs):
 
 
 @pytest.fixture(scope="session")
-def runtime_grpc_server(session_scoped_open_port) -> RuntimeGRPCServer:
+def runtime_grpc_server(
+    session_scoped_open_port,
+) -> RuntimeGRPCServer:
     with runtime_grpc_test_server(
         session_scoped_open_port,
     ) as server:
@@ -182,14 +183,12 @@ def runtime_http_test_server(open_port, *args, **kwargs):
             # 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 tls_config_override := config_overrides.get("runtime", {}).get("tls"):
-                kwargs["tls_config_override"] = tls_config_override
-            else:
+            if "use_in_test" not in config_overrides:
                 config_overrides = {}
-            check_readiness = kwargs.pop("check_readiness", True)
+            else:
+                kwargs["tls_config_override"] = config_overrides["runtime"]["tls"]
             with http_server.RuntimeHTTPServer(*args, **kwargs) as server:
-                if check_readiness:
-                    _check_http_server_readiness(server, config_overrides)
+                _check_http_server_readiness(server, config_overrides)
                 # Give tests access to the workdir
                 server.workdir = workdir
                 yield server
@@ -452,13 +451,11 @@ def _check_http_server_readiness(server, config_overrides: Dict[str, Dict]):
     done = False
     while not done:
         try:
-            with warnings.catch_warnings():
-                warnings.filterwarnings("ignore", module="urllib3")
-                response = requests.get(
-                    f"{mode}://localhost:{server.port}{http_server.HEALTH_ENDPOINT}",
-                    verify=False,
-                    cert=cert,
-                )
+            response = requests.get(
+                f"{mode}://localhost:{server.port}{http_server.HEALTH_ENDPOINT}",
+                verify=False,
+                cert=cert,
+            )
             assert response.status_code == 200
             assert response.text == "OK"
             done = True
diff --git a/tests/runtime/http_server/test_http_server.py b/tests/runtime/http_server/test_http_server.py
index 0d4ab553..d46e830f 100644
--- a/tests/runtime/http_server/test_http_server.py
+++ b/tests/runtime/http_server/test_http_server.py
@@ -99,8 +99,6 @@ def generate_tls_configs(
         config_overrides = {}
         client_keyfile, client_certfile = None, None
         ca_cert, server_cert, server_key = None, None, None
-        use_in_test = config_overrides.setdefault("use_in_test", {})
-        use_in_test["workdir"] = workdir
         if mtls or tls:
             ca_key = tls_test_tools.generate_key()[0]
             ca_cert = tls_test_tools.generate_ca_cert(ca_key)
@@ -126,9 +124,11 @@ def generate_tls_configs(
             # 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)
-            use_in_test["ca_cert"] = ca_certfile
-            use_in_test["server_key"] = server_keyfile
-            use_in_test["server_cert"] = server_certfile
+            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")
@@ -137,7 +137,7 @@ def generate_tls_configs(
                     "-----BEGIN CERTIFICATE-----\nfoobar\n-----END CERTIFICATE-----"
                 )
                 handle.write(bad_cert)
-            use_in_test["bad_ca_cert"] = bad_ca_file
+            config_overrides["use_in_test"]["bad_ca_cert"] = bad_ca_file
 
             if mtls:
                 if separate_client_ca:
@@ -171,8 +171,8 @@ def generate_tls_configs(
                     ),
                 )
                 # need to save the client cert and key in config_overrides so the mtls test below can access it
-                use_in_test["client_cert"] = client_certfile
-                use_in_test["client_key"] = client_keyfile
+                config_overrides["use_in_test"]["client_cert"] = client_certfile
+                config_overrides["use_in_test"]["client_key"] = client_keyfile
 
             config_overrides["runtime"] = {"tls": tls_config.to_dict()}
         config_overrides.setdefault("runtime", {})["http"] = {
diff --git a/tests/runtime/test_caikit_health_probe.py b/tests/runtime/test_caikit_health_probe.py
index a36769e3..cb6e1dba 100644
--- a/tests/runtime/test_caikit_health_probe.py
+++ b/tests/runtime/test_caikit_health_probe.py
@@ -21,10 +21,9 @@ early causing some of the core tests to fail!
 """
 # Standard
 from contextlib import contextmanager
-from dataclasses import dataclass
 from enum import Enum
+from typing import Tuple
 from unittest import mock
-import os
 
 # Third Party
 import pytest
@@ -45,6 +44,21 @@ from tests.runtime.http_server.test_http_server import generate_tls_configs
 log = alog.use_channel("TEST")
 
 
+class TlsMode(Enum):
+    INSECURE = 0
+    TLS = 1
+    TLS_NO_LOCALHOST = 2
+    MTLS_COMMON_CLIENT_CA = 3
+    MTLS_SEPARATE_CLIENT_CA = 4
+    MTLS_NO_LOCALHOST = 5
+
+
+class ServerMode(Enum):
+    HTTP = 0
+    GRPC = 1
+    BOTH = 2
+
+
 @contextmanager
 def maybe_runtime_grpc_test_server(*args, **kwargs):
     if get_config().runtime.grpc.enabled:
@@ -71,101 +85,54 @@ def temp_probe_config(*args, **kwargs):
             yield
 
 
-class TlsMode(Enum):
-    INSECURE = 0
-    TLS = 1
-    MTLS = 3
-
-
-class ServerMode(Enum):
-    HTTP = 0
-    GRPC = 1
-    BOTH = 2
-
-
-@dataclass
-class ProbeTestConfig:
-    tls_mode: TlsMode
-    server_mode: ServerMode
-    # TLS blobs passed as inline strings instead of files
-    inline: bool = False
-    # Run the unix socket grpc server
-    unix_socket: bool = True
-    # Put "localhost" in the SAN list for the server's cert
-    localhost_in_cert: bool = True
-    # Use a common CA for client and server certs (mTLS only)
-    common_client_ca: bool = True
-    # Whether the test should eventually become healthy
-    should_become_healthy: bool = True
-
-
 ## Tests #######################################################################
 
 
 @pytest.mark.parametrize(
     "test_config",
     [
-        # Insecure
-        ProbeTestConfig(TlsMode.INSECURE, ServerMode.HTTP),
-        ProbeTestConfig(TlsMode.INSECURE, ServerMode.GRPC),
-        ProbeTestConfig(TlsMode.INSECURE, ServerMode.BOTH),
-        # TLS
-        ProbeTestConfig(TlsMode.TLS, ServerMode.HTTP),
-        ProbeTestConfig(TlsMode.TLS, ServerMode.GRPC),
-        ProbeTestConfig(TlsMode.TLS, ServerMode.BOTH),
-        ProbeTestConfig(TlsMode.TLS, ServerMode.BOTH, inline=True),
-        ProbeTestConfig(TlsMode.TLS, ServerMode.BOTH, localhost_in_cert=False),
-        # mTLS
-        ProbeTestConfig(TlsMode.MTLS, ServerMode.HTTP),
-        ProbeTestConfig(TlsMode.MTLS, ServerMode.GRPC),
-        ProbeTestConfig(TlsMode.MTLS, ServerMode.BOTH),
-        ProbeTestConfig(TlsMode.MTLS, ServerMode.BOTH, inline=True),
-        ProbeTestConfig(TlsMode.MTLS, ServerMode.BOTH, localhost_in_cert=False),
-        ProbeTestConfig(TlsMode.MTLS, ServerMode.BOTH, common_client_ca=False),
-        # Invalid configs that never pass
-        ProbeTestConfig(
-            TlsMode.TLS,
-            ServerMode.GRPC,
-            localhost_in_cert=False,
-            unix_socket=False,
-            should_become_healthy=False,
-        ),
-        ProbeTestConfig(
-            TlsMode.TLS,
-            ServerMode.BOTH,
-            localhost_in_cert=False,
-            unix_socket=False,
-            should_become_healthy=False,
-        ),
-        ProbeTestConfig(
-            TlsMode.MTLS,
-            ServerMode.GRPC,
-            localhost_in_cert=False,
-            unix_socket=False,
-            should_become_healthy=False,
-        ),
+        (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),
+        # Tests for no "localhost" in the SAN list
+        (TlsMode.TLS_NO_LOCALHOST, False),
+        (TlsMode.MTLS_NO_LOCALHOST, False),
     ],
 )
-def test_health_probe(test_config: ProbeTestConfig):
+@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"""
-    with alog.ContextLog(log.info, "---LOG CONFIG: %s---", test_config):
+    tls_mode, inline = test_config
+    with alog.ContextLog(
+        log.info,
+        "---LOG CONFIG: tls: %s, inline: %s, server: %s---",
+        tls_mode.name,
+        inline,
+        server_mode.name,
+    ):
         # Get ports for both servers
         http_port = tls_test_tools.open_port()
         grpc_port = tls_test_tools.open_port()
 
         # Set up SAN lists if not putting "localhost" in
         server_sans, client_sans = None, None
-        if not test_config.localhost_in_cert:
+        if "NO_LOCALHOST" in tls_mode.name:
             server_sans = ["foo.bar"]
             client_sans = ["baz.bat"]
 
         # Set up tls values if needed
+        tls = tls_mode.name.startswith("TLS")
+        mtls = tls_mode.name.startswith("MTLS")
         with generate_tls_configs(
             port=http_port,
-            tls=test_config.tls_mode == TlsMode.TLS,
-            mtls=test_config.tls_mode == TlsMode.MTLS,
-            inline=test_config.inline,
-            separate_client_ca=not test_config.common_client_ca,
+            tls=tls,
+            mtls=mtls,
+            inline=inline,
+            separate_client_ca=tls_mode == TlsMode.MTLS_SEPARATE_CLIENT_CA,
             server_sans=server_sans,
             client_sans=client_sans,
         ) as config_overrides:
@@ -174,17 +141,11 @@ def test_health_probe(test_config: ProbeTestConfig):
                     "runtime": {
                         "grpc": {
                             "port": grpc_port,
-                            "enabled": test_config.server_mode
+                            "enabled": server_mode
                             in [ServerMode.GRPC, ServerMode.BOTH],
-                            "unix_socket_path": os.path.join(
-                                config_overrides["use_in_test"]["workdir"],
-                                "grpc.sock",
-                            )
-                            if test_config.unix_socket
-                            else None,
                         },
                         "http": {
-                            "enabled": test_config.server_mode
+                            "enabled": server_mode
                             in [ServerMode.HTTP, ServerMode.BOTH],
                         },
                     }
@@ -197,17 +158,11 @@ def test_health_probe(test_config: ProbeTestConfig):
                 with maybe_runtime_grpc_test_server(grpc_port):
                     # If only running gRPC, health probe should pass
                     assert caikit_health_probe.health_probe() == (
-                        test_config.should_become_healthy
-                        and test_config.server_mode == ServerMode.GRPC
+                        server_mode == ServerMode.GRPC
                     )
                     # If booting the HTTP server, do so
                     with maybe_runtime_http_test_server(
-                        http_port,
-                        tls_config_override=config_overrides,
-                        check_readiness=test_config.should_become_healthy,
+                        http_port, tls_config_override=config_overrides
                     ):
                         # Probe should always pass with both possible servers up
-                        assert (
-                            caikit_health_probe.health_probe()
-                            == test_config.should_become_healthy
-                        )
+                        assert caikit_health_probe.health_probe()