Repository: caikit/caikit
Pull request #389 (2023-08-01 20:20:50)

✨ Introduce a common start script
<!--  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**:

For: https://github.com/caikit/caikit/issues/225

This PR:
- Refactors the init function of the `http_server` to align with `grpc_server`. 
- Moves a lot of common code into `server_base.py`.
- Introduces a common `caikt/runtime/__main__.py` which can start both servers based on configurations.
- Refactors the `text-sentiment` example to use this common script.
- Also fixes https://github.com/caikit/caikit/issues/380

**Special notes for your reviewer**:

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/config/config.yml b/caikit/config/config.yml
index 2196bb74..0fbb48b8 100644
--- a/caikit/config/config.yml
+++ b/caikit/config/config.yml
@@ -109,7 +109,7 @@ runtime:
 
     # Configuration for the http server
     http:
-        enabled: false
+        enabled: true
         port: 8080
         # String prefix for API routes
         route_prefix: api/v1
diff --git a/caikit/runtime/__main__.py b/caikit/runtime/__main__.py
new file mode 100644
index 00000000..12d0c06f
--- /dev/null
+++ b/caikit/runtime/__main__.py
@@ -0,0 +1,58 @@
+# Standard
+import signal
+
+# First Party
+import alog
+
+# Local
+from caikit.config.config import get_config
+from caikit.runtime.grpc_server import RuntimeGRPCServer
+from caikit.runtime.http_server import RuntimeHTTPServer
+
+log = alog.use_channel("RUNTIME-MAIN")
+
+
+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
+    #####################
+
+    # Start serving grpc server
+    if get_config().runtime.grpc.enabled:
+        log.debug("Starting up caikit.runtime.grpc_server")
+
+        _grpc_server = RuntimeGRPCServer()
+        _grpc_server.start(blocking=not get_config().runtime.http.enabled)
+
+    # Start serving http server
+    if get_config().runtime.http.enabled:
+        log.debug("Starting up caikit.runtime.http_server")
+
+        _http_server = RuntimeHTTPServer()
+        _http_server.start(blocking=True)  # make http always blocking
+
+
+if __name__ == "__main__":
+    main()
diff --git a/caikit/runtime/dump_services.py b/caikit/runtime/dump_services.py
index 0ebbd1d4..cb13bc74 100644
--- a/caikit/runtime/dump_services.py
+++ b/caikit/runtime/dump_services.py
@@ -20,6 +20,7 @@ import sys
 # Local
 from ..core.data_model import render_dataobject_protos
 from .service_factory import ServicePackageFactory
+from caikit.config.config import get_config
 import caikit
 
 
@@ -71,5 +72,7 @@ if __name__ == "__main__":
     # Set up logging so users can set LOG_LEVEL etc
     caikit.core.toolkit.logging.configure()
 
-    dump_grpc_services(out_dir)
-    dump_http_services(out_dir)
+    if get_config().runtime.grpc.enabled:
+        dump_grpc_services(out_dir)
+    if get_config().runtime.http.enabled:
+        dump_http_services(out_dir)
diff --git a/caikit/runtime/grpc_server.py b/caikit/runtime/grpc_server.py
index 69b2aacc..f0783362 100644
--- a/caikit/runtime/grpc_server.py
+++ b/caikit/runtime/grpc_server.py
@@ -16,12 +16,10 @@
 from concurrent import futures
 from typing import Optional, Union
 import os
-import signal
 
 # Third Party
 from grpc_health.v1 import health, health_pb2_grpc
 from grpc_reflection.v1alpha import reflection
-from prometheus_client import start_http_server
 from py_grpc_prometheus.prometheus_server_interceptor import PromServerInterceptor
 import grpc
 
@@ -49,7 +47,6 @@ from caikit.runtime.servicers.model_train_servicer import ModelTrainServicerImpl
 from caikit.runtime.servicers.training_management_servicer import (
     TrainingManagementServicerImpl,
 )
-from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
 
 # Have pylint ignore broad exception catching in this file so that we can log all
 # unexpected errors using alog.
@@ -65,22 +62,9 @@ class RuntimeGRPCServer(RuntimeServerBase):
 
     def __init__(
         self,
-        inference_service: ServicePackage,
-        training_service: Optional[ServicePackage],
         tls_config_override: Optional[aconfig.Config] = None,
-        handle_terminations: bool = False,
     ):
         super().__init__(get_config().runtime.grpc.port, tls_config_override)
-        self.inference_service = inference_service
-        self.training_service = training_service
-
-        # 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__.
-        if handle_terminations:
-            signal.signal(signal.SIGINT, self.interrupt)
-            signal.signal(signal.SIGTERM, self.interrupt)
 
         # Initialize basic server
         # py_grpc_prometheus.server_metrics.
@@ -90,6 +74,8 @@ class RuntimeGRPCServer(RuntimeServerBase):
             ),
             interceptors=(PROMETHEUS_METRICS_INTERCEPTOR,),
         )
+        # Start metrics server
+        RuntimeServerBase._start_metrics_server()
 
         # Start tracking service names for reflection
         service_names = [reflection.SERVICE_NAME]
@@ -226,18 +212,6 @@ class RuntimeGRPCServer(RuntimeServerBase):
         if blocking:
             self.server.wait_for_termination(None)
 
