diff --git a/pkg/workloads/cortex/onnx_serve/api.py b/pkg/workloads/cortex/onnx_serve/api.py index fbfd940650..59c36e7165 100644 --- a/pkg/workloads/cortex/onnx_serve/api.py +++ b/pkg/workloads/cortex/onnx_serve/api.py @@ -248,6 +248,21 @@ def start(args): local_cache["request_handler"] = ctx.get_request_handler_impl( api["name"], args.project_dir ) + request_handler = local_cache.get("request_handler") + + if request_handler is not None and util.has_function(request_handler, "pre_inference"): + logger.info( + "using pre_inference request handler provided in {}".format(api["request_handler"]) + ) + else: + logger.info("pre_inference request handler not found") + + if request_handler is not None and util.has_function(request_handler, "post_inference"): + logger.info( + "using post_inference request handler provided in {}".format(api["request_handler"]) + ) + else: + logger.info("post_inference request handler not found") sess = rt.InferenceSession(model_path) local_cache["sess"] = sess diff --git a/pkg/workloads/cortex/tf_api/api.py b/pkg/workloads/cortex/tf_api/api.py index bf78f3a4ec..f9fe5e104a 100644 --- a/pkg/workloads/cortex/tf_api/api.py +++ b/pkg/workloads/cortex/tf_api/api.py @@ -375,6 +375,22 @@ def start(args): local_cache["request_handler"] = ctx.get_request_handler_impl( api["name"], args.project_dir ) + request_handler = local_cache.get("request_handler") + + if request_handler is not None and util.has_function(request_handler, "pre_inference"): + logger.info( + "using pre_inference request handler provided in {}".format(api["request_handler"]) + ) + else: + logger.info("pre_inference request handler not found") + + if request_handler is not None and util.has_function(request_handler, "post_inference"): + logger.info( + "using post_inference request handler provided in {}".format(api["request_handler"]) + ) + else: + logger.info("post_inference request handler not found") + except Exception as e: logger.exception("failed to start api") sys.exit(1)