Skip to content

Implement concurrent MadMatrix builds across subprocesses #61

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Qubitol wants to merge 4 commits into
base: main
Choose a base branch
from
+177 −31
Open

Implement concurrent MadMatrix builds across subprocesses #61

Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
8a39247
Implement concurrent compilation of subprocesses
Qubitol Aug 10, 2026
eda3363
Fix race condition for commonlib compilation
Qubitol Aug 10, 2026
0a32282
Make cppauto deduction run before concurrent subprocesses compilation
Qubitol Aug 10, 2026
fa90560
Improve compilation logging by treating common library as a subprocess
Qubitol Aug 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
208 changes: 177 additions & 31 deletions madgraph/iolibs/template_files/mg7/madevent.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
from dataclasses import dataclass
from typing import Literal, NamedTuple
import resource
from concurrent.futures import ThreadPoolExecutor, as_completed

# Locate the madspace installation bundled alongside MadGraph.
# madgraph/__init__.py lives one level below the MadGraph root, so .parents[1]
Expand Down Expand Up @@ -402,6 +403,147 @@ def init_subprocesses(self) -> None:
for subproc_id, meta in enumerate(self.subprocess_data):
self.subprocesses.append(MadgraphSubprocess(self, meta, subproc_id))

verbosity = resolve_verbosity(self.run_card["run"]["verbosity"])
devices = self.run_card["run"]["devices"]
if not isinstance(devices, list):
devices = [devices]

# Resolve cppauto once using a representative subprocess, then reuse the
# resolved backend for the common library and every subprocess build.
self.cppauto_resolved = "cppauto"
if "cppauto" in devices and self.subprocesses:
result = subprocess.run(
["make", "-n", "BACKEND=cppauto", "detect-backend"],
cwd=self.subprocesses[0].meta["path"],
capture_output=True,
text=True,
check=True,
)
match = re.search(r"BACKEND=(\S+) \(was cppauto\)", result.stdout)
if match is None:
raise RuntimeError("Could not deduce the backend for cppauto")
self.cppauto_resolved = match.group(1)
logger.info("Device 'cppauto' deduced as '%s'", self.cppauto_resolved)

box = None
detailed_compile_view = False
completed_compile_count = 0
subprocess_count = len(self.subprocesses)
compile_item_count = subprocess_count + bool(self.subprocesses)
compile_item_label = (
f"{subprocess_count} + commonlib"
if self.subprocesses
else str(subprocess_count)
)
if verbosity == "pretty":
terminal_size = shutil.get_terminal_size(fallback=(91, 24))
available_rows = max(1, terminal_size.lines - 6)
detailed_compile_view = compile_item_count <= available_rows
box_width = min(91, terminal_size.columns)
device_labels = [
f"'{device}' -> '{self.cppauto_resolved}'"
if device == "cppauto"
else f"'{device}'"
for device in devices
]
title = (
"Compiling subprocesses for device"
+ "s" * (len(device_labels) > 1)
+ " "
+ ", ".join(device_labels)
)

if detailed_compile_view:
box = ms.PrettyBox(
title,
compile_item_count,
[0],
box_width=box_width,
)
if self.subprocesses:
box.set_row(0, ["commonlib..."])
for subproc in self.subprocesses:
box.set_row(subproc.id + 1, [f"{subproc.name}..."])
else:
box = ms.PrettyBox(
title,
3,
[16, 0],
box_width=box_width,
)
box.set_column(
0,
["Progress:", "Remaining:", "Last completed:"],
)
box.set_column(
1,
[
f"0 / ({compile_item_label})",
compile_item_label,
"-",
],
)
box.print_first()

def update_compile_status(item_id, name):
nonlocal completed_compile_count
completed_compile_count += 1

if verbosity == "pretty":
if detailed_compile_view:
box.set_row(item_id, [f"{name}...done!"])
else:
progress_label = (
f"{completed_compile_count} / ({compile_item_label}) "
)
progress_width = max(
5, box_width - 20 - len(progress_label)
)
progress = ms.format_progress(
completed_compile_count / compile_item_count,
progress_width,
)
box.set_column(
1,
[
f"{progress_label}{progress}",
str(compile_item_count - completed_compile_count),
name,
],
)
box.print_update()
elif verbosity == "log":
logger.info("Compiling subprocess %s...done!", name)

# Compile the common library as the first build item. Finishing this
# synchronous step is the barrier before subprocess workers are started.
if self.subprocesses:
if verbosity == "log":
logger.info("Compiling subprocess commonlib...")
common_build_path = self.subprocesses[0].meta["path"]
for device in devices:
device = self.cppauto_resolved if device == "cppauto" else device
misc.compile(
arg=[f"BACKEND={device}", "USEBUILDDIR=1", "commonlib"],
cwd=common_build_path,
mode="cpp",
)
update_compile_status(0, "commonlib")

max_workers = self.run_card["run"]["cpu_thread_pool_size"] if self.run_card["run"]["cpu_thread_pool_size"] > 0 else None
with ThreadPoolExecutor(max_workers = max_workers) as executor:
futures = {
executor.submit(subproc.compile): subproc
for subproc in self.subprocesses
}
for future in as_completed(futures):
subproc = futures[future]
future.result()
update_compile_status(subproc.id + 1, subproc.name)