-    def interrupt(self, signal_, _stack_frame):
-        """Interrupt the server. Implements the interface for Python signal.signal
-
-        Reference: https://docs.python.org/3/library/signal.html#signal.signal
-        """
-        log.info(
-            "<RUN81000120I>",
-            "Caikit Runtime received interrupt signal %s, shutting down",
-            signal_,
-        )
-        self.stop()
-
     def stop(self, grace_period_seconds: Union[float, int] = None):
         """Stop the server, with an optional grace period.
 
@@ -245,7 +219,6 @@ 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
@@ -297,43 +270,8 @@ 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()
-    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
-
-    # We should always be able to stand up an inference service
-    inference_service: ServicePackage = ServicePackageFactory.get_service_package(
-        ServicePackageFactory.ServiceType.INFERENCE,
-    )
-
-    # But maybe not always a training service
-    try:
-        training_service: Optional[
-            ServicePackage
-        ] = ServicePackageFactory.get_service_package(
-            ServicePackageFactory.ServiceType.TRAINING,
-        )
-    except CaikitRuntimeException as e:
-        log.warning("Cannot stand up training service, disabling training: %s", e)
-        training_service = None
-
-    server = RuntimeGRPCServer(
-        inference_service=inference_service,
-        training_service=training_service,
-        handle_terminations=handle_terminations,
-    )
+    server = RuntimeGRPCServer()
+    server._intercept_interrupt_signal()
     server.start(blocking)
 
 
diff --git a/caikit/runtime/http_server.py b/caikit/runtime/http_server.py
index c84406ef..6b021fbf 100644
--- a/caikit/runtime/http_server.py
+++ b/caikit/runtime/http_server.py
@@ -50,15 +50,14 @@ from caikit.core.data_model import DataBase
 from caikit.core.data_model.dataobject import make_dataobject
 from caikit.core.toolkit.sync_to_async import async_wrap_iter
 from caikit.runtime.server_base import RuntimeServerBase
-from caikit.runtime.service_factory import ServicePackage, ServicePackageFactory
+from caikit.runtime.service_factory import ServicePackage
 from caikit.runtime.service_generation.rpcs import CaikitRPCBase
 from caikit.runtime.servicers.global_predict_servicer import GlobalPredictServicer
 from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
-import caikit.core.toolkit.logging
 
 ## Globals #####################################################################
 
-log = alog.use_channel("HTTP")
+log = alog.use_channel("SERVR-HTTP")
 
 # Registry of DM -> Pydantic model mapping to avoid errors when reusing messages
 # across endpoints
@@ -111,11 +110,11 @@ class RuntimeHTTPServer(RuntimeServerBase):
 
         self.app = FastAPI()
 
+        # Start metrics server
+        RuntimeServerBase._start_metrics_server()
+
         # Set up the central predict servicer
-        inference_service = ServicePackageFactory().get_service_package(
-            ServicePackageFactory.ServiceType.INFERENCE,
-        )
-        self.global_predict_servicer = GlobalPredictServicer(inference_service)
+        self.global_predict_servicer = GlobalPredictServicer(self.inference_service)
 
         # Set up the central train servicer
         # TODO: uncomment later on
@@ -124,7 +123,9 @@ class RuntimeHTTPServer(RuntimeServerBase):
         # )
         # self.global_train_servicer = GlobalTrainServicer(train_service)
 
-        self.package_name = inference_service.descriptor.full_name.rsplit(".", 1)[0]
+        self.package_name = self.inference_service.descriptor.full_name.rsplit(".", 1)[
+            0
+        ]
 
         # Add the health endpoint
         self.app.get(HEALTH_ENDPOINT, response_class=PlainTextResponse)(
@@ -132,7 +133,7 @@ class RuntimeHTTPServer(RuntimeServerBase):
         )
 
         # Bind all routes to the server
-        self._bind_routes(inference_service)
+        self._bind_routes(self.inference_service)
         # self._bind_routes(train_service)
 
         # Parse TLS configuration
@@ -179,11 +180,12 @@ class RuntimeHTTPServer(RuntimeServerBase):
         self._uvicorn_server_thread = None
 
     def __del__(self):
-        if get_config().runtime.metering.enabled:
+        config = get_config()
+        if config and config.runtime.metering.enabled:
             self.global_predict_servicer.stop_metering()
 
     def start(self, blocking: bool = True):
-        """Boot the gRPC server. Can be non-blocking, or block until shutdown
+        """Boot the http server. Can be non-blocking, or block until shutdown
 
         Args:
             blocking (boolean): Whether to block until shutdown
@@ -200,7 +202,6 @@ class RuntimeHTTPServer(RuntimeServerBase):
             grace_period_seconds (Union[float, int]): Grace period for service shutdown.
                 Defaults to application config
         """
