# AUTOGENERATED - modify shared_anyscale_util in root directory to make changes
import asyncio
import logging
from typing import Awaitable, List, Optional, TypeVar, Union


_T = TypeVar("_T")
logger = logging.getLogger(__file__)


async def gather_in_batches(
    n: int,
    *tasks: Awaitable[_T],
    per_task_timeout_secs: Optional[float] = None,
    return_exceptions: bool = False,
    log_exceptions: bool = False,
) -> List[Union[_T, BaseException]]:
    # n cannot be less than 1
    n = max(n, 1)
    semaphore = asyncio.Semaphore(n)

    async def sem_task(task: Awaitable[_T]) -> _T:
        try:
            async with semaphore:
                return await asyncio.wait_for(task, per_task_timeout_secs)
        except Exception:
            if log_exceptions:
                logger.exception("Error while gathering task")
            raise  # re-raise exception

    return await asyncio.gather(
        *(sem_task(task) for task in tasks), return_exceptions=return_exceptions
    )
