Parallelize nix_eval

nixpkgs_review/cli/__init__.py

@@ -249,6 +249,12 @@ def common_flags() -> list[CommonFlag]:
             default="{ }",
             help="Extra nixpkgs config to pass to `import <nixpkgs>`",
         ),
+        CommonFlag(
+            "--num-procs-eval",
+            type=int,
+            default=1,
+            help="Number of parallel `nix-env` processes to run simultaneously (warning, can imply heavy RAM usage)",
+        ),
     ]
 
 

nixpkgs_review/cli/pr.py

@@ -86,6 +86,7 @@ def pr_command(args: argparse.Namespace) -> str:
                     build_graph=args.build_graph,
                     nixpkgs_config=nixpkgs_config,
                     extra_nixpkgs_config=args.extra_nixpkgs_config,
+                    n_procs_eval=args.num_procs_eval,
                 )
                 contexts.append((pr, builddir.path, review.build_pr(pr)))
             except NixpkgsReviewError as e:

nixpkgs_review/nix.py

@@ -1,9 +1,11 @@
 import json
+import multiprocessing as mp
 import os
 import shlex
 import shutil
 import subprocess
 from dataclasses import dataclass, field
+from functools import partial
 from pathlib import Path
 from sys import platform
 from tempfile import NamedTemporaryFile
@@ -269,6 +271,34 @@ def nix_eval(
             os.unlink(attr_json.name)
 
 
+def nix_eval_thread(
+    system: System,
+    attr_names: set[str],
+    allow: AllowedFeatures,
+    nix_path: str,
+) -> tuple[System, list[Attr]]:
+    return system, nix_eval(attr_names, system, allow, nix_path)
+
+
+def multi_system_eval(
+    attr_names_per_system: dict[System, set[str]],
+    allow: AllowedFeatures,
+    nix_path: str,
+    n_procs: int,
+) -> dict[System, list[Attr]]:
+    nix_eval_partial = partial(
+        nix_eval_thread,
+        allow=allow,
+        nix_path=nix_path,
+    )
+
+    args: list[tuple[System, set[str]]] = list(attr_names_per_system.items())
+    with mp.Pool(n_procs) as pool:
+        results: list[tuple[System, list[Attr]]] = pool.starmap(nix_eval_partial, args)
+
+    return {system: attrs for system, attrs in results}
+
+
 def nix_build(
     attr_names_per_system: dict[System, set[str]],
     args: str,
@@ -279,19 +309,23 @@ def nix_build(
     build_graph: str,
     nix_path: str,
     nixpkgs_config: Path,
+    n_procs_eval: int,
 ) -> dict[System, list[Attr]]:
     if not attr_names_per_system:
         info("Nothing to be built.")
         return {}
 
-    # TODO parallelize evaluation
-    attrs_per_system: dict[System, list[Attr]] = {}
-    filtered_per_system: dict[System, list[str]] = {}
-    for system, attr_names in attr_names_per_system.items():
-        attrs_per_system[system] = nix_eval(attr_names, system, allow, nix_path)
+    attrs_per_system: dict[System, list[Attr]] = multi_system_eval(
+        attr_names_per_system,
+        allow,
+        nix_path,
+        n_procs=n_procs_eval,
+    )
 
+    filtered_per_system: dict[System, list[str]] = {}
+    for system, attrs in attrs_per_system.items():
         filtered_per_system[system] = []
-        for attr in attrs_per_system[system]:
+        for attr in attrs:
             if not (attr.broken or attr.blacklisted):
                 filtered_per_system[system].append(attr.name)
 

nixpkgs_review/review.py

@@ -107,6 +107,7 @@ def __init__(
         skip_packages_regex: list[Pattern[str]] = [],
         checkout: CheckoutOption = CheckoutOption.MERGE,
         sandbox: bool = False,
+        n_procs_eval: int = 1,
     ) -> None:
         self.builddir = builddir
         self.build_args = build_args
@@ -139,6 +140,7 @@ def __init__(
         self.build_graph = build_graph
         self.nixpkgs_config = nixpkgs_config
         self.extra_nixpkgs_config = extra_nixpkgs_config
+        self.n_procs_eval = n_procs_eval
 
     def worktree_dir(self) -> str:
         return str(self.builddir.worktree_dir)
@@ -260,6 +262,7 @@ def build(
             self.build_graph,
             self.builddir.nix_path,
             self.nixpkgs_config,
+            self.n_procs_eval,
         )
 
     def build_pr(self, pr_number: int) -> dict[System, list[Attr]]:
@@ -624,6 +627,7 @@ def review_local_revision(
             build_graph=args.build_graph,
             nixpkgs_config=nixpkgs_config,
             extra_nixpkgs_config=args.extra_nixpkgs_config,
+            n_procs_eval=args.num_procs_eval,
         )
         review.review_commit(builddir.path, args.branch, commit, staged, print_result)
         return builddir.path