filtering with distributask

Author: antbaez9

.DS_Store

@@ -1,38 +1,37 @@
-
-from distributaur.distributaur import create_from_config
+from distributask.distributask import create_from_config
 from .filter import read_json_in_batches
 from .worker import run_job
 from tqdm import tqdm
 import time
 
-
 if __name__ == "__main__":
 
-    input_filename = 'datasets/cap3d_captions.json'
-    batch_size = 1000
+    input_filename = "datasets/cap3d_captions.json"
+    batch_size = 10000
 
-    distributaur = create_from_config()
+    distributask = create_from_config()
 
-    max_price = 0.1
-    max_nodes = 50
+    max_price = 0.25
+    max_nodes = 25
     docker_image = "antbaez/filter-worker:latest"
     module_name = "filtered.worker"
 
-    redis_client = distributaur.get_redis_connection()
+    redis_client = distributask.get_redis_connection()
 
-    rented_nodes = distributaur.rent_nodes(max_price, max_nodes, docker_image, module_name)
+    rented_nodes = distributask.rent_nodes(
+        max_price, max_nodes, docker_image, module_name
+    )
     print("Total nodes rented: ", len(rented_nodes))
 
-    distributaur.register_function(run_job)
+    distributask.register_function(run_job)
 
     while True:
         user_input = input("press r when workers are ready: ")
         if user_input == "r":
             break
 
-
     total_batches = 0
-    
+
     print("Sending tasks")
     tasks = []
 
@@ -43,32 +42,44 @@
     for i in range(num_batches):
 
         batch = json_batches[i]
-
         total_batches += 1
-        task = distributaur.execute_function("run_job", {
-            "batch_index" : total_batches,
-            "batch" : batch
-        })
+
+        print(total_batches)
+        task = distributask.execute_function(
+            "run_job", {"batch_index": total_batches, "batch": batch}
+        )
 
         tasks.append(task)
 
     first_task_done = False
     print("Tasks sent. Starting monitoring")
+
+    inactivity_log = {node["instance_id"]: 0 for node in rented_nodes}
+
+    start_time = time.time()
     with tqdm(total=len(tasks), unit="task") as pbar:
         while not all(task.ready() for task in tasks):
+
             current_tasks = sum([task.ready() for task in tasks])
             pbar.update(current_tasks - pbar.n)
-            if current_tasks > 0:
-                if not first_task_done:
-                    first_task_done = True
-                    first_task_start_time = time.time()
-
-                end_time = time.time()
-                elapsed_time = end_time - first_task_start_time
-                time_per_tasks = elapsed_time / current_tasks
-                time_left = time_per_tasks * (len(tasks) - current_tasks)
-
-                pbar.set_postfix(
-                    elapsed=f"{elapsed_time:.2f}s", time_left=f"{time_left:.2f}"
-                )
-            time.sleep(2)
+
+            time.sleep(1)
+
+            current_time = time.time()
+            if current_time - start_time > 60:
+                start_time = time.time()
+
+                for node in rented_nodes:
+                    log_response = distributask.get_node_log(node)
+                    if log_response.status_code == 200:
+                        try:
+                            last_msg = log_response.text.splitlines()[-1]
+                            if ("Task complete" in last_msg and inactivity_log[node["instance_id"]] == 0):
+                                inactivity_log[node["instance_id"]] = 1
+                            elif ("Task complete" in last_msg and inactivity_log[node["instance_id"]] == 1):
+                                distributask.terminate_nodes([node])
+                                print("node terminated")
+                            else:
+                                inactivity_log[node["instance_id"]] == 0
+                        except:
+                            pass

filtered/__pycache__/distributed.cpython-310.pyc

@@ -1,23 +1,27 @@
 import sys
 from .filter import filter_captions, write_filtered_json
 
-
 def run_job(batch_index, batch):
 
-    output_filename = f"batch_{batch_index}"
+    if len(str(batch_index)) == 1:
+        batch_num = f"0{batch_index}"
+    else:
+        batch_num = f"{batch_index}"
+
+    output_filename = f"batch_{batch_num}"
 
     filtered_batch = filter_captions(batch)
     write_filtered_json(output_filename, filtered_batch)
 
-    distributaur.upload_file(output_filename)
+    distributask.upload_file(output_filename)
 
     return "Task complete"
 
 
 if __name__ == "__main__" or any("celery" in arg for arg in sys.argv):
-    from distributaur.distributaur import create_from_config
+    from distributask.distributask import create_from_config
 
-    distributaur = create_from_config()
-    distributaur.register_function(run_job)
+    distributask = create_from_config()
+    distributask.register_function(run_job)
 
-    celery = distributaur.app
+    celery = distributask.app

filtered/__pycache__/worker.cpython-310.pyc

@@ -7,4 +7,4 @@ python-dotenv
 omegaconf
 tqdm
 gliner
-distributaur
+distributask