Repository: caikit/caikit
Pull request #376 (2023-07-27 00:08:44)

♻️ refactor init for http (see #389 instead)
<!--  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 refactors the `init` function of the `http_server` to align with `grpc_server`. The idea is to eventually have a common `caikt/runtime/__main__.py` which can start both servers with similar arguments.


Alternate implementation here:  https://github.com/caikit/caikit/pull/389

**Special notes for your reviewer**:

Also fixes https://github.com/caikit/caikit/issues/380

**If applicable**:
- [ ] this PR contains documentation
- [ ] this PR contains unit tests
- [ ] this PR has been tested for backwards compatibility



Patch:
diff --git a/caikit/runtime/dump_services.py b/caikit/runtime/dump_services.py
index 0ebbd1d4..410d8579 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
 
 
@@ -55,7 +56,14 @@ def dump_http_services(output_dir: str):
         RuntimeHTTPServer,
     )
 
-    server = RuntimeHTTPServer()
+    inf_svc = ServicePackageFactory.get_service_package(
+        ServicePackageFactory.ServiceType.INFERENCE,
+    )
+    train_svc = ServicePackageFactory.get_service_package(
+        ServicePackageFactory.ServiceType.TRAINING,
+    )
+
+    server = RuntimeHTTPServer(inference_service=inf_svc, training_service=train_svc)
     with TestClient(server.app) as client:
         response = client.get("/openapi.json")
         response.raise_for_status()
@@ -71,5 +79,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..259f8578 100644
--- a/caikit/runtime/grpc_server.py
+++ b/caikit/runtime/grpc_server.py
@@ -66,7 +66,7 @@ class RuntimeGRPCServer(RuntimeServerBase):
     def __init__(
         self,
         inference_service: ServicePackage,
-        training_service: Optional[ServicePackage],
+        training_service: Optional[ServicePackage] = None,
         tls_config_override: Optional[aconfig.Config] = None,
         handle_terminations: bool = False,
     ):
diff --git a/caikit/runtime/http_server.py b/caikit/runtime/http_server.py
index c84406ef..908c67a4 100644
--- a/caikit/runtime/http_server.py
+++ b/caikit/runtime/http_server.py
@@ -106,15 +106,18 @@ class RuntimeHTTPServer(RuntimeServerBase):
     ## Interface ##
     ###############
 
-    def __init__(self, tls_config_override: Optional[aconfig.Config] = None):
+    def __init__(
+        self,
+        inference_service: ServicePackage,
+        #  # pylint: disable=unused-argument
+        training_service: Optional[ServicePackage] = None,
+        tls_config_override: Optional[aconfig.Config] = None,
+    ):
         super().__init__(get_config().runtime.http.port, tls_config_override)
 
         self.app = FastAPI()
 
         # Set up the central predict servicer
-        inference_service = ServicePackageFactory().get_service_package(
-            ServicePackageFactory.ServiceType.INFERENCE,
-        )
         self.global_predict_servicer = GlobalPredictServicer(inference_service)
 
         # Set up the central train servicer
@@ -179,11 +182,11 @@ class RuntimeHTTPServer(RuntimeServerBase):
         self._uvicorn_server_thread = None
 
     def __del__(self):
-        if get_config().runtime.metering.enabled:
+        if get_config() and get_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
@@ -328,6 +331,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 +397,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]
             )
@@ -554,8 +557,32 @@ class RuntimeHTTPServer(RuntimeServerBase):
 
 
 def main():
+
     caikit.core.toolkit.logging.configure()
-    server = RuntimeHTTPServer()
+    log.debug("Starting up caikit.runtime.http_server")
+
+    # This code will live in caikt/runtime/__main__.py eventually
+    # pylint: disable=duplicate-code
+
+    # 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 = RuntimeHTTPServer(
+        inference_service=inference_service, training_service=training_service
+    )
     server.start()
 
 
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


Test patch:
diff --git a/tests/runtime/conftest.py b/tests/runtime/conftest.py
index 36491763..d671a754 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.
@@ -149,6 +161,52 @@ def runtime_grpc_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
+
+
+# This is honestly not required for fastAPI testing
+# as they have a built-in TestClient that's capable
+# of testing http routes. I've kept it here for tests that
+# need a working http server
+@pytest.fixture(scope="session")
+def runtime_http_server(
+    http_session_scoped_open_port, sample_inference_service, sample_train_service
+) -> http_server.RuntimeHTTPServer:
+    """yields an actual running http server"""
+    with runtime_http_test_server(
+        http_session_scoped_open_port,
+        inference_service=sample_inference_service,
+        training_service=sample_train_service,
+    ) 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 +379,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.01)
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_http_server.py b/tests/runtime/test_http_server.py
index 2f5e3bef..2b692dc0 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,24 @@ 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):
+def test_basic_tls_server(sample_inference_service, sample_train_service, 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,
+            inference_service=sample_inference_service,
+            training_service=sample_train_service,
+            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"],
@@ -148,15 +148,19 @@ def test_basic_tls_server(open_port):
             resp.raise_for_status()
 
 
-def test_basic_tls_server_with_wrong_cert(open_port):
+def test_basic_tls_server_with_wrong_cert(
+    sample_inference_service, sample_train_service, 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,
+            inference_service=sample_inference_service,
+            training_service=sample_train_service,
+            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",
@@ -164,15 +168,17 @@ def test_basic_tls_server_with_wrong_cert(open_port):
                 )
 
 
-def test_mutual_tls_server(open_port):
+def test_mutual_tls_server(sample_inference_service, sample_train_service, 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,
+            inference_service=sample_inference_service,
+            training_service=sample_train_service,
+            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"],
@@ -184,15 +190,19 @@ def test_mutual_tls_server(open_port):
             resp.raise_for_status()
 
 
-def test_mutual_tls_server_with_wrong_cert(open_port):
+def test_mutual_tls_server_with_wrong_cert(
+    sample_inference_service, sample_train_service, 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,
+            inference_service=sample_inference_service,
+            training_service=sample_train_service,
+            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 +214,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 +241,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 +259,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 +272,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 +290,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 +300,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 +313,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 +337,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 +351,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