For live reload, retry if model is not ready

truss/templates/control/control/endpoints.py

@@ -1,5 +1,6 @@
 import requests
 from flask import Blueprint, Response, current_app, jsonify, request
+from helpers.errors import ModelNotReady
 from requests.exceptions import ConnectionError
 from tenacity import Retrying, retry_if_exception_type, stop_after_attempt, wait_fixed
 
@@ -23,7 +24,10 @@ def proxy(path):
 
     # Wait a bit for inference server to start
     for attempt in Retrying(
-        retry=retry_if_exception_type(ConnectionError),
+        retry=(
+            retry_if_exception_type(ConnectionError)
+            | retry_if_exception_type(ModelNotReady)
+        ),
         stop=stop_after_attempt(INFERENCE_SERVER_START_WAIT_SECS),
         wait=wait_fixed(1),
     ):
@@ -36,6 +40,8 @@ def proxy(path):
                     cookies=request.cookies,
                     headers=request.headers,
                 )
+                if _is_model_not_ready(resp):
+                    raise ModelNotReady("Model has started running, but not ready yet.")
             except ConnectionError as exp:
                 # This check is a bit expensive so we don't do it before every request, we
                 # do it only if request fails with connection error. If the inference server
@@ -89,3 +95,11 @@ def has_partially_applied_patch():
 def stop_inference_server():
     current_app.config["inference_server_controller"].stop()
     return {"msg": "Inference server stopped successfully"}
+
+
+def _is_model_not_ready(resp) -> bool:
+    return (
+        resp.status_code == 503
+        and resp.content is not None
+        and "model is not ready" in resp.content.decode("utf-8")
+    )

truss/templates/control/control/helpers/errors.py

@@ -34,3 +34,9 @@ class InadmissiblePatch(PatchApplicatonError):
     """Patch does not apply to current state of Truss."""
 
     pass
+
+
+class ModelNotReady(Error):
+    """Model has started running, but not ready yet."""
+
+    pass

truss/test_data/truss_container_fs/app/common/logging.py

@@ -0,0 +1,47 @@
+import logging
+import sys
+
+from pythonjsonlogger import jsonlogger
+
+LEVEL = logging.INFO
+
+JSON_LOG_HANDLER = logging.StreamHandler(stream=sys.stderr)
+JSON_LOG_HANDLER.set_name("json_logger_handler")
+JSON_LOG_HANDLER.setLevel(LEVEL)
+JSON_LOG_HANDLER.setFormatter(
+    jsonlogger.JsonFormatter("%(asctime)s %(levelname)s %(message)s")
+)
+
+
+class HealthCheckFilter(logging.Filter):
+    def filter(self, record: logging.LogRecord) -> bool:
+        # for any health check endpoints, lets skip logging
+        return (
+            record.getMessage().find("GET / ") == -1
+            and record.getMessage().find("GET /v1/models/model ") == -1
+        )
+
+
+def setup_logging():
+    loggers = [logging.getLogger()] + [
+        logging.getLogger(name) for name in logging.root.manager.loggerDict
+    ]
+
+    for logger in loggers:
+        logger.setLevel(LEVEL)
+        logger.propagate = False
+
+        setup = False
+
+        # let's not thrash the handlers unnecessarily
+        for handler in logger.handlers:
+            if handler.name == JSON_LOG_HANDLER.name:
+                setup = True
+
+        if not setup:
+            logger.handlers.clear()
+            logger.addHandler(JSON_LOG_HANDLER)
+
+        # some special handling for request logging
+        if logger.name == "uvicorn.access":
+            logger.addFilter(HealthCheckFilter())

truss/test_data/truss_container_fs/app/common/serialization.py

@@ -3,10 +3,6 @@
 import json
 import uuid
 
-import msgpack
-import msgpack_numpy as mp_np
-import numpy as np
-
 
 # mostly cribbed from django.core.serializer.DjangoJSONEncoder
 def truss_msgpack_encoder(obj, chain=None):
@@ -61,6 +57,8 @@ def truss_msgpack_decoder(obj, chain=None):
 
 # this json object is JSONType + np.array + datetime
 def is_truss_serializable(obj):
+    import numpy as np
+
     # basic JSON types
     if isinstance(obj, (str, int, float, bool, type(None), dict, list)):
         return True
@@ -75,19 +73,27 @@ def is_truss_serializable(obj):
 
 
 def truss_msgpack_serialize(obj):
+    import msgpack
+    import msgpack_numpy as mp_np
+
     return msgpack.packb(
         obj, default=lambda x: truss_msgpack_encoder(x, chain=mp_np.encode)
     )
 
 
 def truss_msgpack_deserialize(obj):
+    import msgpack
+    import msgpack_numpy as mp_np
+
     return msgpack.unpackb(
         obj, object_hook=lambda x: truss_msgpack_decoder(x, chain=mp_np.decode)
     )
 
 
 class DeepNumpyEncoder(json.JSONEncoder):
     def default(self, obj):
+        import numpy as np
+
         if isinstance(obj, np.integer):
             return int(obj)
         elif isinstance(obj, np.floating):