Merge pull request #112 from ISSResearch/temp-main

f

Author: githubering182

storage-app/src/shared/archive_helpers.py

@@ -1,4 +1,4 @@
-from asyncio import create_task, wait, FIRST_COMPLETED, gather
+from asyncio import create_task, wait, FIRST_COMPLETED
 from typing import Optional
 from threading import Thread
 from queue import Queue
@@ -257,7 +257,7 @@ async def produce(self):
             self.iter_count += 1
             if not self.iter_count % GC_FREQ: gc_collect()
 
-        await gather(*tasks)
+        for task in tasks: await task
         self.queue.put(None)
 
         self._done = True

storage-app/src/shared/worker_services.py

@@ -6,7 +6,8 @@
     APP_BACKEND_URL,
     ASYNC_PRODUCER_MAX_CONCURRENT as MAX_CONCURENT
 )
-from asyncio import get_event_loop, sleep as async_stall_for, create_task
+from asyncio import get_event_loop
+from time import sleep as stall_for
 from shared.app_services import Bucket
 from shared.utils import emit_token
 from shared.embedding_db import EmbeddingStorage
@@ -57,12 +58,13 @@ 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()
 
         wait_item = lambda t, n: type("wi", (object,), {"task": t, "next": n})
-        wait_list = wait_item(producer, wait_item(consumer, wait_item(writer, None)))
+        wait_list = wait_item(consumer, wait_item(writer, None))
+
+        await producer.produce()
 
         while wait_list:
             if wait_list.task.ready:
@@ -71,9 +73,8 @@ async def archive_objects(self) -> Optional[bool]:
 
             print(f"ZIP WORK STALL, {producer.iter_count}")
             self._task.update_state(state="PROGRESS")
-            await async_stall_for(5)
+            stall_for(5)
 
-        await producer_task
         await object_set.close()
         consumer.join()
         writer.join()

storage-app/src/worker.py

@@ -18,6 +18,9 @@ def default(self, o) -> Any: return getattr(o, "__json__", super().default)(o)
 worker.conf.broker_url = BROKER_URL
 worker.conf.result_backend = RESULT_URL
 worker.conf.broker_transport_options = {"visibility_timeout": 1 * 60 * 60 * 24}
+worker.conf.task_time_limit = None
+worker.conf.task_soft_time_limit = None
+worker.conf.worker_lost_wait = None
 
 register(
     "custom_encoder",
@@ -29,10 +32,22 @@ def default(self, o) -> Any: return getattr(o, "__json__", super().default)(o)
 
 @worker.task(
     bind=True,
-    name="produce_download_task",
+    name="test",
     time_limit=None,
     soft_time_limit=None
 )
+def test(self):
+    from subprocess import PIPE, Popen
+    from time import sleep
+
+    pc = Popen(["sleep", "40"], stdout=PIPE)
+    while not pc.poll() == 0:
+        sleep(5)
+        print("tick")
+    return pc.stdout.read()
+
+
+@worker.task(bind=True, name="produce_download_task",)
 def produce_download_task(self, bucket_name: str, file_ids: list[str]) -> str | None:
     task = Zipper(bucket_name, file_ids, self)
     run(task.archive_objects())