Skip to content

Commit

Permalink
Use threads instead of processes for parallel 'nix eval'
Browse files Browse the repository at this point in the history
  • Loading branch information
GaetanLepage committed Oct 9, 2024
1 parent 9670630 commit 76308b1
Showing 1 changed file with 18 additions and 22 deletions.
40 changes: 18 additions & 22 deletions nixpkgs_review/nix.py
Original file line number Diff line number Diff line change
@@ -1,11 +1,10 @@
import concurrent.futures
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
Expand Down Expand Up @@ -271,32 +270,29 @@ 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}
results: dict[System, list[Attr]] = {}
with concurrent.futures.ThreadPoolExecutor(max_workers=n_procs) as executor:
future_to_system = {
executor.submit(
nix_eval,
attrs=attrs,
system=system,
allow=allow,
nix_path=nix_path,
): system
for system, attrs in attr_names_per_system.items()
}
for future in concurrent.futures.as_completed(future_to_system):
system = future_to_system[future]
results[system] = future.result()

return results


def nix_build(
Expand Down

0 comments on commit 76308b1

Please sign in to comment.