minor refactoring

pyinfra/callback.py

@@ -31,7 +31,7 @@ class Callback:
         return pipeline
 
     @staticmethod
-    def __run_pipeline(pipeline, body):
+    def __run_pipeline(pipeline, analysis_input: dict):
         """
         TODO: Since data and metadata are passed as singletons, there is no buffering and hence no batching happening
          within the pipeline. However, the queue acknowledgment logic needs to be changed in order to facilitate
@@ -43,19 +43,19 @@ class Callback:
          operates on singletons ([data], [metadata]).
         """
 
-        def combine_storage_item_metadata_with_queue_message_metadata(body):
-            return merge(body["metadata"], omit(body, ["data", "metadata"]))
+        def combine_storage_item_metadata_with_queue_message_metadata(analysis_input):
+            return merge(analysis_input["metadata"], omit(analysis_input, ["data", "metadata"]))
 
-        def remove_queue_message_metadata(result):
-            metadata = omit(result["metadata"], queue_message_keys(body))
-            return {**result, "metadata": metadata}
+        def remove_queue_message_metadata(analysis_result):
+            metadata = omit(analysis_result["metadata"], queue_message_keys(analysis_input))
+            return {**analysis_result, "metadata": metadata}
 
-        def queue_message_keys(body):
-            return {*body.keys()}.difference({"data", "metadata"})
+        def queue_message_keys(analysis_input):
+            return {*analysis_input.keys()}.difference({"data", "metadata"})
 
         try:
-            data = body["data"]
-            metadata = combine_storage_item_metadata_with_queue_message_metadata(body)
+            data = analysis_input["data"]
+            metadata = combine_storage_item_metadata_with_queue_message_metadata(analysis_input)
             analysis_response_stream = pipeline([data], [metadata])
             analysis_response_stream = lmap(remove_queue_message_metadata, analysis_response_stream)
             return analysis_response_stream
@@ -64,13 +64,15 @@ class Callback:
             logger.error(err)
             raise AnalysisFailure from err
 
-    def __call__(self, body: dict):
-        operation = body.get("operation", "submit")
+    def __call__(self, analysis_input: dict):
+        """data_metadata_pack: {'dossierId': ..., 'fileId': ..., 'pages': ..., 'operation': ...}"""
+        # TODO: hardcoding of "submit" is complectig callback with analysis server implementation detail
+        operation = analysis_input.get("operation", "submit")
         endpoint = self.__make_endpoint(operation)
         pipeline = self.__get_pipeline(endpoint)
 
         try:
             logging.debug(f"Requesting analysis from {endpoint}...")
-            return self.__run_pipeline(pipeline, body)
+            return self.__run_pipeline(pipeline, analysis_input)
         except AnalysisFailure:
             logging.warning(f"Exception caught when calling analysis endpoint {endpoint}.")

pyinfra/server/server.py

@@ -1,4 +1,4 @@
-from functools import singledispatch
+from functools import singledispatch, lru_cache
 from typing import Dict, Callable, Union
 
 from flask import Flask, jsonify, request
@@ -51,17 +51,20 @@ def __stream_fn_to_processing_server(operation2stream_fn: dict, buffer_size):
 
 
 def __set_up_processing_server(operation2function: Dict[str, QueuedStreamFunction]):
-    app = Flask(__name__)
-    registry = CollectorRegistry(auto_describe=True)
-
-    operation2processor = {
-        op: LazyRestProcessor(fn, **build_endpoint_suffixes(op)) for op, fn in operation2function.items()
-    }
+    @lru_cache(maxsize=None)
+    def get_registry():
+        return CollectorRegistry(auto_describe=True)
 
     def make_summary_instance(op: str):
         op = op.replace("_pickup", "") if op != "pickup" else "default"
         service_name = CONFIG.service.name
-        return Summary(f"redactmanager_{service_name}_{op}_seconds", f"Time spent on {op}.", registry=registry)
+        return Summary(f"redactmanager_{service_name}_{op}_seconds", f"Time spent on {op}.", registry=get_registry())
+
+    app = Flask(__name__)
+
+    operation2processor = {
+        op: LazyRestProcessor(fn, **build_endpoint_suffixes(op)) for op, fn in operation2function.items()
+    }
 
     submit_operation2processor = {submit_suffix(op): prc for op, prc in operation2processor.items()}
     pickup_operation2processor = {pickup_suffix(op): prc for op, prc in operation2processor.items()}
@@ -83,7 +86,7 @@ def __set_up_processing_server(operation2function: Dict[str, QueuedStreamFunctio
 
     @app.route("/prometheus", methods=["GET"])
     def prometheus():
-        return generate_latest(registry=registry)
+        return generate_latest(registry=get_registry())
 
     @app.route("/<operation>", methods=["POST", "PATCH"])
     def submit(operation):