From 0acbbedfeab7973ca52b1c83da19229f6b2326e4 Mon Sep 17 00:00:00 2001 From: Hofer-Julian <30049909+Hofer-Julian@users.noreply.github.com> Date: Mon, 25 Mar 2024 13:38:37 +0100 Subject: [PATCH] Parallelize generate testmodels (#1309) Fixes #1302 On my machine, that reduces the runtime of `pixi run generate-testmodels` from 1m50s to 42s --- utils/generate-testmodels.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/utils/generate-testmodels.py b/utils/generate-testmodels.py index 2394efe30..516ef2281 100644 --- a/utils/generate-testmodels.py +++ b/utils/generate-testmodels.py @@ -1,8 +1,18 @@ +import multiprocessing import shutil +from functools import partial from pathlib import Path import ribasim_testmodels + +def generate_model(args, datadir): + model_name, model_constructor = args + model = model_constructor() + model.write(datadir / model_name / "ribasim.toml") + return model_name + + if __name__ == "__main__": datadir = Path("generated_testmodels") if datadir.is_dir(): @@ -18,7 +28,10 @@ Don't put important stuff in here, it will be emptied for every run.""" ) - for model_name, model_constructor in ribasim_testmodels.constructors.items(): - print(f"Generating {model_name}") - model = model_constructor() - model.write(datadir / model_name / "ribasim.toml") + generate_model_partial = partial(generate_model, datadir=datadir) + + with multiprocessing.Pool(processes=4) as p: + for model_name in p.imap_unordered( + generate_model_partial, ribasim_testmodels.constructors.items() + ): + print(f"Generated {model_name}")