-        log.info("Shutting down HTTP Server")
         self.server.should_exit = True
         if (
             self._uvicorn_server_thread is not None
@@ -328,6 +329,7 @@ class RuntimeHTTPServer(RuntimeServerBase):
         pydantic_response = self._dataobject_to_pydantic(
             self._get_response_dataobject(rpc)
         )
+
         # pylint: disable=unused-argument
         @self.app.post(self._get_route(rpc), response_model=pydantic_response)
         async def _handler(
@@ -393,7 +395,6 @@ class RuntimeHTTPServer(RuntimeServerBase):
                 route = "/" + route
             return route
         if rpc.name.endswith("Train"):
-
             route = "/".join(
                 [self.config.runtime.http.route_prefix, "{model_id}", rpc.name]
             )
@@ -550,13 +551,10 @@ class RuntimeHTTPServer(RuntimeServerBase):
         return "OK"
 
 
-## Main ########################################################################
-
-
-def main():
-    caikit.core.toolkit.logging.configure()
+def main(blocking: bool = True):
     server = RuntimeHTTPServer()
-    server.start()
+    server._intercept_interrupt_signal()
+    server.start(blocking)
 
 
 if __name__ == "__main__":
diff --git a/caikit/runtime/server_base.py b/caikit/runtime/server_base.py
index 2baf594d..34455c64 100644
--- a/caikit/runtime/server_base.py
+++ b/caikit/runtime/server_base.py
@@ -16,6 +16,10 @@
 # Standard
 from typing import Optional
 import abc
+import signal
+
+# Third Party
+from prometheus_client import start_http_server
 
 # First Party
 import aconfig
@@ -24,12 +28,16 @@ import alog
 # Local
 from caikit.config import get_config
 from caikit.runtime.model_management.model_manager import ModelManager
+from caikit.runtime.service_factory import ServicePackage, ServicePackageFactory
+from caikit.runtime.types.caikit_runtime_exception import CaikitRuntimeException
+import caikit
 
 log = alog.use_channel("SERVR-BASE")
 
 
 class RuntimeServerBase(abc.ABC):
     __doc__ = __doc__
+    _metrics_server_started = False
 
     def __init__(self, base_port: int, tls_config_override: Optional[aconfig.Config]):
         self.config = get_config()
@@ -39,6 +47,54 @@ class RuntimeServerBase(abc.ABC):
         )
         log.debug4("Full caikit config: %s", get_config())
 
+        # Configure using the log level and formatter type specified in config.
+        caikit.core.toolkit.logging.configure()
+
+        # We should always be able to stand up an inference service
+        self.inference_service: ServicePackage = (
+            ServicePackageFactory.get_service_package(
+                ServicePackageFactory.ServiceType.INFERENCE,
+            )
+        )
+
+        # But maybe not always a training service
+        try:
+            training_service: Optional[
+                ServicePackage
+            ] = ServicePackageFactory.get_service_package(
+                ServicePackageFactory.ServiceType.TRAINING,
+            )
+        except CaikitRuntimeException as e:
+            log.warning("Cannot stand up training service, disabling training: %s", e)
+            training_service = None
+
+        self.training_service = training_service
+
+    @classmethod
+    def _start_metrics_server(cls) -> None:
+        """Start a single instance of the metrics server based on configuration"""
+        if not cls._metrics_server_started and get_config().runtime.metrics.enabled:
+            log.info(
+                "Serving prometheus metrics on port %s",
+                get_config().runtime.metrics.port,
+            )
+            with alog.ContextTimer(log.info, "Booted metrics server in "):
+                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>",
+            "Caikit Runtime received interrupt signal %s, shutting down",
+            signal_,
+        )
+        self.stop()
+
     def _shut_down_model_manager(self):
         """Shared utility for shutting down the model manager"""
         ModelManager.get_instance().shut_down()
diff --git a/caikit/runtime/servicers/global_predict_servicer.py b/caikit/runtime/servicers/global_predict_servicer.py
index 13a94ce5..ee58abd8 100644
--- a/caikit/runtime/servicers/global_predict_servicer.py
+++ b/caikit/runtime/servicers/global_predict_servicer.py
@@ -95,8 +95,10 @@ class GlobalPredictServicer:
         inference_service: ServicePackage,
         use_abortable_threads: bool = get_config().runtime.use_abortable_threads,
     ):
+        self._started_metering = False
         self._model_manager = ModelManager.get_instance()
         if get_config().runtime.metering.enabled:
