Parallel map, but with backpressure. If the caller doesn't call `next` fast enough, this will stop calling `func` at some point rather than letting results pile up in memory. Specifically, there is a max of one output value buffered per thread.
(func: Callable[[In], Out], iterable: Iterable[In], concurrency: int, max_workers: int | None = None, use_processpool_executor: bool = False)
| 1074 | |
| 1075 | |
| 1076 | def bounded_parallel_map(func: Callable[[In], Out], iterable: Iterable[In], concurrency: int, max_workers: int | None = None, use_processpool_executor: bool = False) -> Iterable[Out]: |
| 1077 | '''Parallel map, but with backpressure. If the caller doesn't call `next` |
| 1078 | fast enough, this will stop calling `func` at some point rather than |
| 1079 | letting results pile up in memory. Specifically, there is a max of one |
| 1080 | output value buffered per thread.''' |
| 1081 | if concurrency < 2: |
| 1082 | yield from map(func, iterable) |
| 1083 | # Not reached. |
| 1084 | iterable = iter(iterable) |
| 1085 | executor_class: type[ThreadPoolExecutor] | type[ProcessPoolExecutor] |
| 1086 | if use_processpool_executor: |
| 1087 | executor_class = ProcessPoolExecutor |
| 1088 | else: |
| 1089 | executor_class = ThreadPoolExecutor |
| 1090 | with executor_class(max_workers=max_workers) as executor: |
| 1091 | futures: list[concurrent.futures.Future[Out]] = [] |
| 1092 | done = False |
| 1093 | for _ in range(concurrency): |
| 1094 | try: |
| 1095 | futures.append(executor.submit(func, next(iterable))) |
| 1096 | except StopIteration: |
| 1097 | done = True |
| 1098 | break |
| 1099 | |
| 1100 | while futures: |
| 1101 | result = futures.pop(0).result() |
| 1102 | while not done and len(futures) < concurrency: |
| 1103 | try: |
| 1104 | futures.append(executor.submit(func, next(iterable))) |
| 1105 | except StopIteration: |
| 1106 | done = True |
| 1107 | break |
| 1108 | yield result |
| 1109 | |
| 1110 | |
| 1111 | def check_vocab_size(params: Params, vocab: BaseVocab, pad_vocab: bool = False) -> None: |