Merge pull request #2168 from NNPDF/parallelize-evolution

Parallelize Evolution at the level of `evolven3fit`

Author: scarlehoff

conda-recipe/meta.yaml

@@ -40,6 +40,7 @@ requirements:
         - eko >=0.14.2
         - fiatlux
         - sphinx >=5.0.2
+        - joblib
         - sphinx_rtd_theme >0.5
         - sphinxcontrib-bibtex
         - ruamel.yaml <0.18

n3fit/src/evolven3fit/cli.py

@@ -2,7 +2,7 @@
 
 
 def cli_evolven3fit(
-    configuration_folder, q_fin, q_points, op_card_info, theory_card_info, force, load, dump
+    configuration_folder, q_fin, q_points, op_card_info, theory_card_info, force, load, dump, ncores
 ):
     """Evolves the fitted PDFs.
 
@@ -23,5 +23,13 @@ def cli_evolven3fit(
     """
     utils.check_is_a_fit(configuration_folder)
     return evolve.evolve_fit(
-        configuration_folder, q_fin, q_points, op_card_info, theory_card_info, force, load, dump
+        configuration_folder,
+        q_fin,
+        q_points,
+        op_card_info,
+        theory_card_info,
+        force,
+        load,
+        dump,
+        ncores,
     )

n3fit/src/evolven3fit/evolve.py

@@ -5,7 +5,9 @@
 import sys
 
 from ekobox import apply, genpdf, info_file
+from joblib import Parallel, delayed
 import numpy as np
+import psutil
 
 import eko
 from eko import basis_rotation, runner
@@ -22,9 +24,19 @@
     "level": logging.DEBUG,
 }
 
+NUM_CORES = psutil.cpu_count(logical=False)
+
 
 def evolve_fit(
-    fit_folder, q_fin, q_points, op_card_dict, theory_card_dict, force, eko_path, dump_eko=None
+    fit_folder,
+    q_fin,
+    q_points,
+    op_card_dict,
+    theory_card_dict,
+    force,
+    eko_path,
+    dump_eko=None,
+    ncores=1,
 ):
     """
     Evolves all the fitted replica in fit_folder/nnfit
@@ -117,10 +129,16 @@ def evolve_fit(
         info["NumFlavors"] = theory.heavy.num_flavs_max_pdf
         dump_info_file(usr_path, info)
 
-        for replica, pdf_data in initial_PDFs_dict.items():
-            evolved_blocks = evolve_exportgrid(pdf_data, eko_op, x_grid)
+        def _wrap_evolve(pdf, replica):
+            evolved_blocks = evolve_exportgrid(pdf, eko_op, x_grid)
             dump_evolved_replica(evolved_blocks, usr_path, int(replica.removeprefix("replica_")))
 
+        # Choose the number of cores to be the Minimal value
+        nb_cores = min(NUM_CORES, abs(ncores))
+        Parallel(n_jobs=nb_cores)(
+            delayed(_wrap_evolve)(pdf, r) for r, pdf in initial_PDFs_dict.items()
+        )
+
     # remove folder:
     # The function dump_evolved_replica dumps the replica files in a temporary folder
     # We need then to remove it after fixing the position of those replica files

n3fit/src/n3fit/scripts/evolven3fit.py

@@ -7,10 +7,10 @@
 import pathlib
 import sys
 
-from evolven3fit import cli, eko_utils, evolve, utils
 import numpy as np
 
 from eko.runner.managed import solve
+from evolven3fit import cli, eko_utils, evolve, utils
 from n3fit.io.writer import XGRID
 from validphys.loader import FallbackLoader, Loader
 
@@ -169,6 +169,7 @@ def main():
             args.force,
             eko_path,
             None,
+            args.n_cores,
         )
     else:
         # If we are in the business of producing an eko, do some checks before starting:

pyproject.toml

@@ -73,6 +73,7 @@ reportengine = { git = "https://github.com/NNPDF/reportengine" }
 psutil = "*"
 tensorflow = "*"
 eko = "^0.14.1"
+joblib = "*"
 # Hyperopt
 hyperopt = "*"
 seaborn = "*"