Add option to limit the number of mp workers

Due to python/cpython#89240, any use of mp.Pool with a number of workers greater than 60 fails. This means that by using cpu_count(), any system with more than 60 logical cores will crash when attempting to run.

Solve this by adding a flag to allow limiting the number of workers for users with systems with that many cores

Author: ThadHouse

gersemi/__main__.py

@@ -1,5 +1,6 @@
 import argparse
 from dataclasses import fields
+import multiprocessing as mp
 import pathlib
 import sys
 from lark import __version__ as lark_version
@@ -144,6 +145,15 @@ def create_argparser():
             """,
     )
 
+    parser.add_argument(
+        "-w",
+        "--workers",
+        dest="workers",
+        type=int,
+        default=mp.cpu_count(),
+        help = "number of workers to run (default is number of cores)"
+    )
+
     return parser
 
 
@@ -170,8 +180,9 @@ def main():
 
     configuration = make_configuration(args)
     mode = get_mode(args)
+    num_workers = args.workers
 
-    sys.exit(run(mode, configuration, args.sources))
+    sys.exit(run(mode, configuration, num_workers, args.sources))
 
 
 if __name__ == "__main__":

gersemi/runner.py

@@ -137,10 +137,10 @@ def consume_task_result(task_result: TaskResult) -> Tuple[Path, int]:
     return path, return_code
 
 
-def create_pool(is_stdin_in_sources):
+def create_pool(is_stdin_in_sources, num_workers):
     if is_stdin_in_sources:
         return mp_dummy.Pool
-    return partial(mp.Pool, processes=mp.cpu_count())
+    return partial(mp.Pool, processes=num_workers)
 
 
 def filter_already_formatted_files(
@@ -163,12 +163,12 @@ def store_files_in_cache(
         cache.store_files(configuration_summary, files)
 
 
-def run(mode: Mode, configuration: Configuration, sources: Iterable[Path]):
+def run(mode: Mode, configuration: Configuration, num_workers, sources: Iterable[Path]):
     configuration_summary = configuration.summary()
     requested_files = get_files(sources)
     task = select_task(mode, configuration)
 
-    pool_cm = create_pool(Path("-") in requested_files)
+    pool_cm = create_pool(Path("-") in requested_files, num_workers)
     with create_cache() as cache, pool_cm() as pool:
         files_to_format = list(
             filter_already_formatted_files(