Repository: cta-observatory/ctapipe (python)

Issue:
Allow overriding number of cores used by reconstructors on the CLI
**Please describe the use case that requires this feature.**

Useful for GRID processing / other job systems with resource management.


**Describe the solution you'd like**

Add a CLI option that allows overriding the number of cores used by reconstructors e.g. SKLearn based reconstructors to the train and apply tools.

This is needed to be able to set different n_cores for training and application.

**Describe alternatives you've considered**


**Additional context**



Pull request description:
Allow setting n_jobs on tool invocation
- It can be useful to set n_jobs indepently from the config file
- For the sklearn reconstructors this requires setting n_jobs on every model, that is attached to the reconstructor
- Fixes #2307

Patch:
diff --git a/ctapipe/reco/reconstructor.py b/ctapipe/reco/reconstructor.py
index 28402449..57a7a8a6 100644
--- a/ctapipe/reco/reconstructor.py
+++ b/ctapipe/reco/reconstructor.py
@@ -8,7 +8,7 @@ from astropy.coordinates import AltAz, SkyCoord
 
 from ctapipe.containers import ArrayEventContainer, TelescopeImpactParameterContainer
 from ctapipe.core import Provenance, QualityQuery, TelescopeComponent
-from ctapipe.core.traits import List
+from ctapipe.core.traits import Integer, List
 
 from ..compat import StrEnum
 from ..coordinates import shower_impact_distance
@@ -70,9 +70,15 @@ class Reconstructor(TelescopeComponent):
     algorithms should inherit from
     """
 
-    #: ctapipe_rco entry points may provide Reconstructor implementations
+    #: ctapipe_reco entry points may provide Reconstructor implementations
     plugin_entry_point = "ctapipe_reco"
 
+    n_jobs = Integer(
+        default_value=None,
+        allow_none=True,
+        help="Number of threads to use for the reconstruction if supported by the reconstructor.",
+    ).tag(config=True)
+
     def __init__(self, subarray, **kwargs):
         super().__init__(subarray=subarray, **kwargs)
         self.quality_query = StereoQualityQuery(parent=self)
diff --git a/ctapipe/reco/sklearn.py b/ctapipe/reco/sklearn.py
index fec334d9..230f6af3 100644
--- a/ctapipe/reco/sklearn.py
+++ b/ctapipe/reco/sklearn.py
@@ -16,7 +16,7 @@ from sklearn.metrics import accuracy_score, r2_score, roc_auc_score
 from sklearn.model_selection import KFold, StratifiedKFold
 from sklearn.utils import all_estimators
 from tqdm import tqdm
-from traitlets import TraitError
+from traitlets import TraitError, observe
 
 from ctapipe.exceptions import TooFewEvents
 
@@ -103,7 +103,7 @@ class SKLearnReconstructor(Reconstructor):
         help="If given, load serialized model from this path",
     ).tag(config=True)
 
-    def __init__(self, subarray=None, models=None, **kwargs):
+    def __init__(self, subarray=None, models=None, n_jobs=None, **kwargs):
         # Run the Component __init__ first to handle the configuration
         # and make `self.load_path` available
         Component.__init__(self, **kwargs)
@@ -199,7 +199,10 @@ class SKLearnReconstructor(Reconstructor):
         return QTable(self.subarray.to_table("joined"))
 
     def _new_model(self):
-        return SUPPORTED_MODELS[self.model_cls](**self.model_config)
+        cfg = self.model_config
+        if self.n_jobs:
+            cfg["n_jobs"] = self.n_jobs
+        return SUPPORTED_MODELS[self.model_cls](**cfg)
 
     def _table_to_y(self, table, mask=None):
         """
@@ -222,6 +225,15 @@ class SKLearnReconstructor(Reconstructor):
         y = self._table_to_y(table, mask=valid)
         self._models[key].fit(X, y)
 
