made adjustments for image classification with pyinfra 2.x; added related fixmes

config.yaml

@@ -4,6 +4,8 @@ service:
   response_formatter: default  # formats analysis payloads of response messages
   upload_formatter: projecting  # formats analysis payloads of objects uploaded to storage
   # Note: This is not really the right place for this. It should be configured on a per-service basis.
+  operation: default
+#  operation: image_classification  # FIXME: operation needs to be specified in deployment config for services that are called without an operation specifiec
   operations:
     conversion:
       input:
@@ -26,6 +28,13 @@ service:
       output:
         subdir: "table_parses"
         extension: json.gz
+    image_classification:
+      input:
+        subdir: "extracted_images"
+        extension: json.gz
+      output:
+        subdir: ""
+        extension: IMAGE_INFO.json.gz
     default:
       input:
         subdir: ""

pyinfra/file_descriptor_manager.py

@@ -3,6 +3,8 @@ from _operator import itemgetter
 
 from funcy import project
 
+from pyinfra.config import CONFIG
+
 
 class FileDescriptorManager:
     def __init__(self, bucket_name, operation2file_patterns: dict = None):
@@ -28,8 +30,7 @@ class FileDescriptorManager:
 
         return object_name
 
-    @staticmethod
-    def __build_matcher(file_descriptor):
+    def __build_matcher(self, file_descriptor):
         def make_filename(file_id, subdir, suffix):
             return os.path.join(file_id, subdir, suffix) if subdir else f"{file_id}.{suffix}"
 
@@ -37,19 +38,28 @@ class FileDescriptorManager:
             "dossierId", "fileId", "subdir", "pages", "extension"
         )(file_descriptor)
 
+        matcher = os.path.join(
+            dossier_id, make_filename(file_id, subdir, self.__build_page_regex(pages, subdir) + extension)
+        )
+
+        return matcher
+
+    @staticmethod
+    def __build_page_regex(pages, subdir):
         if len(pages) > 1 and subdir:
             page_re = "id:(" + "|".join(map(str, pages)) + ")."
+        elif not pages and subdir:
+            page_re = r"id:\d+."
         elif len(pages) == 1 and subdir:
             page_re = f"id:{pages[0]}."
         else:
             page_re = ""
 
-        matcher = os.path.join(dossier_id, make_filename(file_id, subdir, page_re + extension))
-
-        return matcher
+        return page_re
 
     def build_file_descriptor(self, queue_item_body, end="input"):
-        operation = queue_item_body.get("operation", "default")
+        # FIXME: reference to config is misplaced here. Needs to be abstracted in some way.
+        operation = queue_item_body.get("operation", CONFIG.service.operation)
 
         file_pattern = self.operation2file_patterns[operation][end]
 

pyinfra/server/utils.py

@@ -12,4 +12,5 @@ def make_streamable_and_wrap_in_packing_logic(fn, batched):
 
 
 def make_streamable(fn, batched):
+    # FIXME: something broken with batched == True
     return compose(normalize, (identity if batched else starlift)(fn))

pyinfra/visitor/downloader.py

@@ -30,11 +30,12 @@ class Downloader:
         logger.debug(f"Filtering objects in bucket {self.bucket_name} by pattern...")
 
         names_of_relevant_objects = compose(
+            list,
             self.get_pattern_filter(queue_item_body),
             self.get_names_of_all_associated_objects,
         )(queue_item_body)
 
-        logger.debug(f"Completed filtering.")
+        logger.debug(f"Found {len(names_of_relevant_objects)} objects matching filter.")
 
         return names_of_relevant_objects
 
@@ -51,5 +52,6 @@ class Downloader:
 
     def get_pattern_filter(self, queue_item_body):
         file_pattern = self.file_descriptor_manager.build_input_matcher(queue_item_body)
+        logger.debug(f"Filtering pattern: {file_pattern if len(file_pattern) <= 120 else (file_pattern[:120]+'...')}")
         matches_pattern = flift(file_pattern)
         return matches_pattern

scripts/mock_client.py

@@ -9,7 +9,7 @@ from pyinfra.storage.storages import get_s3_storage
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--bucket_name", "-b", required=True)
+    parser.add_argument("--bucket_name", "-b")
     parser.add_argument(
         "--analysis_container",
         "-a",
@@ -51,7 +51,7 @@ def make_connection() -> pika.BlockingConnection:
 def build_message_bodies(analysis_container, bucket_name):
     def update_message(message_dict):
         if analysis_container == "detr" or analysis_container == "image":
-            message_dict.update({"targetFileExtension": "ORIGIN.pdf.gz", "responseFileExtension": "IMAGE_INFO.json.gz"})
+            message_dict.update({"pages": []})
         if analysis_container == "conversion":
             message_dict.update(
                 {
@@ -97,7 +97,9 @@ def main(args):
     declare_queue(channel, CONFIG.rabbitmq.queues.input)
     declare_queue(channel, CONFIG.rabbitmq.queues.output)
 
-    for body in build_message_bodies(args.analysis_container, args.bucket_name):
+    bucket_name = args.bucket_name or parse_disjunction_string(CONFIG.storage.bucket)
+
+    for body in build_message_bodies(args.analysis_container, bucket_name):
         channel.basic_publish("", CONFIG.rabbitmq.queues.input, body)
         print(f"Put {body} on {CONFIG.rabbitmq.queues.input}")
 

scripts/show_compressed_json.py

@@ -2,25 +2,23 @@ import argparse
 import gzip
 import json
 
-from pyinfra.server.packing import bytes_to_string
-
 
 def parse_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--compressed_json", "-j", required=True)
+    parser.add_argument("compressed_json_path", help="Path to compressed JSON file")
     return parser.parse_args()
 
 
 def main(fp):
     with open(fp, "rb") as f:
-        compressed_json = f.read()
+        compressed_json_path = f.read()
 
-    json_str = gzip.decompress(compressed_json)
+    json_str = gzip.decompress(compressed_json_path)
     json_dict = json.loads(json_str)
 
     print(json.dumps(json_dict, indent=2))
 
 
 if __name__ == "__main__":
-    fp = parse_args().compressed_json
-    main(fp)
+    args = parse_args()
+    main(args.compressed_json_path)