| """ |
| Async utilities for SkyDiscover |
| """ |
|
|
| import asyncio |
| import logging |
| from typing import Any, Callable, List, Optional, Sequence, Tuple |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class TaskPool: |
| """ |
| A simple task pool for managing and limiting concurrent tasks |
| """ |
|
|
| def __init__(self, max_concurrency: int = 10): |
| self.max_concurrency = max_concurrency |
| self._semaphore: Optional[asyncio.Semaphore] = None |
| self.tasks: List[asyncio.Task] = [] |
|
|
| @property |
| def semaphore(self) -> asyncio.Semaphore: |
| """Lazy-initialize the semaphore when first needed.""" |
| if self._semaphore is None: |
| self._semaphore = asyncio.Semaphore(self.max_concurrency) |
| return self._semaphore |
|
|
| async def run(self, coro: Callable, *args: Any, **kwargs: Any) -> Any: |
| """Run a single coroutine function under the concurrency semaphore.""" |
| async with self.semaphore: |
| return await coro(*args, **kwargs) |
|
|
| def create_task(self, coro: Callable, *args: Any, **kwargs: Any) -> asyncio.Task: |
| """Create, track, and return an ``asyncio.Task`` bounded by the pool.""" |
| task = asyncio.create_task(self.run(coro, *args, **kwargs)) |
| self.tasks.append(task) |
| task.add_done_callback(lambda t: self.tasks.remove(t)) |
| return task |
|
|
| async def gather( |
| self, |
| coros: Sequence[Callable], |
| args_list: Sequence[Tuple[Any, ...]] = (), |
| kwargs_list: Sequence[dict] = (), |
| return_exceptions: bool = False, |
| ) -> List[Any]: |
| """Run *coros* concurrently (bounded by the semaphore), return results in order.""" |
| n = len(coros) |
| _args = args_list if args_list else [() for _ in range(n)] |
| _kwargs = kwargs_list if kwargs_list else [{} for _ in range(n)] |
|
|
| if len(_args) != n: |
| raise ValueError(f"args_list length ({len(_args)}) must match coros length ({n})") |
| if len(_kwargs) != n: |
| raise ValueError(f"kwargs_list length ({len(_kwargs)}) must match coros length ({n})") |
|
|
| tasks = [ |
| self.create_task(coro, *args, **kwargs) |
| for coro, args, kwargs in zip(coros, _args, _kwargs) |
| ] |
| return await asyncio.gather(*tasks, return_exceptions=return_exceptions) |
|
|
