ref archive task to mimic work load

Author: githubering182

storage-app/src/shared/archive_helpers.py

@@ -36,13 +36,17 @@ def __init__(self, file_name: str, as_file: bool = False):
         self.queue = Queue()
         self._archive_id = None
         self._as_file = as_file
+        self._done = False
 
     def stop(self): self.queue.put(None)
 
     def write(self, task: tuple[BytesIO, int]): self.queue.put(task)
 
     def result(self) -> Optional[str]: return self._archive_id
 
+    @property
+    def ready(self) -> bool: return self._done
+
     def _get_write_in(self):
         return (
             open(self.file_name, "wb")
@@ -75,6 +79,7 @@ def run(self):
             self._archive_id = dest._id
             SyncDataBase.close_connection()
 
+        self._done = True
 
 class ZipConsumer(Thread):
     DUMP_THRESHOLD: int = 10 << 20
@@ -91,9 +96,13 @@ def __init__(
         self.file_list = []
         self.writer = writer
         self._local_dir_end = 0
+        self._done = False
 
     def tell(self) -> int: return self._local_dir_end
 
+    @property
+    def ready(self) -> bool: return self._done
+
     def _dump_buffer(self, buffer: BytesIO, zip_buffer: ZipFile):
         dest_offset = self.tell()
 
@@ -214,6 +223,8 @@ def run(self):
 
         self._finalize()
 
+        self._done = True
+
 
 class FileProducer:
     def __init__(
@@ -225,6 +236,10 @@ def __init__(
         self.object_set = object_set
         self.queue = queue
         self.max_concurrent = max_concurrent
+        self._done = False
+
+    @property
+    def ready(self) -> bool: return self._done
 
     async def produce(self):
         tasks = []
@@ -244,6 +259,8 @@ async def produce(self):
         await gather(*tasks)
         self.queue.put(None)
 
+        self._done = True
+
     async def next(self, file: GridOut):
         try: self.queue.put((self._get_file_name(file), await file.read()))
         except Exception as e: print(f"next err {str(e)}")

storage-app/src/shared/worker_services.py

@@ -4,16 +4,14 @@
     SECRET_KEY,
     SECRET_ALGO,
     APP_BACKEND_URL,
-    TEMP_ZIP,
     ASYNC_PRODUCER_MAX_CONCURRENT as MAX_CONCURENT
 )
-from asyncio import get_event_loop
+from asyncio import get_event_loop, sleep as async_stall_for, create_task
 from shared.app_services import Bucket
 from shared.utils import emit_token
 from shared.embedding_db import EmbeddingStorage
 from typing import Any, Optional
 import requests
-from os import mkdir, path
 from .hasher import VHash, IHash
 from queue import Queue
 from .archive_helpers import FileProducer, ZipConsumer, ZipWriter
@@ -44,8 +42,6 @@ def __init__(self, bucket_name: str, file_ids: list[str]) -> None:
         self.bucket_name = bucket_name
 
     async def archive_objects(self) -> Optional[bool]:
-        if not path.exists(TEMP_ZIP): mkdir(TEMP_ZIP)
-
         json_data: Any = ("annotation.json", dumps(self.annotation, indent=4).encode("utf-8"))
 
         queue = Queue()
@@ -54,12 +50,21 @@ async def archive_objects(self) -> Optional[bool]:
         writer = ZipWriter(f"{self.bucket_name}_dataset")
         consumer = ZipConsumer(queue, [json_data], writer)
 
+        producer_task = create_task(producer.produce())
         consumer.start()
         writer.start()
 
-        await producer.produce()
-        await self.object_set.close()
+        wait_item = lambda t, n: type("wi", (object,), {"task": t, "next": n})
+        wait_list = wait_item(producer, wait_item(consumer, wait_item(writer, None)))
 
+        while wait_list:
+            if wait_list.task.ready:
+                wait_list = wait_list.next
+                continue
+            await async_stall_for(1)
+
+        await producer_task
+        await self.object_set.close()
         consumer.join()
         writer.join()