Merge branch 'master' into containerized_tests

banner.txt

@@ -1,8 +1,6 @@
-______     _____       __
-| ___ \   |_   _|     / _|
-| |_/ /   _ | | _ __ | |_ _ __ __ _
-|  __/ | | || || '_ \|  _| '__/ _` |
-| |  | |_| || || | | | | | | | (_| |
-\_|   \__, \___/_| |_|_| |_|  \__,_|
-       __/ |
-      |___/
+             ___    _  _    ___               __
+    o O O   | _ \  | || |  |_ _|   _ _       / _|    _ _   __ _
+   o        |  _/   \_, |   | |   | ' \     |  _|   | '_| / _` |
+  TS__[O]  _|_|_   _|__/   |___|  |_||_|   _|_|_   _|_|_  \__,_|
+ {======|_| ``` |_| ````|_|`````|_|`````|_|`````|_|`````|_|`````|
+./o--000' `-0-0-' `-0-0-' `-0-0-' `-0-0-' `-0-0-' `-0-0-' `-0-0-'

pyinfra/flask.py

@@ -6,7 +6,6 @@ from waitress import serve
 
 from pyinfra.config import CONFIG
 
-
 logger = logging.getLogger(__file__)
 logger.setLevel(CONFIG.service.logging_level)
 
@@ -31,6 +30,7 @@ def run_probing_webserver(app, host=None, port=None, mode=None):
 def set_up_probing_webserver():
     # TODO: implement meaningful checks
     app = Flask(__name__)
+    informed_about_missing_prometheus_endpoint = False
 
     @app.route("/ready", methods=["GET"])
     def ready():
@@ -46,12 +46,19 @@ def set_up_probing_webserver():
 
     @app.route("/prometheus", methods=["GET"])
     def get_metrics_from_analysis_endpoint():
+        nonlocal informed_about_missing_prometheus_endpoint
         try:
             resp = requests.get(f"{CONFIG.rabbitmq.callback.analysis_endpoint}/prometheus")
             resp.raise_for_status()
-            return resp.text
-        except Exception as err:
-            logger.warning(f"Got no metrics from analysis prometheus endpoint: {err}")
-            return resp
+        except ConnectionError:
+            return ""
+        except requests.exceptions.HTTPError as err:
+            if resp.status_code == 404:
+                if not informed_about_missing_prometheus_endpoint:
+                    logger.warning(f"Got no metrics from analysis prometheus endpoint: {err}")
+                    informed_about_missing_prometheus_endpoint = True
+            else:
+                logging.warning(f"Caught {err}")
+        return resp.text
 
     return app

pyinfra/queue/queue_manager/pika_queue_manager.py

@@ -78,8 +78,8 @@ class PikaQueueManager(QueueManager):
 
         args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": dead_letter_queue}
 
-        self.channel.queue_declare(input_queue, arguments=args, auto_delete=False)
-        self.channel.queue_declare(output_queue, arguments=args, auto_delete=False)
+        self.channel.queue_declare(input_queue, arguments=args, auto_delete=False, durable=True)
+        self.channel.queue_declare(output_queue, arguments=args, auto_delete=False, durable=True)
 
     def republish(self, body, n_current_attempts, frame):
         self.channel.basic_publish(
@@ -100,7 +100,7 @@ class PikaQueueManager(QueueManager):
 
     def publish_response(self, message, callback, max_attempts=3):
 
-        logger.debug(f"Publishing response for {message}.")
+        logger.debug(f"Processing {message}.")
 
         frame, properties, body = message
 

scripts/manage_minio.py

@@ -37,8 +37,11 @@ def upload_compressed_response(storage, bucket_name, dossier_id, file_id, result
 
 
 def add_file_compressed(storage, bucket_name, dossier_id, path) -> None:
-
-    path_gz = combine_dossier_id_and_file_id_and_extension(dossier_id, Path(path).stem, ".ORIGIN.pdf.gz")
+    if Path(path).suffix == ".pdf":
+        suffix_gz = ".ORIGIN.pdf.gz"
+    if Path(path).suffix == ".json":
+        suffix_gz = ".TEXT.json.gz"
+    path_gz = combine_dossier_id_and_file_id_and_extension(dossier_id, Path(path).stem, suffix_gz)
 
     with open(path, "rb") as f:
         data = gzip.compress(f.read())

scripts/mock_client.py

@@ -1,3 +1,4 @@
+import argparse
 import json
 
 import pika
@@ -6,6 +7,13 @@ from pyinfra.config import CONFIG
 from pyinfra.storage.storages import get_s3_storage
 
 
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--analysis_container", "-a", choices=["detr", "ner", "image"], required=True)
+    args = parser.parse_args()
+    return args
+
+
 def read_connection_params():
     credentials = pika.PlainCredentials(CONFIG.rabbitmq.user, CONFIG.rabbitmq.password)
     parameters = pika.ConnectionParameters(
@@ -25,7 +33,7 @@ def make_channel(connection) -> pika.adapters.blocking_connection.BlockingChanne
 
 def declare_queue(channel, queue: str):
     args = {"x-dead-letter-exchange": "", "x-dead-letter-routing-key": CONFIG.rabbitmq.queues.dead_letter}
-    return channel.queue_declare(queue=queue, auto_delete=False, arguments=args)
+    return channel.queue_declare(queue=queue, auto_delete=False, durable=True, arguments=args)
 
 
 def make_connection() -> pika.BlockingConnection:
@@ -34,32 +42,42 @@ def make_connection() -> pika.BlockingConnection:
     return connection
 
 
-def build_message_bodies():
+def build_message_bodies(analyse_container_type):
+    def update_message(message_dict):
+        if analyse_container_type == "detr" or analyse_container_type == "image":
+            message_dict.update({"targetFileExtension": "ORIGIN.pdf.gz", "responseFileExtension": "IMAGE_INFO.json.gz"})
+        if analyse_container_type == "ner":
+            message_dict.update(
+                {"targetFileExtension": "TEXT.json.gz", "responseFileExtension": "NER_ENTITIES.json.gz"}
+            )
+        return message_dict
+
     storage = get_s3_storage()
     for bucket_name, pdf_name in storage.get_all_object_names(CONFIG.storage.bucket):
+        if "pdf" not in pdf_name:
+            continue
         file_id = pdf_name.split(".")[0]
         dossier_id, file_id = file_id.split("/")
-        yield json.dumps(
-            {
-                "dossierId": dossier_id,
-                "fileId": file_id,
-                "targetFileExtension": "ORIGIN.pdf.gz",
-                "responseFileExtension": "detr.json.gz",
-            }
-        ).encode()
+        message_dict = {"dossierId": dossier_id, "fileId": file_id}
+        update_message(message_dict)
+        yield json.dumps(message_dict).encode()
 
 
-if __name__ == "__main__":
-
+def main(args):
     connection = make_connection()
     channel = make_channel(connection)
     declare_queue(channel, CONFIG.rabbitmq.queues.input)
     declare_queue(channel, CONFIG.rabbitmq.queues.output)
 
-    for body in build_message_bodies():
+    for body in build_message_bodies(args.analysis_container):
         channel.basic_publish("", CONFIG.rabbitmq.queues.input, body)
         print(f"Put {body} on {CONFIG.rabbitmq.queues.input}")
 
     for method_frame, _, body in channel.consume(queue=CONFIG.rabbitmq.queues.output):
         print(f"Received {json.loads(body)}")
         channel.basic_ack(method_frame.delivery_tag)
+    channel.close()
+
+
+if __name__ == "__main__":
+    main(parse_args())

src/serve.py

@@ -32,6 +32,8 @@ def make_callback(analysis_endpoint):
         operations = message.get("operations", ["/"])
         results = map(perform_operation, operations)
         result = dict(zip(operations, results))
+        if list(result.keys()) == ["/"]:
+            result = list(result.values())[0]
         return result
 
     return callback