+    @observe("n_jobs")
+    def _set_n_jobs(self, n_jobs):
+        """
+        Update n_jobs of all associated models.
+        """
+        if hasattr(self, "_models"):
+            for model in self._models.values():
+                model.n_jobs = n_jobs.new
+
 
 class SKLearnRegressionReconstructor(SKLearnReconstructor):
     """
@@ -562,7 +574,6 @@ class DispReconstructor(Reconstructor):
 
             # to verify settings
             self._new_models()
-
             self._models = {} if models is None else models
             self.unit = None
             self.stereo_combiner = StereoCombiner.from_name(
@@ -584,8 +595,13 @@ class DispReconstructor(Reconstructor):
             self.subarray = subarray
 
     def _new_models(self):
-        norm_regressor = SUPPORTED_REGRESSORS[self.norm_cls](**self.norm_config)
-        sign_classifier = SUPPORTED_CLASSIFIERS[self.sign_cls](**self.sign_config)
+        norm_cfg = self.norm_config
+        sign_cfg = self.sign_config
+        if self.n_jobs:
+            norm_cfg["n_jobs"] = self.n_jobs
+            sign_cfg["n_jobs"] = self.n_jobs
+        norm_regressor = SUPPORTED_REGRESSORS[self.norm_cls](**norm_cfg)
+        sign_classifier = SUPPORTED_CLASSIFIERS[self.sign_cls](**sign_cfg)
         return norm_regressor, sign_classifier
 
     def _table_to_y(self, table, mask=None):
@@ -803,6 +819,16 @@ class DispReconstructor(Reconstructor):
             ReconstructionProperty.GEOMETRY: altaz_result,
         }
 
+    @observe("n_jobs")
+    def _set_n_jobs(self, n_jobs):
+        """
+        Update n_jobs of all associated models.
+        """
+        if hasattr(self, "_models"):
+            for (disp, sign) in self._models.values():
+                disp.n_jobs = n_jobs.new
+                sign.n_jobs = n_jobs.new
+
 
 class CrossValidator(Component):
     """Class to train sklearn based reconstructors in a cross validation"""
diff --git a/ctapipe/resources/train_disp_reconstructor.yaml b/ctapipe/resources/train_disp_reconstructor.yaml
index 499c0684..1a798782 100644
--- a/ctapipe/resources/train_disp_reconstructor.yaml
+++ b/ctapipe/resources/train_disp_reconstructor.yaml
@@ -19,13 +19,15 @@ TrainDispReconstructor:
     # prefix:  # Add a prefix of the output here, if you want to apply multiple
                # DispReconstructors on the same file (e.g. for comparing different settings)
 
+    # How many cores to use. Overwrites model config 
+    n_jobs: -1
+
     # All regression algorithms in scikit-learn are supported
     # (https://scikit-learn.org/stable/modules/classes.html)
     norm_cls: ExtraTreesRegressor
     norm_config:
       n_estimators: 10
       max_depth: 10
-      n_jobs: -1
 
     log_target: True
 
@@ -35,7 +37,6 @@ TrainDispReconstructor:
     sign_config:
       n_estimators: 10
       max_depth: 10
-      n_jobs: -1
 
     QualityQuery:  # Event Selection performed before training the models
       quality_criteria:
diff --git a/ctapipe/resources/train_energy_regressor.yaml b/ctapipe/resources/train_energy_regressor.yaml
index a2ecdcfa..e33a37e5 100644
--- a/ctapipe/resources/train_energy_regressor.yaml
+++ b/ctapipe/resources/train_energy_regressor.yaml
@@ -20,6 +20,7 @@ TrainEnergyRegressor:
                # EnergyRegressors on the same file (e.g. for comparing different settings)
 
     log_target: True
+    n_jobs: -1
 
     # All regression algorithms in scikit-learn are supported
     # (https://scikit-learn.org/stable/modules/classes.html)
@@ -27,7 +28,6 @@ TrainEnergyRegressor:
     model_config:
       n_estimators: 10
       max_depth: 10
-      n_jobs: -1
 
     QualityQuery:  # Event Selection performed before training the models
       quality_criteria:
diff --git a/ctapipe/resources/train_particle_classifier.yaml b/ctapipe/resources/train_particle_classifier.yaml
index ae4a8127..8b999a8e 100644
--- a/ctapipe/resources/train_particle_classifier.yaml
+++ b/ctapipe/resources/train_particle_classifier.yaml
@@ -22,13 +22,15 @@ TrainParticleClassifier:
     # prefix:  # Add a prefix of the output here, if you want to apply multiple
                # ParticleClassifiers on the same file (e.g. for comparing different settings)
 
+    # How many cores to use. Overwrites model config 
+    n_jobs: -1
+
     # All classification algorithms in scikit-learn are supported 
     # (https://scikit-learn.org/stable/modules/classes.html)
     model_cls: ExtraTreesClassifier
     model_config:
       n_estimators: 10
       max_depth: 10
-      n_jobs: -1
     
     QualityQuery:  # Event Selection performed before training the models
       quality_criteria:
diff --git a/ctapipe/tools/apply_models.py b/ctapipe/tools/apply_models.py
index 2df691d8..dee234f6 100644
--- a/ctapipe/tools/apply_models.py
+++ b/ctapipe/tools/apply_models.py
@@ -68,6 +68,12 @@ class ApplyModels(Tool):
         help="How many subarray events to load at once for making predictions.",
     ).tag(config=True)
 
+    n_jobs = Integer(
+        default_value=None,
+        allow_none=True,
+        help="Number of threads to use for the reconstruction. This overwrites the values in the config",
+    ).tag(config=True)
+
     progress_bar = Bool(
         help="show progress bar during processing",
         default_value=True,
@@ -77,6 +83,7 @@ class ApplyModels(Tool):
         ("i", "input"): "ApplyModels.input_url",
         ("r", "reconstructor"): "ApplyModels.reconstructor_paths",
         ("o", "output"): "ApplyModels.output_path",
+        "n-jobs": "ApplyModels.n_jobs",
         "chunk-size": "ApplyModels.chunk_size",
     }
 
@@ -146,10 +153,12 @@ class ApplyModels(Tool):
             )
         )
 
-        self._reconstructors = [
-            Reconstructor.read(path, parent=self, subarray=self.loader.subarray)
-            for path in self.reconstructor_paths
-        ]
+        self._reconstructors = []
+        for path in self.reconstructor_paths:
+            r = Reconstructor.read(path, parent=self, subarray=self.loader.subarray)
+            if self.n_jobs:
+                r.n_jobs = self.n_jobs
+            self._reconstructors.append(r)
 
     def start(self):
         """Apply models to input tables"""
diff --git a/ctapipe/tools/train_disp_reconstructor.py b/ctapipe/tools/train_disp_reconstructor.py
index a5a968d2..89ac2eeb 100644
--- a/ctapipe/tools/train_disp_reconstructor.py
+++ b/ctapipe/tools/train_disp_reconstructor.py
@@ -66,6 +66,12 @@ class TrainDispReconstructor(Tool):
         default_value=0, help="Random seed for sampling and cross validation"
     ).tag(config=True)
 
+    n_jobs = Int(
+        default_value=None,
+        allow_none=True,
+        help="Number of threads to use for the reconstruction. This overwrites the values in the config of each reconstructor.",
+    ).tag(config=True)
+
     project_disp = Bool(
         default_value=False,
         help=(
@@ -80,6 +86,7 @@ class TrainDispReconstructor(Tool):
         ("i", "input"): "TableLoader.input_url",
         ("o", "output"): "TrainDispReconstructor.output_path",
         "n-events": "TrainDispReconstructor.n_events",
+        "n-jobs": "DispReconstructor.n_jobs",
         "cv-output": "CrossValidator.output_path",
     }
 
@@ -103,6 +110,7 @@ class TrainDispReconstructor(Tool):
         self.n_events.attach_subarray(self.loader.subarray)
 
         self.models = DispReconstructor(self.loader.subarray, parent=self)
+
         self.cross_validate = CrossValidator(parent=self, model_component=self.models)
         self.rng = np.random.default_rng(self.random_seed)
         self.check_output(self.output_path, self.cross_validate.output_path)
@@ -182,6 +190,7 @@ class TrainDispReconstructor(Tool):
         Write-out trained models and cross-validation results.
         """
         self.log.info("Writing output")
