Refresh logger after loading user module (#563)

* Refresh logs after loading modules and calling init functions

* Fix ONNX import

* Create loggers at nanotime

(cherry picked from commit 01808ed)

Author: vishalbollu

pkg/workloads/cortex/downloader/download.py

@@ -19,9 +19,7 @@
 
 from cortex.lib import util
 from cortex.lib.storage import S3
-from cortex.lib.log import get_logger
-
-logger = get_logger()
+from cortex.lib.log import cx_logger
 
 
 def start(args):
@@ -34,12 +32,12 @@ def start(args):
         s3_client = S3(bucket_name, client_config={})
 
         if item_name != "":
-            logger.info("downloading {} from {}".format(item_name, from_path))
+            cx_logger().info("downloading {} from {}".format(item_name, from_path))
         s3_client.download(prefix, to_path)
 
         if download_arg.get("unzip", False):
             if item_name != "":
-                logger.info("unzipping {}".format(item_name))
+                cx_logger().info("unzipping {}".format(item_name))
             util.extract_zip(
                 os.path.join(to_path, os.path.basename(from_path)), delete_zip_file=True
             )

pkg/workloads/cortex/lib/api_utils.py

@@ -18,9 +18,7 @@
 import time
 
 from cortex.lib.exceptions import UserException, CortexException
-from cortex.lib.log import get_logger
-
-logger = get_logger()
+from cortex.lib.log import cx_logger
 
 
 def get_classes(ctx, api_name):
@@ -157,10 +155,10 @@ def post_request_metrics(ctx, api, response, prediction_payload, start_time, cla
 
                 metrics_list += prediction_metrics(api_dimensions, api, prediction)
             except Exception as e:
-                logger.warn("unable to record prediction metric", exc_info=True)
+                cx_logger().warn("unable to record prediction metric", exc_info=True)
 
     metrics_list += latency_metric(api_dimensions, start_time)
     try:
         ctx.publish_metrics(metrics_list)
     except Exception as e:
-        logger.warn("failure encountered while publishing metrics", exc_info=True)
+        cx_logger().warn("failure encountered while publishing metrics", exc_info=True)

pkg/workloads/cortex/lib/context.py

@@ -22,12 +22,10 @@
 
 from cortex import consts
 from cortex.lib import util
+from cortex.lib.log import refresh_logger
 from cortex.lib.storage import S3, LocalStorage
 from cortex.lib.exceptions import CortexException, UserException
 from cortex.lib.resources import ResourceMap
-from cortex.lib.log import get_logger
-
-logger = get_logger()
 
 
 class Context:
@@ -149,6 +147,8 @@ def get_request_handler_impl(self, api_name, project_dir):
         except CortexException as e:
             e.wrap("api " + api_name, "failed to load request_handler", request_handler_path)
             raise
+        finally:
+            refresh_logger()
 
         try:
             _validate_impl(impl, REQUEST_HANDLER_IMPL_VALIDATION)
@@ -163,9 +163,12 @@ def get_predictor_impl(self, api_name, project_dir):
             impl = self.load_module(
                 "predictor", api["name"], os.path.join(project_dir, api["predictor"]["path"])
             )
+
         except CortexException as e:
             e.wrap("api " + api_name, "failed to load predictor", api["predictor"]["path"])
             raise
+        finally:
+            refresh_logger()
 
         try:
             _validate_impl(impl, PREDICTOR_IMPL_VALIDATION)

pkg/workloads/cortex/lib/log.py

@@ -14,6 +14,7 @@
 
 import logging
 import sys
+import time
 
 from cortex.lib import stringify
 import datetime as dt
@@ -28,23 +29,39 @@ def formatTime(self, record, datefmt):
         return s
 
 
-logger = logging.getLogger("cortex")
-handler = logging.StreamHandler(stream=sys.stdout)
-formatter = MyFormatter(
-    fmt="%(asctime)s:%(name)s:%(levelname)s:%(message)s", datefmt="%Y-%m-%d %H:%M:%S.%f"
-)
-handler.setFormatter(formatter)
+current_logger = None
 
-logger.addHandler(handler)
-logger.setLevel(logging.DEBUG)
+
+def register_logger(name):
+    logger = logging.getLogger(name)
+    handler = logging.StreamHandler(stream=sys.stdout)
+    formatter = MyFormatter(
+        fmt="%(asctime)s:cortex:%(levelname)s:%(message)s", datefmt="%Y-%m-%d %H:%M:%S.%f"
+    )
+    handler.setFormatter(formatter)
+
+    logger.propagate = False
+    logger.addHandler(handler)
+    logger.setLevel(logging.DEBUG)
+    return logger
+
+
+def refresh_logger():
+    global current_logger
+    if current_logger is not None:
+        current_logger.disabled = True
+    current_logger = register_logger("{}-cortex".format(int(time.time() * 1000000)))
+
+
+def cx_logger():
+    return current_logger
 
 
 def debug_obj(name, sample, debug):
     if not debug:
         return
 
-    logger.info("{}: {}".format(name, stringify.truncate(sample)))
+    cx_logger().info("{}: {}".format(name, stringify.truncate(sample)))
 
 
-def get_logger():
-    return logging.getLogger("cortex")
+refresh_logger()

pkg/workloads/cortex/lib/util.py

@@ -27,14 +27,10 @@
 from copy import deepcopy
 from datetime import datetime
 
-from cortex.lib.log import get_logger
 from cortex.lib import stringify
 import json_tricks
 
 
-logger = get_logger()
-
-
 def isclose(a, b, rel_tol=1e-09, abs_tol=0.0):
     return abs(a - b) <= max(rel_tol * max(abs(a), abs(b)), abs_tol)
 

pkg/workloads/cortex/onnx_serve/api.py

@@ -24,13 +24,10 @@
 import numpy as np
 
 from cortex.lib import util, Context, api_utils
-from cortex.lib.log import get_logger, debug_obj
+from cortex.lib.log import cx_logger, debug_obj
 from cortex.lib.exceptions import CortexException, UserRuntimeException, UserException
 from cortex.lib.stringify import truncate
 
-logger = get_logger()
-logger.propagate = False  # prevent double logging (flask modifies root logger)
-
 app = Flask(__name__)
 
 app.json_encoder = util.json_tricks_encoder
@@ -79,7 +76,7 @@ def after_request(response):
     api = local_cache["api"]
     ctx = local_cache["ctx"]
 
-    logger.info(response.status)
+    cx_logger().info(response.status)
 
     prediction = None
     if "prediction" in g:
@@ -94,7 +91,7 @@ def after_request(response):
 
 def prediction_failed(reason):
     message = "prediction failed: {}".format(reason)
-    logger.error(message)
+    cx_logger().error(message)
     return message, status.HTTP_406_NOT_ACCEPTABLE
 
 
@@ -215,7 +212,7 @@ def predict():
 
             debug_obj("post_inference", result, debug)
     except Exception as e:
-        logger.exception("prediction failed")
+        cx_logger().exception("prediction failed")
         return prediction_failed(str(e))
 
     g.prediction = result
@@ -238,7 +235,7 @@ def get_signature():
 
 @app.errorhandler(Exception)
 def exceptions(e):
-    logger.exception(e)
+    cx_logger().exception(e)
     return jsonify(error=str(e)), 500
 
 
@@ -262,45 +259,45 @@ def start(args):
         request_handler = local_cache.get("request_handler")
 
         if request_handler is not None and util.has_function(request_handler, "pre_inference"):
-            logger.info(
+            cx_logger().info(
                 "using pre_inference request handler provided in {}".format(
                     api["onnx"]["request_handler"]
                 )
             )
         else:
-            logger.info("pre_inference request handler not found")
+            cx_logger().info("pre_inference request handler not found")
 
         if request_handler is not None and util.has_function(request_handler, "post_inference"):
-            logger.info(
+            cx_logger().info(
                 "using post_inference request handler provided in {}".format(
                     api["onnx"]["request_handler"]
                 )
             )
         else:
-            logger.info("post_inference request handler not found")
+            cx_logger().info("post_inference request handler not found")
 
         sess = rt.InferenceSession(model_path)
         local_cache["sess"] = sess
         local_cache["input_metadata"] = sess.get_inputs()
-        logger.info(
+        cx_logger().info(
             "input_metadata: {}".format(truncate(extract_signature(local_cache["input_metadata"])))
         )
         local_cache["output_metadata"] = sess.get_outputs()
-        logger.info(
+        cx_logger().info(
             "output_metadata: {}".format(
                 truncate(extract_signature(local_cache["output_metadata"]))
             )
         )
 
     except Exception as e:
-        logger.exception("failed to start api")
+        cx_logger().exception("failed to start api")
         sys.exit(1)
 
     if api.get("tracker") is not None and api["tracker"].get("model_type") == "classification":
         try:
             local_cache["class_set"] = api_utils.get_classes(ctx, api["name"])
         except Exception as e:
-            logger.warn("an error occurred while attempting to load classes", exc_info=True)
+            cx_logger().warn("an error occurred while attempting to load classes", exc_info=True)
 
     serve(app, listen="*:{}".format(args.port))
 

pkg/workloads/cortex/predictor_serve/api.py

@@ -22,12 +22,9 @@
 from waitress import serve
 
 from cortex.lib import util, Context, api_utils
-from cortex.lib.log import get_logger, debug_obj
+from cortex.lib.log import cx_logger, debug_obj, refresh_logger
 from cortex.lib.exceptions import CortexException, UserRuntimeException
 
-logger = get_logger()
-logger.propagate = False  # prevent double logging (flask modifies root logger)
-
 app = Flask(__name__)
 
 app.json_encoder = util.json_tricks_encoder
@@ -51,7 +48,7 @@ def after_request(response):
     api = local_cache["api"]
     ctx = local_cache["ctx"]
 
-    logger.info(response.status)
+    cx_logger().info(response.status)
 
     prediction = None
     if "prediction" in g:
@@ -66,7 +63,7 @@ def after_request(response):
 
 def prediction_failed(reason):
     message = "prediction failed: {}".format(reason)
-    logger.error(message)
+    cx_logger().error(message)
     return message, status.HTTP_406_NOT_ACCEPTABLE
 
 
@@ -95,7 +92,7 @@ def predict():
         except Exception as e:
             raise UserRuntimeException(api["predictor"]["path"], "predict", str(e)) from e
     except Exception as e:
-        logger.exception("prediction failed")
+        cx_logger().exception("prediction failed")
         return prediction_failed(str(e))
 
     g.prediction = output
@@ -104,7 +101,7 @@ def predict():
 
 @app.errorhandler(Exception)
 def exceptions(e):
-    logger.exception(e)
+    cx_logger().exception(e)
     return jsonify(error=str(e)), 500
 
 
@@ -133,16 +130,18 @@ def start(args):
                 local_cache["predictor"].init(model_path, api["predictor"]["metadata"])
             except Exception as e:
                 raise UserRuntimeException(api["predictor"]["path"], "init", str(e)) from e
-        logger.info("init ran successfully")
+            finally:
+                refresh_logger()
     except:
-        logger.exception("failed to start api")
+        cx_logger().exception("failed to start api")
         sys.exit(1)
 
+    cx_logger().info("init ran successfully")
     if api.get("tracker") is not None and api["tracker"].get("model_type") == "classification":
         try:
             local_cache["class_set"] = api_utils.get_classes(ctx, api["name"])
         except Exception as e:
-            logger.warn("an error occurred while attempting to load classes", exc_info=True)
+            cx_logger().warn("an error occurred while attempting to load classes", exc_info=True)
 
     serve(app, listen="*:{}".format(args.port))