download_dataset command

new_experiment/queue_base.py

@@ -23,7 +23,7 @@ def ack_message(channel, delivery_tag):
         pass
 
 
-def process_queue(queue_name: str, process_message: Callable[[bytes], None]):
+def process_queue(queue_name: str, process_message: Callable[[bytes], None], prefetch_count: int = 1):
     def do_work(connection, channel, delivery_tag, body):
         process_message(body)
         cb = functools.partial(ack_message, channel, delivery_tag)
@@ -40,7 +40,7 @@ def process_queue(queue_name: str, process_message: Callable[[bytes], None]):
     parameters = pika.URLParameters(_RABBIT_URL)
     connection = pika.BlockingConnection(parameters)
     channel = connection.channel()
-    channel.basic_qos(prefetch_count=1)
+    channel.basic_qos(prefetch_count=prefetch_count)
 
     threads = []
     on_message_callback = functools.partial(on_message, args=(connection, threads))

new_experiment/worker_pipeline.py

+import argparse
 import json
 
 from new_experiment.pipeline.pipeline_process_flair_upos import run_flair_upos_multi_pipeline
@@ -37,4 +38,8 @@ def process_message(body: bytes):
 
 
 if __name__ == '__main__':
-    process_queue('asr_benchmark_experiments', process_message)
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--prefetch_count")
+    args = parser.parse_args()
+    process_queue('asr_benchmark_experiments', process_message,
+                  1 if args.prefetch_count in [None, ''] else int(args.prefetch_count))