135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
import asyncio
|
|
import logging
|
|
from typing import Callable, Any
|
|
import psutil # For CPU and memory monitoring
|
|
|
|
# Set up the logger
|
|
logger = logging.getLogger("AsyncTaskQueue")
|
|
logger.setLevel(logging.INFO)
|
|
handler = logging.StreamHandler()
|
|
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
|
handler.setFormatter(formatter)
|
|
logger.addHandler(handler)
|
|
|
|
class AsyncTaskQueue:
|
|
def __init__(
|
|
self,
|
|
min_workers: int = 1,
|
|
max_workers: int = 10,
|
|
cpu_threshold: float = 90.0,
|
|
memory_threshold: float = 90.0,
|
|
):
|
|
self.queue = asyncio.Queue()
|
|
self.min_workers = min_workers
|
|
self.max_workers = max_workers
|
|
self.cpu_threshold = cpu_threshold
|
|
self.memory_threshold = memory_threshold
|
|
self.workers = []
|
|
self._stop_event = asyncio.Event()
|
|
self._scale_lock = asyncio.Lock() # Prevent race conditions
|
|
|
|
async def _scale_up(self):
|
|
"""Add a worker if under max limit and resources allow."""
|
|
async with self._scale_lock:
|
|
if (
|
|
len(self.workers) < self.max_workers
|
|
and psutil.cpu_percent() < self.cpu_threshold
|
|
and psutil.virtual_memory().percent < self.memory_threshold
|
|
):
|
|
worker = asyncio.create_task(self.worker_loop(f"Worker-{len(self.workers) + 1}"))
|
|
self.workers.append(worker)
|
|
logger.info(f"Scaled up. Workers: {len(self.workers)}")
|
|
|
|
async def _scale_down(self):
|
|
"""Remove a worker if above min limit."""
|
|
async with self._scale_lock:
|
|
if len(self.workers) > self.min_workers:
|
|
worker = self.workers.pop()
|
|
worker.cancel()
|
|
logger.info(f"Scaled down. Workers: {len(self.workers)}")
|
|
|
|
async def _monitor_resources(self):
|
|
"""Monitor CPU and memory usage, and adjust workers."""
|
|
while not self._stop_event.is_set():
|
|
cpu_usage = psutil.cpu_percent()
|
|
memory_usage = psutil.virtual_memory().percent
|
|
|
|
if (
|
|
cpu_usage > self.cpu_threshold
|
|
or memory_usage > self.memory_threshold
|
|
):
|
|
await self._scale_down()
|
|
elif (
|
|
len(self.workers) < self.max_workers
|
|
and self.queue.qsize() > 2 # Only scale up if there's work
|
|
):
|
|
await self._scale_up()
|
|
|
|
await asyncio.sleep(2) # Check every 2 seconds
|
|
|
|
async def start_workers(self):
|
|
"""Start initial workers and the resource monitor."""
|
|
for i in range(self.min_workers):
|
|
worker = asyncio.create_task(self.worker_loop(f"Worker-{i+1}"))
|
|
self.workers.append(worker)
|
|
# Start the resource monitor
|
|
asyncio.create_task(self._monitor_resources())
|
|
logger.info(f"Started {self.min_workers} workers and resource monitor.")
|
|
|
|
async def stop_workers(self):
|
|
"""Stop all workers and the resource monitor."""
|
|
logger.info("Stopping workers...")
|
|
self._stop_event.set()
|
|
await self.queue.join() # Wait for all tasks to complete
|
|
for worker in self.workers:
|
|
worker.cancel()
|
|
await asyncio.gather(*self.workers, return_exceptions=True)
|
|
logger.info("All workers stopped.")
|
|
|
|
async def worker_loop(self, name: str):
|
|
"""Worker loop: Process tasks from the queue."""
|
|
logger.info(f"{name} started.")
|
|
while not self._stop_event.is_set():
|
|
try:
|
|
task = await self.queue.get()
|
|
logger.info(f"{name} processing: {task['name']}")
|
|
await task['func'](*task['args'])
|
|
except Exception as e:
|
|
logger.error(f"Error in {name}: {e}", exc_info=True)
|
|
finally:
|
|
self.queue.task_done()
|
|
logger.info(f"{name} exited.")
|
|
|
|
async def enqueue(self, name: str, func: Callable, *args: Any):
|
|
"""Add a task to the queue."""
|
|
logger.info(f"Enqueuing task: {name}")
|
|
await self.queue.put({'name': name, 'func': func, 'args': args})
|
|
|
|
async def run_sync_task_in_thread(func: Callable, *args: Any):
|
|
"""Run a synchronous function in a separate thread."""
|
|
await asyncio.to_thread(func, *args)
|
|
"""import asyncio
|
|
|
|
async def example_async_task(name: str, duration: int):
|
|
print(f"{name} started, will sleep for {duration} seconds")
|
|
await asyncio.sleep(duration)
|
|
print(f"{name} finished")
|
|
|
|
async def main():
|
|
queue = AsyncTaskQueue(
|
|
min_workers=2,
|
|
max_workers=10,
|
|
cpu_threshold=90.0,
|
|
memory_threshold=90.0,
|
|
)
|
|
await queue.start_workers()
|
|
|
|
# Enqueue tasks
|
|
for i in range(20):
|
|
await queue.enqueue(f"Task{i}", example_async_task, f"Task{i}", 1)
|
|
|
|
await asyncio.sleep(10) # Let tasks run
|
|
await queue.stop_workers()
|
|
|
|
asyncio.run(main())
|
|
""" |