fix: threadpool for local task queue supports 10 workers

Author: k11kirky

README.md

@@ -199,6 +199,9 @@ redis:
 server:
   host: 0.0.0.0
   port: 8000
+
+queue: # if using local task queue
+  max_workers: 10 # default
 ```
 
 #### Using Environment Variables

agentserve/agent_server.py

@@ -33,7 +33,7 @@ def _initialize_task_queue(self):
                 return RedisTaskQueue(self.config)
             else:
                 from .queues.local_task_queue import LocalTaskQueue
-                return LocalTaskQueue()
+                return LocalTaskQueue(self.config)
         except Exception as e:
             self.logger.error(f"Failed to initialize task queue: {str(e)}")
             raise

agentserve/config.py

@@ -42,6 +42,9 @@ def _load_config(self):
                 server_config['host'] = server_host
             if server_port:
                 server_config['port'] = int(server_port)
+                
+        queue_config = config.setdefault('queue', {})
+        queue_config['max_workers'] = int(os.getenv('AGENTSERVE_QUEUE_MAX_WORKERS', queue_config.get('max_workers', 10)))
 
         return config
 

agentserve/logging_config.py

@@ -2,7 +2,12 @@
 import sys
 from typing import Optional
 
+_loggers = {}
+
 def setup_logger(name: str = "agentserve", level: Optional[str] = None) -> logging.Logger:
+    if name in _loggers:
+        return _loggers[name]
+        
     logger = logging.getLogger(name)
 
     if not logger.handlers:
@@ -12,8 +17,10 @@ def setup_logger(name: str = "agentserve", level: Optional[str] = None) -> loggi
         )
         handler.setFormatter(formatter)
         logger.addHandler(handler)
+        logger.propagate = False  # Prevent duplicate logging
 
     log_level = getattr(logging, (level or "DEBUG").upper())
     logger.setLevel(log_level)
 
+    _loggers[name] = logger
     return logger

agentserve/queues/local_task_queue.py

@@ -5,44 +5,62 @@
 from .task_queue import TaskQueue
 import threading
 from ..logging_config import setup_logger
+import concurrent.futures
 
 class LocalTaskQueue(TaskQueue):
     def __init__(self):
         self.logger = setup_logger("agentserve.queue.local")
         self.results = {}
         self.statuses = {}
-        self.loop = asyncio.new_event_loop()
+        max_workers = 10  # default
+        if config:
+            max_workers = config.get('queue', {}).get('max_workers', 10)
+        self.thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=max_workers)
+        self.lock = threading.Lock()
         self.logger.info("LocalTaskQueue initialized")
 
     def enqueue(self, agent_function, task_data: Dict[str, Any], task_id: str):
         self.logger.debug(f"Enqueueing task {task_id}")
-        self.statuses[task_id] = 'queued'
-        threading.Thread(target=self._run_task, args=(agent_function, task_data, task_id)).start()
+        with self.lock:
+            self.statuses[task_id] = 'queued'
+        self.thread_pool.submit(self._run_task, agent_function, task_data, task_id)
 
     def _run_task(self, agent_function, task_data: Dict[str, Any], task_id: str):
         self.logger.debug(f"Starting task {task_id}")
-        self.statuses[task_id] = 'in_progress'
+        with self.lock:
+            self.statuses[task_id] = 'in_progress'
+        
         try:
             if getattr(agent_function, '_is_async', False):
-                asyncio.set_event_loop(self.loop)
-                result = self.loop.run_until_complete(agent_function(task_data))
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
+                try:
+                    result = loop.run_until_complete(agent_function(task_data))
+                finally:
+                    loop.close()
             else:
                 result = agent_function(task_data)
-            self.results[task_id] = result
-            self.statuses[task_id] = 'completed'
+            
+            with self.lock:
+                self.results[task_id] = result
+                self.statuses[task_id] = 'completed'
             self.logger.info(f"Task {task_id} completed successfully")
+        
         except Exception as e:
             self.logger.error(f"Task {task_id} failed: {str(e)}")
-            self.results[task_id] = e
-            self.statuses[task_id] = 'failed'
+            with self.lock:
+                self.results[task_id] = e
+                self.statuses[task_id] = 'failed'
 
     def get_status(self, task_id: str) -> str:
-        return self.statuses.get(task_id, 'not_found')
+        with self.lock:
+            return self.statuses.get(task_id, 'not_found')
 
     def get_result(self, task_id: str) -> Any:
-        if task_id not in self.results:
-            return None
-        result = self.results[task_id]
-        if isinstance(result, Exception):
-            raise result
-        return result
+        with self.lock:
+            if task_id not in self.results:
+                return None
+            result = self.results[task_id]
+            if isinstance(result, Exception):
+                raise result
+            return result

async_example.py

@@ -12,7 +12,7 @@ class MyInputSchema(BaseModel):
 
 @app.agent(input_schema=MyInputSchema)
 async def my_agent(task_data):
-    await asyncio.sleep(1)
+    await asyncio.sleep(20)
     return task_data
 
 app.run()

example.py

@@ -2,7 +2,7 @@
 from pydantic import BaseModel
 
 # Configure logging level
-agentserve.setup_logger(level="DEBUG")  # or "INFO", "WARNING", "ERROR"
+agentserve.setup_logger(level="INFO")  # or "INFO", "WARNING", "ERROR"
 
 app = agentserve.app()