tmp

Author: Nick Tsai

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "worker_dispatcher"
-version = "0.0.1"
+version = "0.0.2"
 authors = [
   { name="Nick Tsai", email="myintaer@gmail.com" },
 ]

src/worker_dispatcher/worker_dispatcher.py

@@ -1,10 +1,7 @@
-import copy
-import queue
-import threading
-import multiprocessing
-import concurrent.futures
-import time
+import multiprocessing, concurrent.futures
+import time, copy
 
+# Sample callback function
 def callback_sample(config, task_id: int, task_data):
     if task_id == 1:
         print("Runing sample task function, please customize yours according to the actual usage.")
@@ -13,11 +10,13 @@ def callback_sample(config, task_id: int, task_data):
     }
     return result
 
+# Sample result callback function
 def result_callback_sample(config, task_id: int, task_result, log):
     if task_id == 1:
         print("Runing sample result function, please customize yours according to the actual usage.")
     return task_result
 
+# Default configuration
 default_config = {
     'debug': False,
     'task': {
@@ -27,22 +26,16 @@ def result_callback_sample(config, task_id: int, task_result, log):
         'result_callback': False
     },
     'worker': {
-        'number': 1,
+        'number': multiprocessing.cpu_count(),
         'per_second': 0,     
-        'multiprocessing': False,
-    },
-    'output': {
-        'enabled': False,
-        'format': 'csv',
-        'filepath': './',
-        'filename': 'output.csv',
-        'header': [],
+        'multiprocessing': False
     }
 }
 
 results = []
 logs = []
 
+# Start function
 def start(userConfig: dict):
 
     global results, logs
@@ -64,10 +57,10 @@ def start(userConfig: dict):
         in_child_process = (multiprocessing.current_process().name != 'MainProcess')
         # Return False if is in worker process to let caller handle
         if in_child_process:
-            print("Exit procedure due to the child process")
+            # print("Exit procedure due to the child process")
             return False
 
-
+    # Debug mode
     if config['debug']:
         print("Configuration:")
         print(config)
@@ -77,7 +70,6 @@ def start(userConfig: dict):
         exit("Callback function is invalied")
 
     # Task list to queue
-    task_queue = multiprocessing.Queue() if use_multiprocessing else queue.Queue()
     task_list = []
     user_task_list = config['task']['list']
     if isinstance(user_task_list, list):
@@ -101,19 +93,22 @@ def start(userConfig: dict):
     # Worker dispatch
     worker_num = config['worker']['number']
     worker_num = worker_num if isinstance(worker_num, int) else 1
+    worker_per_second = config['worker']['per_second'] if config['worker']['per_second'] else 0
+    max_workers = len(task_list) if worker_per_second else worker_num
     # result_queue = multiprocessing.Queue() if use_multiprocessing else queue.Queue()
     pool_executor_class = concurrent.futures.ProcessPoolExecutor if use_multiprocessing else concurrent.futures.ThreadPoolExecutor
     print("Start to dispatch workers ---\n")
 
     # Pool Executor
-    with pool_executor_class(max_workers=worker_num) as executor:
+    with pool_executor_class(max_workers=max_workers) as executor:
         pool_results = []
+        # Task dispatch
         for task in task_list:
             pool_result = executor.submit(consume_task, task, config)
             pool_results.append(pool_result)
             # Worker per_second setting
-            if config['worker']['per_second'] and task['id'] % worker_num == 0:
-                time.sleep(float(config['worker']['per_second']))
+            if worker_per_second and task['id'] % worker_num == 0:
+                time.sleep(float(worker_per_second))
         # Get results from the async results
         for pool_result in concurrent.futures.as_completed(pool_results):
             log = pool_result.result()
@@ -127,20 +122,23 @@ def start(userConfig: dict):
     print("End of worker dispatch ---\n")
     return results
 
-def get_results():
-    return results
-
-def get_logs():
-    return logs
-
+# Worker function
 def consume_task(data, config):
     started_at = time.time()
     return_value = config['task']['callback'](config['task']['config'], data['id'], data['task'])
     ended_at = time.time()
+    duration = ended_at - started_at
     log = {
         'task_id': data['id'],
         'started_at': started_at,
         'ended_at': ended_at,
+        'duration': duration,
         'result': return_value
     }
-    return log
+    return log
+
+def get_results():
+    return results
+
+def get_logs():
+    return logs