We now use the 2 bytes for padding in AMSMessage for sending the message ID from AMSlib

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

Author: lpottier

benchmarks/ams_bench_db.cpp

@@ -294,10 +294,6 @@ int main(int argc, char **argv)
     return -1;
   }
 
-  if (dbType != AMSDBType::AMS_RMQ) {
-    AMSConfigureFSDatabase(dbType, db_config);
-  }
-
   // -------------------------------------------------------------------------
   // AMS allocators setup
   // -------------------------------------------------------------------------

src/AMSWorkflow/ams/monitor.py

@@ -206,7 +206,7 @@ def reset(cls):
 
     def start_monitor(self, *args, **kwargs):
         self.start_time = time.time_ns()
-        self.internal_ts = str(time.time_ns())
+        self.internal_ts = time.time_ns()
 
     def stop_monitor(self):
         end = time.time_ns()
@@ -233,7 +233,7 @@ def __call__(self, func: Callable):
         """
 
         def wrapper(*args, **kwargs):
-            ts = str(time.time_ns())
+            ts = time.time_ns()
             start = time.time_ns()
 
             if self.use_arrays:
@@ -364,11 +364,11 @@ def _filter(self, data: dict, keys: List[str]) -> dict:
             return data
         return {k: v for k, v in data.items() if k in keys}
 
-    def _get_ts(self, class_name: str, func: str) -> str:
+    def _get_ts(self, class_name: str, func: str) -> int:
         """
         Return initial timestamp for a given monitored function.
         """
-        ts = str(time.time_ns())
+        ts = time.time_ns()
         if class_name not in AMSMonitor._stats or func not in AMSMonitor._stats[class_name]:
             return ts
 

src/AMSWorkflow/ams/rmq.py

@@ -34,12 +34,17 @@ def __init__(self, body: str):
         self.domain_names = []
         self.input_dim = None
         self.output_dim = None
+        self.message_id = None
 
     def __str__(self):
         dt = "float" if self.dtype_byte == 4 else 8
         if not self.dtype_byte:
             dt = None
-        return f"AMSMessage(domain={self.domain_names}, #mpi={self.mpi_rank}, num_elements={self.num_elements}, datatype={dt}, input_dim={self.input_dim}, output_dim={self.output_dim})"
+        s = f"AMSMessage(mpi_rank={self.mpi_rank}, "
+        s += f"message_id={self.message_id}, domain={self.domain_names}, "
+        s += f"num_elements={self.num_elements}, datatype={dt}, "
+        s += f"input_dim={self.input_dim}, output_dim={self.output_dim})"
+        return s
 
     def __repr__(self):
         return self.__str__()
@@ -55,9 +60,9 @@ def header_format(self) -> str:
         - 4 bytes are the number of elements in the message. Limit max: 2^32 - 1
         - 2 bytes are the input dimension. Limit max: 65535
         - 2 bytes are the output dimension. Limit max: 65535
-        - 2 bytes are for aligning memory to 8
+        - 2 bytes are the message ID given by AMSlib (local to each MPI rank). Limit max: 65535
 
-            |_Header_|_Datatype_|_Rank_|_DomainSize_|_#elems_|_InDim_|_OutDim_|_Pad_|_DomainName_|.Real_Data.|
+            |_Header_|_Datatype_|_Rank_|_DomainSize_|_#elems_|_InDim_|_OutDim_|_MessageID_|_DomainName_|.Real_Data.|
 
         Then the data starts at byte 16 with the domain name, then the real data and
         is structured as pairs of input/outputs. Let K be the total number of elements,
@@ -75,7 +80,7 @@ def endianness(self) -> str:
         """
         return "="
 
-    def encode(self, num_elem: int, domain_name: str, input_dim: int, output_dim: int, dtype_byte: int = 4) -> bytes:
+    def encode(self, num_elem: int, domain_name: str, input_dim: int, output_dim: int, message_id: int , dtype_byte: int = 4) -> bytes:
         """
         For debugging and testing purposes, this function encode a message identical to what AMS would send
         """
@@ -87,8 +92,7 @@ def encode(self, num_elem: int, domain_name: str, input_dim: int, output_dim: in
         data = np.random.rand(num_elem * (input_dim + output_dim))
         domain_name_size = len(domain_name)
         domain_name = bytes(domain_name, "utf-8")
-        padding = 0
-        header_content = (hsize, dtype_byte, mpi_rank, domain_name_size, data.size, input_dim, output_dim, padding)
+        header_content = (hsize, dtype_byte, mpi_rank, domain_name_size, data.size, input_dim, output_dim, message_id)
         # float or double
         msg_format = f"{header_format}{domain_name_size}s{data.size}{dt}"
         return struct.pack(msg_format, *header_content, domain_name, *data)
@@ -113,7 +117,7 @@ def _parse_header(self, body: str) -> dict:
             res["num_element"],
             res["input_dim"],
             res["output_dim"],
-            res["padding"],
+            res["message_id"],
         ) = struct.unpack(fmt, body[:hsize])
         assert hsize == res["hsize"]
         assert res["datatype"] in [4, 8]