+            self._started_metering = True
             self.rpc_meter = RPCMeter()
             log.info(
                 "<RUN76773775I>",
@@ -267,9 +269,10 @@ class GlobalPredictServicer:
             return response
 
     def stop_metering(self):
-        if get_config().runtime.metering.enabled:
+        if get_config().runtime.metering.enabled and self._started_metering:
             self.rpc_meter.flush_metrics()
             self.rpc_meter.end_writer_thread()
+            self._started_metering = False
 
     ## Implementation Details ##################################################
 
diff --git a/caikit/runtime/servicers/training_management_servicer.py b/caikit/runtime/servicers/training_management_servicer.py
index b1941316..81dfbded 100644
--- a/caikit/runtime/servicers/training_management_servicer.py
+++ b/caikit/runtime/servicers/training_management_servicer.py
@@ -37,13 +37,14 @@ class TrainingManagementServicerImpl:
     service in Model Mesh as a Model-Runtime."""
 
     def GetTrainingStatus(self, request, context):  # pylint: disable=unused-argument
-        """Get the status of a training by ID"""
+        """Missing associated documentation comment in .proto file."""
         training_info = TrainingInfoRequest.from_proto(request)
-        model_future = self._get_model_future(
-            training_info.training_id, operation="get_status"
-        )
 
         try:
+            model_future = MODEL_MANAGER.get_model_future(
+                training_id=training_info.training_id
+            )
+
             reasons = []
             if model_future.get_info().errors:
                 reasons = [str(error) for error in model_future.get_info().errors]
@@ -53,10 +54,10 @@ class TrainingManagementServicerImpl:
                 state=model_future.get_info().status,
                 reasons=reasons,
             ).to_proto()
-        except Exception as err:
+        except ValueError as err:
             raise CaikitRuntimeException(
-                grpc.StatusCode.INTERNAL,
-                "Failed to get status for training id {}".format(
+                grpc.StatusCode.NOT_FOUND,
+                "{} not found in the list of currently running training jobs".format(
                     training_info.training_id,
                 ),
             ) from err
@@ -64,11 +65,12 @@ class TrainingManagementServicerImpl:
     def CancelTraining(self, request, context):  # pylint: disable=unused-argument
         """Cancel a training future."""
         training_info = TrainingInfoRequest.from_proto(request)
-        model_future = self._get_model_future(
-            training_info.training_id, operation="cancel"
-        )
 
         try:
+            model_future = MODEL_MANAGER.get_model_future(
+                training_id=training_info.training_id
+            )
+
             model_future.cancel()
 
             reasons = []
@@ -80,46 +82,19 @@ class TrainingManagementServicerImpl:
                 state=model_future.get_info().status,
                 reasons=reasons,
             ).to_proto()
-        except Exception as err:
-            log.debug2(
-                "Unexpected error trying to cancel training id %s: [%s]",
-                training_info.training_id,
-                err,
-            )
-            raise CaikitRuntimeException(
-                grpc.StatusCode.INTERNAL,
-                "Failed to cancel training id {}".format(
-                    training_info.training_id,
-                ),
-            ) from err
 
-    @staticmethod
-    def _get_model_future(training_id: str, operation: str):
-        """Returns a model future, or raises 404 caikit runtime exception on error.
-        Wrapped here so that we only catch errors directly in the `trainer.get_model_future` call
-        """
-        try:
-            return MODEL_MANAGER.get_model_future(training_id=training_id)
         except ValueError as err:
-            log.debug2(
-                "Caught ValueError while trying to look up model future for id %s: [%s]",
-                training_id,
-                err,
-            )
             raise CaikitRuntimeException(
                 grpc.StatusCode.NOT_FOUND,
-                "{} not found in the list of currently running training jobs. "
-                "Could not perform {}".format(training_id, operation),
+                "{} not found in the list of currently running training jobs. \
+                    Did not perform cancel.".format(
+                    training_info.training_id,
+                ),
             ) from err
         except Exception as err:
-            log.debug2(
-                "Caught unexpected exception while trying to look up model future for id %s: [%s]",
-                training_id,
-                err,
-            )
             raise CaikitRuntimeException(
                 grpc.StatusCode.INTERNAL,
-                "Unexpected error with training id {}. Could not perform {}".format(
-                    training_id, operation
+                "Failed to cancel training id {}".format(
+                    training_info.training_id,
                 ),
             ) from err
diff --git a/examples/text-sentiment/README.md b/examples/text-sentiment/README.md
index 8b128fff..c5569844 100644
--- a/examples/text-sentiment/README.md
+++ b/examples/text-sentiment/README.md
@@ -15,16 +15,13 @@ Install the dependencies: `pip install -r requirements.txt`
 
 ## Running the Caikit runtime
 
-In one terminal, start the runtime grpc server:
+In one terminal, start the runtime server:
 
 ```shell
 python3 start_runtime.py
 ```
 
-You can alternatively start the REST server with
-```shell
-python3 start_runtime.py --protocol http
-```
+This should start the runtime server with both grpc and http enabled.
 
 You should see output similar to the following:
 
@@ -63,11 +60,6 @@ In another terminal, run the client code:
 ```shell
 python3 client.py
 ```
-Or, for the REST client:
-```shell
-python3 client.py --protocol http
-```
-
 The client code calls the model and queries it for sentiment analysis on 2 different pieces of text.
 
 You should see output similar to the following:
@@ -75,16 +67,36 @@ You should see output similar to the following:
 ```command
 $ python3 client.py
 
-<function register_backend_type at 0x7fe930bdbdc0> is still in the BETA phase and subject to change!
 Text: I am not feeling well today!
-RESPONSE: classes {
+RESPONSE from gRPC: classes {
   class_name: "NEGATIVE"
-  confidence: 0.99977594614028931
+  confidence: 0.9997759461402893
 }
 
 Text: Today is a nice sunny day
-RESPONSE: classes {
+RESPONSE from gRPC: classes {
   class_name: "POSITIVE"
   confidence: 0.999869704246521
 }
+
+
+Text: I am not feeling well today!
+RESPONSE from HTTP: {
+    "classes": [
+        {
+            "class_name": "NEGATIVE",
+            "confidence": 0.9997759461402893
+        }
+    ]
+}
+
+Text: Today is a nice sunny day
+RESPONSE from HTTP: {
+    "classes": [
+        {
+            "class_name": "POSITIVE",
+            "confidence": 0.999869704246521
+        }
+    ]
+}
 ```
diff --git a/examples/text-sentiment/client.py b/examples/text-sentiment/client.py
index c3b1d186..b339d594 100644
--- a/examples/text-sentiment/client.py
+++ b/examples/text-sentiment/client.py
@@ -16,11 +16,11 @@
 import json
 
 # Third Party
-from start_runtime import protocol_arg
 import grpc
 import requests
 
 # Local
+from caikit.config.config import get_config
 from caikit.runtime.service_factory import ServicePackageFactory
 import caikit
 
@@ -30,6 +30,8 @@ if __name__ == "__main__":
             "merge_strategy": "merge",
             "runtime": {
                 "library": "text_sentiment",
+                "grpc": {"enabled": True},
+                "http": {"enabled": True},
             },
         }
     )
@@ -38,11 +40,9 @@ if __name__ == "__main__":
         ServicePackageFactory.ServiceType.INFERENCE,
     )
 
-    protocol = protocol_arg()
-
     model_id = "text_sentiment"
 
-    if protocol == "grpc":
+    if get_config().runtime.grpc.enabled:
         # Setup the client
         port = 8085
         channel = grpc.insecure_channel(f"localhost:{port}")
@@ -57,9 +57,9 @@ if __name__ == "__main__":
                 request, metadata=[("mm-model-id", model_id)], timeout=1
             )
             print("Text:", text)
-            print("RESPONSE:", response)
+            print("RESPONSE from gRPC:", response)
 
-    elif protocol == "http":
+    if get_config().runtime.http.enabled:
         port = 8080
         # Run inference for two sample prompts
         for text in ["I am not feeling well today!", "Today is a nice sunny day"]:
@@ -70,7 +70,4 @@ if __name__ == "__main__":
                 timeout=1,
             )
             print("\nText:", text)
-            print("RESPONSE:", json.dumps(response.json(), indent=4))
-
-    else:
-        print("--protocol must be one of [grpc, http]")
+            print("RESPONSE from HTTP:", json.dumps(response.json(), indent=4))
diff --git a/examples/text-sentiment/start_runtime.py b/examples/text-sentiment/start_runtime.py
index ea2ffb37..68dba668 100644
--- a/examples/text-sentiment/start_runtime.py
+++ b/examples/text-sentiment/start_runtime.py
@@ -14,28 +14,15 @@
 
 # Standard
 from os import path
-import argparse
 import sys
 
 # First Party
 import alog
 
 # Local
-from caikit.runtime import grpc_server, http_server
+from caikit.runtime.__main__ import main
 import caikit
 
-
-def protocol_arg():
-    parser = argparse.ArgumentParser(description="protocol switch")
-    parser.add_argument(
-        "--protocol", type=str, default="grpc", help="Specify a protocol: grpc or http"
-    )
-
-    args = parser.parse_args()
-    print(f"The specified protocol is: {args.protocol}")
-    return args.protocol
-
-
 if __name__ == "__main__":
     models_directory = path.abspath(path.join(path.dirname(__file__), "models"))
     caikit.config.configure(
@@ -44,6 +31,8 @@ if __name__ == "__main__":
             "runtime": {
                 "local_models_dir": models_directory,
                 "library": "text_sentiment",
+                "grpc": {"enabled": True},
+                "http": {"enabled": True},
             },
         }
     )
@@ -55,10 +44,4 @@ if __name__ == "__main__":
 
     alog.configure(default_level="debug")
 
-    protocol = protocol_arg()
-    if protocol == "grpc":
-        grpc_server.main()
-    elif protocol == "http":
-        http_server.main()
-    else:
-        print("--protocol must be one of [grpc, http]")
+    main()


Test patch:
diff --git a/tests/runtime/conftest.py b/tests/runtime/conftest.py
index 36491763..36364cc7 100644
--- a/tests/runtime/conftest.py
+++ b/tests/runtime/conftest.py
@@ -4,7 +4,7 @@ This sets up global test configs when pytest starts
 
 # Standard
 from contextlib import contextmanager
-from typing import Type, Union
+from typing import Dict, Type, Union
 import os
 import shlex
 import socket
@@ -15,9 +15,11 @@ import threading
 import time
 
 # Third Party
+from fastapi.testclient import TestClient
 from grpc_health.v1 import health_pb2, health_pb2_grpc
 import grpc
 import pytest
+import requests
 
 # First Party
 import alog
@@ -25,6 +27,7 @@ import alog
 # Local
 from caikit.core import MODEL_MANAGER
 from caikit.core.data_model.dataobject import render_dataobject_protos
+from caikit.runtime import http_server
 from caikit.runtime.grpc_server import RuntimeGRPCServer
 from caikit.runtime.model_management.loaded_model import LoadedModel
 from caikit.runtime.model_management.model_manager import ModelManager
@@ -56,6 +59,15 @@ def session_scoped_open_port():
     return _open_port()
 
 
+@pytest.fixture(scope="session")
+def http_session_scoped_open_port():
+    """Get an open port on localhost
+    Returns:
+        int: Available port
+    """
+    return _open_port()
+
+
 def _open_port():
     # 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.
@@ -138,17 +150,57 @@ def runtime_grpc_test_server(open_port, *args, **kwargs):
 
 @pytest.fixture(scope="session")
 def runtime_grpc_server(
-    session_scoped_open_port, sample_inference_service, sample_train_service
+    session_scoped_open_port,
 ) -> RuntimeGRPCServer:
     with runtime_grpc_test_server(
         session_scoped_open_port,
-        inference_service=sample_inference_service,
-        training_service=sample_train_service,
     ) as server:
         _check_server_readiness(server)
         yield server
 
 
+@contextmanager
+def runtime_http_test_server(open_port, *args, **kwargs):
+    """Helper to wrap creation of RuntimeHTTPServer in temporary configurations"""
+    with tempfile.TemporaryDirectory() as workdir:
+        temp_log_dir = os.path.join(workdir, "metering_logs")
+        temp_save_dir = os.path.join(workdir, "training_output")
+        os.makedirs(temp_log_dir)
+        os.makedirs(temp_save_dir)
+        with temp_config(
+            {
+                "runtime": {
+                    "metering": {"log_dir": temp_log_dir},
+                    "training": {"output_dir": temp_save_dir},
+                    "http": {"port": open_port},
+                }
+            },
+            "merge",
+        ):
+            config_overrides = {}
+            if "tls_config_override" in kwargs:
+                config_overrides = kwargs["tls_config_override"]
+                kwargs["tls_config_override"] = config_overrides["runtime"]["tls"]
+            with http_server.RuntimeHTTPServer(*args, **kwargs) as server:
+                _check_http_server_readiness(server, config_overrides)
+                # Give tests access to the workdir
+                server.workdir = workdir
+                yield server
+
+
+# For tests that need a working http server
+# For other http tests, we can use FastAPI TestClient
+@pytest.fixture(scope="session")
+def runtime_http_server(
+    http_session_scoped_open_port,
+) -> http_server.RuntimeHTTPServer:
+    """yields an actual running http server"""
+    with runtime_http_test_server(
+        http_session_scoped_open_port,
+    ) as server:
+        yield server
+
+
 @pytest.fixture(scope="session")
 def inference_stub(sample_inference_service, runtime_grpc_server) -> Type:
     inference_stub = sample_inference_service.stub_class(
@@ -321,5 +373,39 @@ def _check_server_readiness(server):
             stub.Check(health_check_request)
             done = True
         except grpc.RpcError:
-            log.debug("[RpcError]; will try to reconnect to test server in 0.1 second.")
-            time.sleep(0.1)
+            log.debug(
+                "[RpcError]; will try to reconnect to test server in 0.01 second."
+            )
+            time.sleep(0.01)
+
+
+def _check_http_server_readiness(server, config_overrides: Dict[str, Dict]):
+    mode = "http"
+    verify = None
+    cert = None
+    # tls
+    if config_overrides:
+        mode = "https"
+        verify = config_overrides["use_in_test"]["ca_cert"]
+        # mtls
+        if "client_cert" and "client_key" in config_overrides["use_in_test"]:
+            cert = (
+                config_overrides["use_in_test"]["client_cert"],
+                config_overrides["use_in_test"]["client_key"],
+            )
+    done = False
+    while not done:
+        try:
+            response = requests.get(
+                f"{mode}://localhost:{server.port}{http_server.HEALTH_ENDPOINT}",
+                verify=verify,
+                cert=cert,
+            )
+            assert response.status_code == 200
+            assert response.text == "OK"
+            done = True
+        except AssertionError:
+            log.debug(
+                "[HTTP server not ready]; will try to reconnect to test server in 0.01 second."
+            )
+            time.sleep(0.001)
diff --git a/tests/runtime/servicers/test_training_management_servicer.py b/tests/runtime/servicers/test_training_management_servicer.py
index 98e76122..7a541bd5 100644
--- a/tests/runtime/servicers/test_training_management_servicer.py
+++ b/tests/runtime/servicers/test_training_management_servicer.py
@@ -209,7 +209,7 @@ def test_training_raises_when_cancel_on_incorrect_id(training_management_service
     assert (
         "some_random_id not found in the list of currently running training jobs."
         in context.value.message
-        and "Could not perform cancel" in context.value.message
+        and "Did not perform cancel" in context.value.message
     )
 
 
diff --git a/tests/runtime/test_grpc_server.py b/tests/runtime/test_grpc_server.py
index 03ec0f58..9cbf3982 100644
--- a/tests/runtime/test_grpc_server.py
+++ b/tests/runtime/test_grpc_server.py
@@ -187,9 +187,7 @@ def test_model_train(runtime_grpc_server):
     assert result.learning_rate == 0.0015
 
 
-def test_components_preinitialized(
-    reset_globals, open_port, sample_inference_service, sample_train_service
-):
+def test_components_preinitialized(reset_globals, open_port):
     """Test that all model management components get pre-initialized when the
     server is instantiated
     """
@@ -198,8 +196,6 @@ def test_components_preinitialized(
     assert not MODEL_MANAGER._initializers
     with runtime_grpc_test_server(
         open_port,
-        inference_service=sample_inference_service,
-        training_service=sample_train_service,
     ):
         assert MODEL_MANAGER._trainers
         assert MODEL_MANAGER._finders
@@ -963,7 +959,7 @@ def test_canceling_model_loads_causes_exceptions(runtime_grpc_server):
         assert request_finished.is_set()
 
 
-def test_tls(sample_inference_service, open_port):
+def test_tls(open_port):
     """Boot up a server with TLS enabled and ping it on a secure channel"""
     ca_key = tls_test_tools.generate_key()[0]
     ca_cert = tls_test_tools.generate_ca_cert(ca_key)
@@ -974,14 +970,12 @@ def test_tls(sample_inference_service, open_port):
     )
     with runtime_grpc_test_server(
         open_port,
-        inference_service=sample_inference_service,
-        training_service=None,
         tls_config_override=tls_config,
     ) as server:
         _assert_connection(_make_secure_channel(server, ca_cert))
 
 
-def test_mtls(sample_inference_service, open_port):
+def test_mtls(open_port):
     """Boot up a server with mTLS enabled and ping it on a secure channel"""
     ca_key = tls_test_tools.generate_key()[0]
     ca_cert = tls_test_tools.generate_ca_cert(ca_key)
@@ -992,8 +986,6 @@ def test_mtls(sample_inference_service, open_port):
     )
     with runtime_grpc_test_server(
         open_port,
-        inference_service=sample_inference_service,
-        training_service=None,
         tls_config_override=tls_config,
     ) as server:
         client_key, client_cert = tls_test_tools.generate_derived_key_cert_pair(ca_key)
@@ -1009,7 +1001,7 @@ def test_mtls(sample_inference_service, open_port):
             stub.Check(health_check_request)
 
 
-def test_certs_can_be_loaded_as_files(sample_inference_service, tmp_path, open_port):
+def test_certs_can_be_loaded_as_files(tmp_path, open_port):
     """mTLS test with all tls configs loaded from files"""
     ca_key = tls_test_tools.generate_key()[0]
     ca_cert = tls_test_tools.generate_ca_cert(ca_key)
@@ -1033,8 +1025,6 @@ def test_certs_can_be_loaded_as_files(sample_inference_service, tmp_path, open_p
     )
     with runtime_grpc_test_server(
         open_port,
-        inference_service=sample_inference_service,
-        training_service=None,
         tls_config_override=tls_config,
     ) as server:
         client_key, client_cert = tls_test_tools.generate_derived_key_cert_pair(ca_key)
@@ -1050,8 +1040,6 @@ def test_metrics_stored_after_server_interrupt(
 
     with runtime_grpc_test_server(
         open_port,
-        inference_service=sample_inference_service,
-        training_service=None,
     ) as server:
         stub = sample_inference_service.stub_class(server.make_local_channel())
         predict_request = sample_inference_service.messages.SampleTaskRequest(
diff --git a/tests/runtime/test_http_server.py b/tests/runtime/test_http_server.py
index 2f5e3bef..4c3203dc 100644
--- a/tests/runtime/test_http_server.py
+++ b/tests/runtime/test_http_server.py
@@ -17,7 +17,7 @@ Tests for the caikit HTTP server
 # Standard
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import List
+from typing import Dict, List
 import json
 import os
 import signal
@@ -34,7 +34,7 @@ from caikit.core import DataObjectBase, dataobject
 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
+from tests.runtime.conftest import ModuleSubproc, runtime_http_test_server
 
 ## Helpers #####################################################################
 
@@ -67,7 +67,7 @@ class TLSConfig:
 @contextmanager
 def generate_tls_configs(
     port: int, tls: bool = False, mtls: bool = False, **http_config_overrides
-):
+) -> Dict[str, Dict]:
     """Helper to generate tls configs"""
     with tempfile.TemporaryDirectory() as workdir:
         config_overrides = {}
@@ -123,24 +123,22 @@ def generate_tls_configs(
 ## Insecure and TLS Tests #######################################################################
 
 
-def test_insecure_server(open_port):
+def test_insecure_server(runtime_http_server, open_port):
     with generate_tls_configs(open_port):
-        insecure_http_server = http_server.RuntimeHTTPServer()
         # start a non-blocking http server
-        with insecure_http_server:
-            resp = requests.get(f"http://localhost:{insecure_http_server.port}/docs")
-            resp.raise_for_status()
+        resp = requests.get(f"http://localhost:{runtime_http_server.port}/docs")
+        resp.raise_for_status()
 
 
 def test_basic_tls_server(open_port):
     with generate_tls_configs(
         open_port, tls=True, mtls=False, http_config_overrides={}
     ) as config_overrides:
-        http_server_with_tls = http_server.RuntimeHTTPServer(
-            tls_config_override=config_overrides["runtime"]["tls"]
-        )
-        # start a non-blocking http server with basic tls
-        with http_server_with_tls:
+        with runtime_http_test_server(
+            open_port,
+            tls_config_override=config_overrides,
+        ) as http_server_with_tls:
+            # start a non-blocking http server with basic tls
             resp = requests.get(
                 f"https://localhost:{http_server_with_tls.port}/docs",
                 verify=config_overrides["use_in_test"]["ca_cert"],
@@ -152,11 +150,11 @@ 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:
-        http_server_with_tls = http_server.RuntimeHTTPServer(
-            tls_config_override=config_overrides["runtime"]["tls"]
-        )
-        # start a non-blocking http server with basic tls
-        with http_server_with_tls:
+        with runtime_http_test_server(
+            open_port,
+            tls_config_override=config_overrides,
+        ) as http_server_with_tls:
+            # start a non-blocking http server with basic tls
             with pytest.raises(requests.exceptions.SSLError):
                 requests.get(
                     f"https://localhost:{http_server_with_tls.port}/docs",
@@ -168,11 +166,11 @@ def test_mutual_tls_server(open_port):
     with generate_tls_configs(
         open_port, tls=True, mtls=True, http_config_overrides={}
     ) as config_overrides:
-        http_server_with_mtls = http_server.RuntimeHTTPServer(
-            tls_config_override=config_overrides["runtime"]["tls"]
-        )
-        # start a non-blocking http server with mutual tls
-        with http_server_with_mtls:
+        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"],
@@ -188,11 +186,11 @@ 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:
-        http_server_with_mtls = http_server.RuntimeHTTPServer(
-            tls_config_override=config_overrides["runtime"]["tls"]
-        )
-        # start a non-blocking http server with mutual tls
-        with http_server_with_mtls:
+        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
             with pytest.raises(requests.exceptions.SSLError):
                 requests.get(
                     f"https://localhost:{http_server_with_mtls.port}/docs",
@@ -204,18 +202,23 @@ def test_mutual_tls_server_with_wrong_cert(open_port):
                 )
 
 
-def test_docs():
+def test_docs(runtime_http_server):
     """Simple check that pinging /docs returns 200"""
-    server = http_server.RuntimeHTTPServer()
-    with TestClient(server.app) as client:
+    with TestClient(runtime_http_server.app) as client:
         response = client.get("/docs")
         assert response.status_code == 200
 
 
-def test_inference_sample_task(sample_task_model_id):
+def test_docs_using_running_http_server(runtime_http_server):
+    """Simple check that pinging /docs returns 200
+    but pints the actual running server"""
+    response = requests.get(f"http://localhost:{runtime_http_server.port}/docs")
+    assert response.status_code == 200
+
+
+def test_inference_sample_task(sample_task_model_id, runtime_http_server):
     """Simple check that we can ping a model"""
-    server = http_server.RuntimeHTTPServer()
-    with TestClient(server.app) as client:
+    with TestClient(runtime_http_server.app) as client:
         json_input = {"inputs": {"name": "world"}}
         response = client.post(
             f"/api/v1/{sample_task_model_id}/task/sample",
@@ -226,10 +229,11 @@ def test_inference_sample_task(sample_task_model_id):
         assert json_response["greeting"] == "Hello world"
 
 
-def test_inference_sample_task_optional_field(sample_task_model_id):
+def test_inference_sample_task_optional_field(
+    sample_task_model_id, runtime_http_server
+):
     """Simple check for optional fields"""
-    server = http_server.RuntimeHTTPServer()
-    with TestClient(server.app) as client:
+    with TestClient(runtime_http_server.app) as client:
         json_input = {
             "inputs": {"name": "world"},
             "parameters": {"throw": True},
@@ -243,10 +247,9 @@ def test_inference_sample_task_optional_field(sample_task_model_id):
         assert response.status_code == 500
 
 
-def test_inference_other_task(other_task_model_id):
+def test_inference_other_task(other_task_model_id, runtime_http_server):
     """Simple check that we can ping a model"""
-    server = http_server.RuntimeHTTPServer()
-    with TestClient(server.app) as client:
+    with TestClient(runtime_http_server.app) as client:
         json_input = {"inputs": {"name": "world"}}
         response = client.post(
             f"/api/v1/{other_task_model_id}/task/other",
@@ -257,10 +260,9 @@ def test_inference_other_task(other_task_model_id):
         assert json_response["farewell"] == "goodbye: world 42 times"
 
 
-def test_inference_streaming_sample_module(sample_task_model_id):
+def test_inference_streaming_sample_module(sample_task_model_id, runtime_http_server):
     """Simple check for testing a happy path unary-stream case"""
-    server = http_server.RuntimeHTTPServer()
-    with TestClient(server.app) as client:
+    with TestClient(runtime_http_server.app) as client:
         json_input = {"inputs": {"name": "world"}}
         stream = client.post(
             f"/api/v1/{sample_task_model_id}/task/server-streaming-sample",
@@ -276,10 +278,9 @@ def test_inference_streaming_sample_module(sample_task_model_id):
         )
 
 
-def test_model_not_found():
+def test_model_not_found(runtime_http_server):
     """Simple check that we can ping a model"""
-    server = http_server.RuntimeHTTPServer()
-    with TestClient(server.app) as client:
+    with TestClient(runtime_http_server.app) as client:
         response = client.post(
             f"/api/v1/this_is_not_a_model/task/sample",
             json={"inputs": {"name": "world"}},
@@ -287,10 +288,11 @@ def test_model_not_found():
         assert response.status_code == 404
 
 
-def test_inference_sample_task_throws_incorrect_input(sample_task_model_id):
+def test_inference_sample_task_throws_incorrect_input(
+    sample_task_model_id, runtime_http_server
+):
     """error check for a request with incorrect input"""
-    server = http_server.RuntimeHTTPServer()
-    with TestClient(server.app) as client:
+    with TestClient(runtime_http_server.app) as client:
         json_input = {"blah": {"sample_input": {"name": "world"}}}
         response = client.post(
             f"/api/v1/{sample_task_model_id}/task/sample",
@@ -299,10 +301,9 @@ def test_inference_sample_task_throws_incorrect_input(sample_task_model_id):
         assert response.status_code == 400
 
 
-def test_health_check_ok():
+def test_health_check_ok(runtime_http_server):
     """Make sure the health check returns OK"""
-    server = http_server.RuntimeHTTPServer()
-    with TestClient(server.app) as client:
+    with TestClient(runtime_http_server.app) as client:
         response = client.get(http_server.HEALTH_ENDPOINT)
         assert response.status_code == 200
         assert response.text == "OK"
@@ -324,7 +325,7 @@ def test_pydantic_wrapping_with_enums():
     assert token.text == "foo"
 
 
-def test_pydantic_wrapping_with_lists():
+def test_pydantic_wrapping_with_lists(runtime_http_server):
     """Check that pydantic wrapping works on data models with lists"""
 
     @dataobject(package="http")
@@ -338,7 +339,7 @@ def test_pydantic_wrapping_with_lists():
     foo = FooTest(bars=[BarTest(1)])
     assert foo.bars[0].baz == 1
 
-    http_server.RuntimeHTTPServer._dataobject_to_pydantic(FooTest)
+    runtime_http_server._dataobject_to_pydantic(FooTest)
 
     foo = FooTest(bars=[BarTest(1)])
     assert foo.bars[0].baz == 1