Fixed a wrong commit previously pushed, revert AMS stager to what it was (#53)

Signed-off-by: Loic Pottier <[email protected]>

Author: lpottier

src/AMSWorkflow/ams/stage.py

@@ -17,7 +17,6 @@
 from typing import Callable, List, Tuple
 import struct
 import signal
-import os
 
 import numpy as np
 
@@ -93,15 +92,10 @@ class Task(ABC):
     the staging mechanism.
     """
 
-    def __init__(self):
-        self.statistics = {"datasize" : 0, "duration" : 0}
-
     @abstractmethod
     def __call__(self):
         pass
 
-    def stats(self):
-        return self.statistics
 
 class ForwardTask(Task):
     """
@@ -118,7 +112,7 @@ def __init__(self, i_queue, o_queue, callback):
         """
         initializes a ForwardTask class with the queues and the callback.
         """
-        super().__init__()
+
         if not isinstance(callback, Callable):
             raise TypeError(f"{callback} argument is not Callable")
 
@@ -148,7 +142,6 @@ def __call__(self):
         the output to the output queue. In the case of receiving a 'termination' messages informs
         the tasks waiting on the output queues about the terminations and returns from the function.
         """
-        start = time.time()
 
         while True:
             # This is a blocking call
@@ -159,14 +152,9 @@ def __call__(self):
             elif item.is_process():
                 inputs, outputs = self._action(item.data())
                 self.o_queue.put(QueueMessage(MessageType.Process, DataBlob(inputs, outputs)))
-                self.statistics["datasize"] += (inputs.nbytes + outputs.nbytes)
             elif item.is_new_model():
                 # This is not handled yet
                 continue
-
-        end = time.time()
-        self.statistics["duration"] = end - start
-        print(f"Spend {end - start} at {self.__class__.__name__} ({self.statistics})")
         return
 
 
@@ -183,7 +171,6 @@ class FSLoaderTask(Task):
     """
 
     def __init__(self, o_queue, loader, pattern):
-        super().__init__()
         self.o_queue = o_queue
         self.pattern = pattern
         self.loader = loader
@@ -206,13 +193,10 @@ def __call__(self):
                 output_batches = np.array_split(output_data, num_batches)
                 for j, (i, o) in enumerate(zip(input_batches, output_batches)):
                     self.o_queue.put(QueueMessage(MessageType.Process, DataBlob(i, o)))
-                self.statistics["datasize"] += (input_data.nbytes + output_data.nbytes)
-
         self.o_queue.put(QueueMessage(MessageType.Terminate, None))
 
         end = time.time()
-        self.statistics["duration"] += (end - start)
-        print(f"Spend {end - start} at {self.__class__.__name__} ({self.statistics})")
+        print(f"Spend {end - start} at {self.__class__.__name__}")
 
 
 class RMQMessage(object):
@@ -359,8 +343,7 @@ class RMQLoaderTask(Task):
         prefetch_count: Number of messages prefected by RMQ (impact performance)
     """
 
-    def __init__(self, o_queue, credentials, cacert, rmq_queue, prefetch_count = 1):
-        super().__init__()
+    def __init__(self, o_queue, credentials, cacert, rmq_queue, prefetch_count=1):
         self.o_queue = o_queue
         self.credentials = credentials
         self.cacert = cacert
@@ -404,8 +387,6 @@ def callback_message(self, ch, basic_deliver, properties, body):
         input_batches = np.array_split(input_data, num_batches)
         output_batches = np.array_split(output_data, num_batches)
 
-        self.statistics["datasize"] += (input_data.nbytes + output_data.nbytes)
-
         for j, (i, o) in enumerate(zip(input_batches, output_batches)):
             self.o_queue.put(QueueMessage(MessageType.Process, DataBlob(i, o)))
 
@@ -416,14 +397,15 @@ def handler(signum, frame):
             print(f"Received SIGNUM={signum} for {name}[pid={pid}]: stopping process")
             self.rmq_consumer.stop()
             self.o_queue.put(QueueMessage(MessageType.Terminate, None))
-            self.statistics["duration"] += (end - start)
-            print(f"Spend {self.total_time} at {self.__class__.__name__} ({self.statistics})")
+            print(f"Spend {self.total_time} at {self.__class__.__name__}")
 
         return handler
 
     def __call__(self):
         """
-        Busy loop of consuming messages from RMQ queue
+        Busy loop of reading all files matching the pattern and creating
+        '100' batches which will be pushed on the queue. Upon reading all files
+        the Task pushes a 'Terminate' message to the queue and returns.
         """
         self.rmq_consumer.run()
 
@@ -444,7 +426,6 @@ def __init__(self, i_queue, o_queue, writer_cls, out_dir):
         initializes the writer task to read data from the i_queue write them using
         the writer_cls and store the data in the out_dir.
         """
-        super().__init__()
         self.data_writer_cls = writer_cls
         self.out_dir = out_dir
         self.i_queue = i_queue
@@ -489,9 +470,7 @@ def __call__(self):
                 break
 
         end = time.time()
-        self.statistics["datasize"] = total_bytes_written
-        self.statistics["duration"] += (end - start)
-        print(f"Spend {end - start} {total_bytes_written} at {self.__class__.__name__} ({self.statistics})")
+        print(f"Spend {end - start} {total_bytes_written} at {self.__class__.__name__}")
 
 
 class PushToStore(Task):
@@ -512,7 +491,7 @@ def __init__(self, i_queue, ams_config, db_path, store):
         is not under db_path, it copies the file to this location and if store defined
         it makes the kosh-store aware about the existence of the file.
         """
-        super().__init__()
+
         self.ams_config = ams_config
         self.i_queue = i_queue
         self.dir = Path(db_path).absolute()
@@ -542,12 +521,9 @@ def __call__(self):
 
                 if self._store:
                     db_store.add_candidates([str(dest_file)])
-            
-                self.statistics["datasize"] += os.stat(src_fn).st_size
 
         end = time.time()
-        self.statistics["duration"] += (end - start)
-        print(f"Spend {end - start} at {self.__class__.__name__} ({self.statistics})")
+        print(f"Spend {end - start} at {self.__class__.__name__}")
 
 
 class Pipeline(ABC):

src/AMSlib/ml/uq.hpp

@@ -140,7 +140,8 @@ class UQ
     }
   }
 
-  void updateModel(std::string model_path, std::string uq_path = "")
+  void updateModel(const std::string &model_path,
+                   const std::string &uq_path = "")
   {
     if (uqPolicy != AMSUQPolicy::RandomUQ &&
         uqPolicy != AMSUQPolicy::DeltaUQ_Max &&