for subproc in self.subprocesses:
subproc.load_matrix_element()

def build_event_generator(self, phasespaces: list[PhaseSpace]) -> ms.EventGenerator:
channel_generators = []
for i, (subproc, phasespace) in enumerate(zip(self.subprocesses, phasespaces)):
Expand Down Expand Up @@ -604,7 +746,7 @@ def train_madnis_old(self) -> None:
phasespace = subproc.build_madnis(phasespace)
if len(self.subprocesses) > 1:
status_func = lambda *args: self.update_madnis_status_multi(
subproc.subproc_id, *args
subproc.id, *args
)
else:
status_func = self.update_madnis_status_single
Expand Down Expand Up @@ -956,32 +1098,10 @@ class MadgraphSubprocess:
def __init__(self, process: MadgraphProcess, meta: dict, subproc_id: int):
self.process = process
self.meta = meta
self.subproc_id = subproc_id
self.name = os.path.basename(meta["path"].rstrip("/"))
self.id = subproc_id
self.multi_channel_data = None

api_path_format = self.meta["me_path"]
subproc_path = self.meta["path"]
devices = self.process.run_card["run"]["devices"]
api_paths = []
if not isinstance(devices, list):
devices = [devices]
for device in devices:
subproc_dir = os.path.dirname(subproc_path)
# 'cppauto' resolve quick fix
resolved = device
if device == "cppauto":
out = subprocess.run(
["make", "-n", "BACKEND=cppauto", "detect-backend"],
cwd=subproc_path, capture_output=True, text=True,
).stdout
match = re.search(r"BACKEND=(\S+) \(was cppauto\)", out)
if match:
resolved = match.group(1)
api_path = api_path_format.format(device=resolved)
if not os.path.isfile(api_path):
logger.info(f"Compiling subprocess {subproc_dir}, for device '{device}'")
misc.compile(arg = [f"BACKEND={device}", "USEBUILDDIR=1"], cwd = subproc_path)
api_paths.append(api_path)
self.api_paths = []

self.incoming_masses = [
self.process.get_mass(pid) for pid in clean_pids(self.meta["incoming"])
Expand Down Expand Up @@ -1022,10 +1142,36 @@ def __init__(self, process: MadgraphProcess, meta: dict, subproc_id: int):
particle_count=self.particle_count, **self.process.scale_kwargs
)

def compile(self):
api_path_format = self.meta["me_path"]
subproc_path = self.meta["path"]
verbosity = resolve_verbosity(self.process.run_card["run"]["verbosity"])
devices = self.process.run_card["run"]["devices"]
if not isinstance(devices, list):
devices = [devices]

if verbosity == "log":
logger.info("Compiling subprocess %s...", self.name)

for device in devices:
device = (
self.process.cppauto_resolved if device == "cppauto" else device
)
api_path = api_path_format.format(device=device)

if not os.path.isfile(api_path):
misc.compile(
arg=[f"BACKEND={device}", "USEBUILDDIR=1"],
cwd=subproc_path,
mode="cpp",
)
self.api_paths.append(api_path)

def load_matrix_element(self):
if self.process.run_card["run"]["dummy_matrix_element"]:
self.matrix_element = None
else:
for context, api_path in zip(self.process.contexts, api_paths):
for context, api_path in zip(self.process.contexts, self.api_paths):
self.matrix_element = context.load_matrix_element(
api_path, self.process.param_card_path
)
Expand Down Expand Up @@ -1147,7 +1293,7 @@ def build_multichannel_phasespace(self) -> PhaseSpace:
permutations=chan_permutations,
leptonic=self.process.leptonic,
)
prefix = f"subproc{self.subproc_id}.channel{channel_id}"
prefix = f"subproc{self.id}.channel{channel_id}"
if topo_count > 1:
prefix += f".subchan{topo_index}"
discrete_before, discrete_after = self.build_discrete(
Expand Down Expand Up @@ -1201,7 +1347,7 @@ def build_flat_phasespace(self) -> PhaseSpace:
cuts=self.cuts,
leptonic=self.process.leptonic,
)
prefix = f"subproc{self.subproc_id}.flat"
prefix = f"subproc{self.id}.flat"
discrete_before, discrete_after = self.build_discrete(
1, len(self.meta["flavors"]), prefix
)
Expand Down Expand Up @@ -1303,7 +1449,7 @@ def build_madnis(self, phasespace: PhaseSpace) -> PhaseSpace:
madnis_args = self.process.run_card["madnis"]
channels = []
for channel_id, channel in enumerate(phasespace.channels):
prefix = f"subproc{self.subproc_id}.channel{channel_id}"
prefix = f"subproc{self.id}.channel{channel_id}"
cond_dim = 0

discrete_before = channel.discrete_before
Expand Down Expand Up @@ -1418,7 +1564,7 @@ def build_cwnet(self, channel_count: int) -> ms.ChannelWeightNetwork:
hidden_dim=madnis_args["cwnet_hidden_dim"],
layers=madnis_args["cwnet_layers"],
activation=self.activation(madnis_args["cwnet_activation"]),
prefix=f"subproc{self.subproc_id}.cwnet",
prefix=f"subproc{self.id}.cwnet",
)
cwnet.initialize_globals(self.process.contexts[0])
return cwnet
Expand Down
Loading