tokenintelligence/LiveCodeBench-SnapShot-0406

1""" Utilities for running functions in parallel processes. """2import sys3import resource4import multiprocessing as mp5import queue6import traceback7from enum import Enum8from typing import Callable, Optional, Dict, Any, List, Iterator9from concurrent.futures import TimeoutError10 11import attrs12import tqdm13from pebble import concurrent, ProcessPool, ProcessExpired14 15 16class FuncTimeoutError(TimeoutError):17    pass18 19 20def generate_queue() -> mp.Queue:21    """22    Generates a queue that can be shared amongst processes23    Returns:24        (multiprocessing.Queue): A queue instance25    """26    manager = mp.Manager()27    return manager.Queue()28 29 30QueueEmptyException = queue.Empty31 32 33def run_func_in_process(34    func: Callable,35    *args,36    _timeout: Optional[int] = None,37    _use_spawn: bool = True,38    **kwargs,39):40    """41    Runs the provided function in a separate process with the supplied args42    and kwargs. The args, kwargs, and43    return values must all be pickle-able.44    Args:45        func: The function to run.46        *args: Positional args, if any.47        _timeout: A timeout to use for the function.48        _use_spawn: The 'spawn' multiprocess context is used.'fork' otherwise.49        **kwargs: Keyword args, if any.50    Returns:51        The result of executing the function.52    """53    mode = "spawn" if _use_spawn else "fork"54    c_func = concurrent.process(timeout=_timeout, context=mp.get_context(mode))(func)55    future = c_func(*args, **kwargs)56 57    try:58        result = future.result()59        return result60 61    except TimeoutError:62        raise FuncTimeoutError63 64 65class TaskRunStatus(Enum):66    SUCCESS = 067    EXCEPTION = 168    TIMEOUT = 269    PROCESS_EXPIRED = 370 71 72@attrs.define(eq=False, repr=False)73class TaskResult:74    status: TaskRunStatus75 76    result: Optional[Any] = None77    exception_tb: Optional[str] = None78 79    def is_success(self) -> bool:80        return self.status == TaskRunStatus.SUCCESS81 82    def is_timeout(self) -> bool:83        return self.status == TaskRunStatus.TIMEOUT84 85    def is_exception(self) -> bool:86        return self.status == TaskRunStatus.EXCEPTION87 88    def is_process_expired(self) -> bool:89        return self.status == TaskRunStatus.PROCESS_EXPIRED90 91 92def initializer(limit):93    """Set maximum amount of memory each worker process can allocate."""94    soft, hard = resource.getrlimit(resource.RLIMIT_AS)95    resource.setrlimit(resource.RLIMIT_AS, (limit, hard))96 97 98def run_tasks_in_parallel_iter(99    func: Callable,100    tasks: List[Any],101    num_workers: int = 2,102    timeout_per_task: Optional[int] = None,103    use_progress_bar: bool = False,104    progress_bar_desc: Optional[str] = None,105    max_tasks_per_worker: Optional[int] = None,106    use_spawn: bool = True,107    max_mem: int = 1024 * 1024 * 1024 * 4,108) -> Iterator[TaskResult]:109    """110    Args:111        func: The function to run. The function must accept a single argument.112        tasks: A list of tasks i.e. arguments to func.113        num_workers: Maximum number of parallel workers.114        timeout_per_task: The timeout, in seconds, to use per task.115        use_progress_bar: Whether to use a progress bar. Default False.116        progress_bar_desc: String to display in the progress bar. Default None.117        max_tasks_per_worker: Maximum number of tasks assigned118        to a single process / worker. None means infinite.119            Use 1 to force a restart.120        use_spawn: The 'spawn' multiprocess context is used. 'fork' otherwise.121    Returns:122        A list of TaskResult objects, one per task.123    """124 125    mode = "spawn" if use_spawn else "fork"126 127    with ProcessPool(128        max_workers=num_workers,129        max_tasks=0 if max_tasks_per_worker is None else max_tasks_per_worker,130        context=mp.get_context(mode),131    ) as pool:132        future = pool.map(func, tasks, timeout=timeout_per_task)133 134        iterator = future.result()135        if use_progress_bar:136            pbar = tqdm.tqdm(137                desc=progress_bar_desc,138                total=len(tasks),139                dynamic_ncols=True,140                file=sys.stdout,141            )142        else:143            pbar = None144 145        succ = timeouts = exceptions = expirations = 0146 147        while True:148            try:149                result = next(iterator)150 151            except StopIteration:152                break153 154            except TimeoutError as error:155                yield TaskResult(156                    status=TaskRunStatus.TIMEOUT,157                )158 159                timeouts += 1160 161            except ProcessExpired as error:162                yield TaskResult(163                    status=TaskRunStatus.PROCESS_EXPIRED,164                )165                expirations += 1166 167            except Exception as error:168                exception_tb = traceback.format_exc()169 170                yield TaskResult(171                    status=TaskRunStatus.EXCEPTION,172                    exception_tb=exception_tb,173                )174                exceptions += 1175 176            else:177                yield TaskResult(178                    status=TaskRunStatus.SUCCESS,179                    result=result,180                )181 182                succ += 1183 184            if pbar is not None:185                pbar.update(1)186                pbar.set_postfix(187                    succ=succ, timeouts=timeouts, exc=exceptions, p_exp=expirations188                )189                sys.stdout.flush()190                sys.stderr.flush()191 192 193def run_tasks_in_parallel(194    func: Callable,195    tasks: List[Any],196    num_workers: int = 2,197    timeout_per_task: Optional[int] = None,198    use_progress_bar: bool = False,199    progress_bar_desc: Optional[str] = None,200    max_tasks_per_worker: Optional[int] = None,201    use_spawn: bool = True,202) -> List[TaskResult]:203    """204    Args:205        func: The function to run. The function must accept a single argument.206        tasks: A list of tasks i.e. arguments to func.207        num_workers: Maximum number of parallel workers.208        timeout_per_task: The timeout, in seconds, to use per task.209        use_progress_bar: Whether to use a progress bar. Defaults False.210        progress_bar_desc: String to display in the progress bar. Default None.211        max_tasks_per_worker: Maximum number of tasks assigned to a single212        process / worker. None means infinite.213            Use 1 to force a restart.214        use_spawn: The 'spawn' multiprocess context is used. 'fork' otherwise.215    Returns:216        A list of TaskResult objects, one per task.217    """218 219    task_results: List[TaskResult] = list(220        run_tasks_in_parallel_iter(221            func=func,222            tasks=tasks,223            num_workers=num_workers,224            timeout_per_task=timeout_per_task,225            use_progress_bar=use_progress_bar,226            progress_bar_desc=progress_bar_desc,227            max_tasks_per_worker=max_tasks_per_worker,228            use_spawn=use_spawn,229        )230    )231 232    return task_results233