+        self.models.n_jobs = None
         self.models.write(self.output_path, overwrite=self.overwrite)
         if self.cross_validate.output_path:
             self.cross_validate.write(overwrite=self.overwrite)
diff --git a/ctapipe/tools/train_energy_regressor.py b/ctapipe/tools/train_energy_regressor.py
index 0d210762..383dedc7 100644
--- a/ctapipe/tools/train_energy_regressor.py
+++ b/ctapipe/tools/train_energy_regressor.py
@@ -63,11 +63,18 @@ class TrainEnergyRegressor(Tool):
         default_value=0, help="Random seed for sampling and cross validation"
     ).tag(config=True)
 
+    n_jobs = Int(
+        default_value=None,
+        allow_none=True,
+        help="Number of threads to use for the reconstruction. This overwrites the values in the config of each reconstructor.",
+    ).tag(config=True)
+
     aliases = {
         ("i", "input"): "TableLoader.input_url",
         ("o", "output"): "TrainEnergyRegressor.output_path",
         "n-events": "TrainEnergyRegressor.n_events",
         "chunk-size": "TrainEnergyRegressor.chunk_size",
+        "n-jobs": "EnergyRegressor.n_jobs",
         "cv-output": "CrossValidator.output_path",
     }
 
@@ -94,6 +101,7 @@ class TrainEnergyRegressor(Tool):
         self.n_events.attach_subarray(self.loader.subarray)
 
         self.regressor = EnergyRegressor(self.loader.subarray, parent=self)
+        self.log.warning(f"{self.regressor._models}")
         self.cross_validate = CrossValidator(
             parent=self, model_component=self.regressor
         )
