From 4e59624df18fc3ba543f53b0da76f44943218777 Mon Sep 17 00:00:00 2001 From: Silia Taider Date: Wed, 21 Oct 2026 15:19:57 +0200 Subject: [PATCH] [Python][ML] Raise if ImplicitMT is enabled for the RDataLoader --- .../ROOT/_pythonization/_ml_dataloader.py | 14 ++++++++++++++ .../pyroot/pythonizations/test/ml_dataloader.py | 17 +++++++++++++++++ 2 files changed, 31 insertions(+) diff --git a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py index 2762ca18e64aa..bef5ea1683f02 100644 --- a/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py +++ b/bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py @@ -44,6 +44,16 @@ def _exclude_target_and_weights(columns: list[str], target_columns: list[str], w return [str(f) for f in columns if str(f) not in excluded] +def _raise_if_implicit_mt(rdfs: list) -> None: + """Raise if ImplicitMT is enabled, or was when one of the RDataFrames was created.""" + import ROOT + + if ROOT.IsImplicitMTEnabled() or any(rdf.GetNSlots() > 1 for rdf in rdfs): + raise RuntimeError( + "Cannot use RDataLoader with implicit multi-threading enabled, call ROOT.DisableImplicitMT() before creating the data loader and its inputs." + ) + + class _RDataLoader: def _resolve_columns( self, @@ -205,6 +215,7 @@ def __init__( if not hasattr(rdataframes, "__iter__"): rdataframes = [rdataframes] self.noded_rdfs = [RDF.AsRNode(rdf) for rdf in rdataframes] + _raise_if_implicit_mt(self.noded_rdfs) self.target_columns = _normalize_target(target) self.weights_column = weights @@ -272,10 +283,12 @@ def DeActivate(self): def ActivateTrainingEpoch(self): """Activate the training epoch of the generator""" + _raise_if_implicit_mt(self.noded_rdfs) self.engine.ActivateTrainingEpoch() def ActivateValidationEpoch(self): """Activate the validation epoch of the generator""" + _raise_if_implicit_mt(self.noded_rdfs) self.engine.ActivateValidationEpoch() def DeActivateTrainingEpoch(self): @@ -493,6 +506,7 @@ def GetValidationBatch(self) -> Any: return batch if (batch and batch.GetSize() > 0) else None def Save(self, dataset_name: str, filename: str, is_training: bool, output_format: str) -> None: + _raise_if_implicit_mt(self.noded_rdfs) self.engine.Save(dataset_name, filename, is_training, output_format) diff --git a/bindings/pyroot/pythonizations/test/ml_dataloader.py b/bindings/pyroot/pythonizations/test/ml_dataloader.py index 7d0cf6bd07f5f..98200fcc245e9 100644 --- a/bindings/pyroot/pythonizations/test/ml_dataloader.py +++ b/bindings/pyroot/pythonizations/test/ml_dataloader.py @@ -8123,5 +8123,22 @@ def test08_round_trip(self): np.testing.assert_array_equal(w.flatten(), np.full(20, 0.5)) +class DataLoaderImplicitMT(unittest.TestCase): + file_name = "dataloader_imt.root" + tree_name = "mytree" + + def tearDown(self): + ROOT.DisableImplicitMT() + os.remove(self.file_name) + + def test01_implicit_mt_not_supported(self): + ROOT.RDataFrame(10).Define("b1", "(int) rdfentry_").Snapshot(self.tree_name, self.file_name) + ROOT.EnableImplicitMT(2) + df = ROOT.RDataFrame(self.tree_name, self.file_name) + + with self.assertRaisesRegex(RuntimeError, "Cannot use RDataLoader with implicit multi-threading enabled"): + ROOT.Experimental.ML.RDataLoader(df, batch_size=2).as_numpy() + + if __name__ == "__main__": unittest.main()