diff --git a/banner.txt b/banner.txt index e7851e0..fa0940f 100644 --- a/banner.txt +++ b/banner.txt @@ -1,8 +1,6 @@ -______ _____ __ -| ___ \ |_ _| / _| -| |_/ / _ | | _ __ | |_ _ __ __ _ -| __/ | | || || '_ \| _| '__/ _` | -| | | |_| || || | | | | | | | (_| | -\_| \__, \___/_| |_|_| |_| \__,_| - __/ | - |___/ \ No newline at end of file + ___ _ _ ___ __ + o O O | _ \ | || | |_ _| _ _ / _| _ _ __ _ + o | _/ \_, | | | | ' \ | _| | '_| / _` | + TS__[O] _|_|_ _|__/ |___| |_||_| _|_|_ _|_|_ \__,_| + {======|_| ``` |_| ````|_|`````|_|`````|_|`````|_|`````|_|`````| +./o--000' `-0-0-' `-0-0-' `-0-0-' `-0-0-' `-0-0-' `-0-0-' `-0-0-' \ No newline at end of file diff --git a/pyinfra/flask.py b/pyinfra/flask.py index 35e3d3a..549ef0c 100644 --- a/pyinfra/flask.py +++ b/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 diff --git a/pyinfra/queue/queue_manager/pika_queue_manager.py b/pyinfra/queue/queue_manager/pika_queue_manager.py index 984b29f..4346658 100644 --- a/pyinfra/queue/queue_manager/pika_queue_manager.py +++ b/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 diff --git a/scripts/manage_minio.py b/scripts/manage_minio.py index 3a72632..eaa921e 100644 --- a/scripts/manage_minio.py +++ b/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()) diff --git a/scripts/mock_client.py b/scripts/mock_client.py index 3590d3e..e7699fa 100644 --- a/scripts/mock_client.py +++ b/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()) diff --git a/src/serve.py b/src/serve.py index 850933f..236d42f 100644 --- a/src/serve.py +++ b/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