Add --ncores flag to specify the nb of cores used to parallelize evolution

Author: Radonirinaunimi

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

@@ -1,13 +1,13 @@
 from collections import defaultdict
 import json
 import logging
-import multiprocessing
 import pathlib
 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
@@ -24,11 +24,19 @@
     "level": logging.DEBUG,
 }
 
-NUM_CORES = multiprocessing.cpu_count()
+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
@@ -125,7 +133,9 @@ 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_")))
 
-        Parallel(n_jobs=NUM_CORES)(
+        # Choose the number of cores to be the Minimal value
+        nb_cores = min(NUM_CORES, ncores)
+        Parallel(n_jobs=nb_cores)(
             delayed(_wrap_evolve)(pdf, r) for r, pdf in initial_PDFs_dict.items()
         )
 

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
 
@@ -82,6 +82,13 @@ def construct_evolven3fit_parser(subparsers):
     parser.add_argument(
         "-d", "--dump", type=pathlib.Path, default=None, help="Path where the EKO is dumped"
     )
+    parser.add_argument(
+        "-n",
+        "--ncores",
+        type=int,
+        default=1,
+        help="Specify the number of cores to parallelize evolution",
+    )
     parser.add_argument(
         "-f",
         "--force",
@@ -169,6 +176,7 @@ def main():
             args.force,
             eko_path,
             None,
+            args.ncores,
         )
     else:
         # If we are in the business of producing an eko, do some checks before starting: