diff --git a/madgraph/iolibs/template_files/mg7/madevent.py b/madgraph/iolibs/template_files/mg7/madevent.py index 6114a4871..0dd23ffdd 100644 --- a/madgraph/iolibs/template_files/mg7/madevent.py +++ b/madgraph/iolibs/template_files/mg7/madevent.py @@ -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] @@ -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)): @@ -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 @@ -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"]) @@ -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 ) @@ -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( @@ -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 ) @@ -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 @@ -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