[ML] Data Loader: Introduce asynchronous lazy loading

bindings/pyroot/pythonizations/python/ROOT/_pythonization/_ml_dataloader.py

@@ -3,6 +3,7 @@
 # Author: Nopphakorn Subsa-Ard, King Mongkut's University of Technology Thonburi (KMUTT) (TH) 08/2024
 # Author: Vincenzo Eduardo Padulano, CERN 10/2024
 # Author: Martin Føll, University of Oslo (UiO) & CERN 01/2026
+# Author: Silia Taider, CERN 02/2026
 
 ################################################################################
 # Copyright (C) 1995-2026, Rene Brun and Fons Rademakers.                      #
@@ -270,42 +271,45 @@ def __init__(
         atexit.register(self.DeActivate)
 
     @property
-    def is_active(self):
+    def isActive(self):
         return self.generator.IsActive()
 
-    def is_training_active(self):
-        return self.generator.TrainingIsActive()
+    def isTrainingActive(self):
+        return self.generator.IsTrainingActive()
+
+    def isValidationActive(self):
+        return self.generator.IsValidationActive()
 
     def Activate(self):
-        """Initialize the generator to be used for a loop"""
+        """Initialize the generator to be used for a loop, this spawns the loading thread"""
         self.generator.Activate()
 
     def DeActivate(self):
         """Deactivate the generator"""
         self.generator.DeActivate()
 
     def ActivateTrainingEpoch(self):
-        """Activate the generator"""
+        """Activate the training epoch of the generator"""
         self.generator.ActivateTrainingEpoch()
 
     def ActivateValidationEpoch(self):
-        """Activate the generator"""
+        """Activate the validation epoch of the generator"""
         self.generator.ActivateValidationEpoch()
 
     def DeActivateTrainingEpoch(self):
-        """Deactivate the generator"""
+        """Deactivate the training epoch of the generator"""
         self.generator.DeActivateTrainingEpoch()
 
     def DeActivateValidationEpoch(self):
-        """Deactivate the generator"""
+        """Deactivate the validation epoch of the generator"""
         self.generator.DeActivateValidationEpoch()
 
     def CreateTrainBatches(self):
-        """Deactivate the generator"""
+        """Create the first training batches from the first chunk"""
         self.generator.CreateTrainBatches()
 
     def CreateValidationBatches(self):
-        """Deactivate the generator"""
+        """Create the first validation batches from the first chunk"""
         self.generator.CreateValidationBatches()
 
     def GetSample(self):
@@ -501,15 +505,17 @@ def GetValidationBatch(self) -> Any:
 class LoadingThreadContext:
     def __init__(self, base_generator: BaseGenerator):
         self.base_generator = base_generator
+        # init loading thread
+        self.base_generator.Activate()
         # create training batches from the first chunk
         self.base_generator.CreateTrainBatches()
 
     def __enter__(self):
         self.base_generator.ActivateTrainingEpoch()
+        return self
 
     def __exit__(self, type, value, traceback):
         self.base_generator.DeActivateTrainingEpoch()
-        return True
 
 
 class TrainDataLoader:
@@ -593,15 +599,17 @@ def __call__(self) -> Any:
 class LoadingThreadContextVal:
     def __init__(self, base_generator: BaseGenerator):
         self.base_generator = base_generator
+        # init loading thread
+        self.base_generator.Activate()
         # create validation batches from the first chunk
         self.base_generator.CreateValidationBatches()
 
     def __enter__(self):
         self.base_generator.ActivateValidationEpoch()
+        return self
 
     def __exit__(self, type, value, traceback):
         self.base_generator.DeActivateValidationEpoch()
-        return True
 
 
 class ValidationDataLoader:
@@ -668,7 +676,6 @@ def __call__(self) -> Any:
             while True:
                 batch = self.base_generator.GetValidationBatch()
                 if batch is None:
-                    self.base_generator.DeActivateValidationEpoch()
                     break
                 yield self.conversion_function(batch)
 
@@ -807,7 +814,7 @@ def CreateNumPyGenerators(
     train_generator = TrainDataLoader(base_generator, base_generator.ConvertBatchToNumpy)
 
     if validation_split == 0.0:
-        return train_generator, None
+        return train_generator
 
     validation_generator = ValidationDataLoader(base_generator, base_generator.ConvertBatchToNumpy)
 
@@ -946,7 +953,6 @@ def CreateTFDatasets(
     )
 
     train_generator = TrainDataLoader(base_generator, base_generator.ConvertBatchToTF)
-    validation_generator = ValidationDataLoader(base_generator, base_generator.ConvertBatchToTF)
 
     num_train_columns = len(train_generator.train_columns)
     num_target_columns = len(train_generator.target_columns)
@@ -982,6 +988,8 @@ def CreateTFDatasets(
     if validation_split == 0.0:
         return ds_train
 
+    validation_generator = ValidationDataLoader(base_generator, base_generator.ConvertBatchToTF)
+
     ds_validation = tf.data.Dataset.from_generator(validation_generator, output_signature=batch_signature)
 
     # Give access to the columns function of the validation set

bindings/pyroot/pythonizations/test/ml_dataloader.py

@@ -540,7 +540,7 @@ def test09_filtered_last_chunk(self):
 
             dff = df.Filter("b1 % 2 == 0", "name")
 
-            gen_train, _ = ROOT.Experimental.ML.CreateNumPyGenerators(
+            gen_train = ROOT.Experimental.ML.CreateNumPyGenerators(
                 dff,
                 batch_size=3,
                 chunk_size=9,
@@ -1622,7 +1622,7 @@ def test09_filtered_last_chunk(self):
 
             dff = df.Filter("b1 % 2 == 0", "name")
 
-            gen_train, _ = ROOT.Experimental.ML.CreateNumPyGenerators(
+            gen_train = ROOT.Experimental.ML.CreateNumPyGenerators(
                 dff, batch_size=3, target="b2", validation_split=0, shuffle=False, drop_remainder=False, load_eager=True
             )
 
@@ -2784,7 +2784,7 @@ def test09_filtered_last_chunk(self):
             dff1 = df1.Filter("b1 % 2 == 0", "name")
             dff2 = df2.Filter("b1 % 2 == 0", "name")
 
-            gen_train, _ = ROOT.Experimental.ML.CreateNumPyGenerators(
+            gen_train = ROOT.Experimental.ML.CreateNumPyGenerators(
                 [dff1, dff2],
                 batch_size=3,
                 target="b2",