@@ -134,6 +138,7 @@ def _parse_header(self, body: str) -> dict:
         self.domain_name_size = int(res["domain_size"])
         self.input_dim = int(res["input_dim"])
         self.output_dim = int(res["output_dim"])
+        self.message_id = int(res["message_id"])
 
         return res
 

src/AMSWorkflow/ams/stage.py

@@ -215,7 +215,7 @@ def __init__(self, o_queue, loader, pattern):
         self.loader = loader
         self.datasize = 0
 
-    @AMSMonitor(record=["datasize"])
+    @AMSMonitor(array=["msgs"], record=["datasize"])
     def __call__(self):
         """
         Busy loop of reading all files matching the pattern and creating
@@ -226,6 +226,7 @@ def __call__(self):
         start = time.time()
         files = list(glob.glob(self.pattern))
         for fn in files:
+            start_time_fs = time.time_ns()
             with self.loader(fn) as fd:
                 domain_name, input_data, output_data = fd.load()
                 print("Domain Name is", domain_name)
@@ -237,14 +238,28 @@ def __call__(self):
                 for j, (i, o) in enumerate(zip(input_batches, output_batches)):
                     self.o_queue.put(QueueMessage(MessageType.Process, DataBlob(i, o, domain_name)))
                 self.datasize += input_data.nbytes + output_data.nbytes
+
+                end_time_fs = time.time_ns()
+                msg = {
+                    "file": fn,
+                    "domain_name": domain_name,
+                    "row_size": row_size,
+                    "batch_size": BATCH_SIZE,
+                    "rows_per_batch": rows_per_batch,
+                    "num_batches": num_batches,
+                    "size_bytes": input_data.nbytes + output_data.nbytes,
+                    "process_time_ns": end_time_fs - start_time_fs,
+                }
+                # Msgs is the array (list) we push to (managed by AMSMonitor)
+                msgs.append(msg)
+
             print(f"Sending Delete Message Type {self.__class__.__name__}")
             self.o_queue.put(QueueMessage(MessageType.Delete, fn))
         self.o_queue.put(QueueMessage(MessageType.Terminate, None))
 
         end = time.time()
         print(f"Spend {end - start} at {self.__class__.__name__}")
 
-
 class RMQDomainDataLoaderTask(Task):
     """
     A RMQDomainDataLoaderTask consumes 'AMSMessages' from RabbitMQ bundles the data of
@@ -283,8 +298,6 @@ def __init__(
         self.orig_sig_handlers = {}
         self.policy = policy
 
-        # Counter that get incremented when we receive a message
-        self.internal_msg_cnt = 0
 
         # Signals can only be used within the main thread
         if self.policy != "thread":
@@ -340,12 +353,10 @@ def callback_message(self, ch, basic_deliver, properties, body):
 
         self.total_time += (end_time - start_time)
         # TODO: Improve the code to manage potentially multiple messages per AMSMessage
-        # TODO: Right now the ID is not encoded in the AMSMessage by AMSlib
-        #       If order of messages matters we might have to encode it
         msg = {
-            "id": self.internal_msg_cnt,
             "delivery_tag": basic_deliver.delivery_tag,
             "mpi_rank": msg.mpi_rank,
+            "message_id": msg.message_id,
             "domain_name": domain_name,
             "num_elements": msg.num_elements,
             "input_dim": msg.input_dim,
@@ -356,7 +367,6 @@ def callback_message(self, ch, basic_deliver, properties, body):
         }
         # Msgs is the array (list) we push to (managed by AMSMonitor)
         msgs.append(msg)
-        self.internal_msg_cnt += 1
 
     def signal_wrapper(self, name, pid):
         def handler(signum, frame):
@@ -621,11 +631,15 @@ def handler(signum, frame):
             # This should only trigger RMQDomainDataLoaderTask
 
             # TODO: I don't like this system to shutdown the pipeline on demand
-            # It's extremely easy to mess thing up with signals.. and it's
+            # It's extremely easy to mess things up with signals.. and it's
             # not a robust solution (if a task is not managing correctly SIGINT
             # the pipeline can explode)
             for e in self._executors:
-                os.kill(e.pid, signal.SIGINT)
+                if e is not None:
+                    try:
+                        os.kill(e.pid, signal.SIGINT)
+                    except Exception as e:
+                        print(f"Error: {e}")
             self.release_signals()
         return handler
 

src/AMSWorkflow/ams_wf/AMSDBStage.py

@@ -93,7 +93,7 @@ def main():
     print(f"End to End time spend : {end - start}")
 
     if args.output_json is not None:
-        print(f"{AMSMonitor.info()}")
+        # print(f"{AMSMonitor.info()}")
         # Output profiling output to JSON (just as an example)
         AMSMonitor.json(args.output_json)