@@ -137,6 +145,7 @@ class TrainEnergyRegressor(Tool):
         Write-out trained models and cross-validation results.
         """
         self.log.info("Writing output")
+        self.regressor.n_jobs = None
         self.regressor.write(self.output_path, overwrite=self.overwrite)
         if self.cross_validate.output_path:
             self.cross_validate.write(overwrite=self.overwrite)
diff --git a/ctapipe/tools/train_particle_classifier.py b/ctapipe/tools/train_particle_classifier.py
index 0002c683..caf4a455 100644
--- a/ctapipe/tools/train_particle_classifier.py
+++ b/ctapipe/tools/train_particle_classifier.py
@@ -92,11 +92,18 @@ class TrainParticleClassifier(Tool):
         help="Random number seed for sampling and the cross validation splitting",
     ).tag(config=True)
 
+    n_jobs = Int(
+        default_value=None,
+        allow_none=True,
+        help="Number of threads to use for the reconstruction. This overwrites the values in the config of each reconstructor.",
+    ).tag(config=True)
+
     aliases = {
         "signal": "TrainParticleClassifier.input_url_signal",
         "background": "TrainParticleClassifier.input_url_background",
         "n-signal": "TrainParticleClassifier.n_signal",
         "n-background": "TrainParticleClassifier.n_background",
+        "n-jobs": "ParticleClassifier.n_jobs",
         ("o", "output"): "TrainParticleClassifier.output_path",
         "cv-output": "CrossValidator.output_path",
     }
@@ -207,6 +214,7 @@ class TrainParticleClassifier(Tool):
         Write-out trained models and cross-validation results.
         """
         self.log.info("Writing output")
+        self.classifier.n_jobs = None
         self.classifier.write(self.output_path, overwrite=self.overwrite)
         self.signal_loader.close()
         self.background_loader.close()
diff --git a/docs/changes/2430.feature.rst b/docs/changes/2430.feature.rst
new file mode 100644
index 00000000..97917d6f
--- /dev/null
+++ b/docs/changes/2430.feature.rst
@@ -0,0 +1,3 @@
+Allow setting n_jobs on the command line for the
+train_* and apply_models tools using a new ``n_jobs`` flag.
+This temporarily overwrites any settings in the (model) config(s).


Test patch:
diff --git a/ctapipe/reco/tests/test_sklearn.py b/ctapipe/reco/tests/test_sklearn.py
index 50b4a016..e893cf94 100644
--- a/ctapipe/reco/tests/test_sklearn.py
+++ b/ctapipe/reco/tests/test_sklearn.py
@@ -9,6 +9,7 @@ from traitlets.config import Config
 from ctapipe.core import Component
 from ctapipe.reco import EnergyRegressor, ParticleClassifier
 from ctapipe.reco.reconstructor import ReconstructionProperty
+from ctapipe.reco.sklearn import DispReconstructor
 
 KEY = "LST_LST_LSTCam"
 
@@ -166,6 +167,49 @@ def test_regressor_single_event(model_cls, example_table, example_subarray):
     assert valid[0] == False
 
 
+def test_set_n_jobs(example_subarray):
+    config = Config(
+        {
+            "EnergyRegressor": {
+                "model_cls": "RandomForestRegressor",
+                "model_config": {"n_estimators": 20, "max_depth": 15, "n_jobs": -1},
+            }
+        }
+    )
+    regressor = EnergyRegressor(
+        example_subarray,
+        config=config,
+    )
+
+    regressor._models["telescope"] = regressor._new_model()
+    assert regressor._models["telescope"].n_jobs == -1
+    regressor.n_jobs = 42
+    assert regressor._models["telescope"].n_jobs == 42
+
+    # DISP has two models per telescope, check that aswell
+    config = Config(
+        {
+            "DispReconstructor": {
+                "norm_cls": "RandomForestRegressor",
+                "norm_config": {"n_estimators": 20, "max_depth": 15, "n_jobs": -1},
+                "sign_cls": "RandomForestClassifier",
+                "sign_config": {"n_estimators": 20, "max_depth": 15, "n_jobs": -1},
+            }
+        }
+    )
+    disp = DispReconstructor(
+        example_subarray,
+        config=config,
+    )
+
+    disp._models["telescope"] = disp._new_models()
+    assert disp._models["telescope"][0].n_jobs == -1
+    assert disp._models["telescope"][1].n_jobs == -1
+    disp.n_jobs = 42
+    assert disp._models["telescope"][0].n_jobs == 42
+    assert disp._models["telescope"][1].n_jobs == 42
+
+
 @pytest.mark.parametrize(
     "model_cls", ["KNeighborsClassifier", "RandomForestClassifier"]
 )