add aggreg callback in aggreg (#255)

Co-authored-by: Thomas Applencourt <[email protected]>
argonne-lcf · Jul 12, 2024 · f2066a9 · f2066a9
1 parent b3c3cee
commit f2066a9
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 5 deletions.
diff --git a/xprof/Makefile.am b/xprof/Makefile.am
@@ -69,7 +69,7 @@ $(BTX_TALLY_GENERATED) &: $(srcdir)/btx_aggreg_model.yaml $(srcdir)/btx_tally_pa
 	$(METABABEL) -u $(srcdir)/btx_aggreg_model.yaml -t SINK -o btx_sink_tally -p btx_tally -c tally --params $(srcdir)/btx_tally_params.yaml
 
 $(BTX_AGGREG_GENERATED) &: $(srcdir)/btx_interval_model.yaml $(srcdir)/btx_aggreg_model.yaml $(srcdir)/btx_aggreg_params.yaml
-	$(METABABEL) -u $(srcdir)/btx_interval_model.yaml -d $(srcdir)/btx_aggreg_model.yaml -t FILTER -o btx_filter_aggreg -p btx_aggreg -c aggreg --params $(srcdir)/btx_aggreg_params.yaml
+	$(METABABEL) -u $(srcdir)/btx_interval_model.yaml,$(srcdir)/btx_aggreg_model.yaml -d $(srcdir)/btx_aggreg_model.yaml -t FILTER -o btx_filter_aggreg -p btx_aggreg -c aggreg --params $(srcdir)/btx_aggreg_params.yaml
 
 $(BTX_STRIPPER_GENERATED) &:
 	$(METABABEL) -t FILTER -o btx_filter_stripper -p btx_stripper -c stripper --enable-callbacks on_downstream

diff --git a/xprof/btx_aggreg.cpp b/xprof/btx_aggreg.cpp
@@ -19,9 +19,7 @@ struct aggreg_s {
 };
 using aggreg_t = struct aggreg_s;
 
-static void initialize_component_callback(void **usr_data) {
-  *usr_data = new aggreg_t;
-}
+static void initialize_component_callback(void **usr_data) { *usr_data = new aggreg_t; }
 
 static void finalize_component_callback(void *usr_data) {
   auto *data = static_cast<aggreg_t *>(usr_data);
@@ -40,7 +38,8 @@ static void finalize_processing_callback(void *btx_handle, void *usr_data) {
   for (const auto &[hptb_function_name, t] : data->traffic) {
     const auto &[hostname, vpid, vtid, backend_id, name, metadata] = hptb_function_name;
     btx_push_message_aggreg_traffic(btx_handle, hostname.c_str(), vpid, vtid, name.c_str(), t.min,
-                                    t.max, t.duration, t.count, (int64_t)backend_id, metadata.c_str());
+                                    t.max, t.duration, t.count, (int64_t)backend_id,
+                                    metadata.c_str());
   }
 
   // Device
@@ -75,6 +74,34 @@ static void device_usr_callback(void *btx_handle, void *usr_data, const char *ho
   data->device[{hostname, vpid, vtid, did, sdid, name, metadata}] += {dur, (bool)err};
 }
 
+// Merge all the aggreg messages in the same trace
+static void aggreg_host_callback(void *btx_handle, void *usr_data, const char *hostname,
+                                 int64_t vpid, uint64_t vtid, const char *name, uint64_t min,
+                                 uint64_t max, uint64_t total, uint64_t count, uint64_t backend,
+                                 uint64_t err) {
+
+  btx_push_message_aggreg_host(btx_handle, hostname, vpid, vtid, name, min, max, total, count,
+                               (int64_t)backend, err);
+}
+
+static void aggreg_device_callback(void *btx_handle, void *usr_data, const char *hostname,
+                                   int64_t vpid, uint64_t vtid, const char *name, uint64_t min,
+                                   uint64_t max, uint64_t total, uint64_t count, uint64_t did,
+                                   uint64_t sdid, const char *metadata) {
+
+  btx_push_message_aggreg_device(btx_handle, hostname, vpid, vtid, name, min, max, total, count,
+                                 did, sdid, metadata);
+}
+
+static void aggreg_traffic_callback(void *btx_handle, void *usr_data, const char *hostname,
+                                    int64_t vpid, uint64_t vtid, const char *name, uint64_t min,
+                                    uint64_t max, uint64_t total, uint64_t count, uint64_t backend,
+                                    const char *metadata) {
+
+  btx_push_message_aggreg_traffic(btx_handle, hostname, vpid, vtid, name, min, max, total, count,
+                                  (int64_t)backend, metadata);
+}
+
 void btx_register_usr_callbacks(void *btx_handle) {
   btx_register_callbacks_initialize_component(btx_handle, &initialize_component_callback);
   btx_register_callbacks_finalize_processing(btx_handle, &finalize_processing_callback);
@@ -83,4 +110,8 @@ void btx_register_usr_callbacks(void *btx_handle) {
   btx_register_callbacks_lttng_host(btx_handle, &host_usr_callback);
   btx_register_callbacks_lttng_traffic(btx_handle, &traffic_usr_callback);
   btx_register_callbacks_lttng_device(btx_handle, &device_usr_callback);
+
+  btx_register_callbacks_aggreg_host(btx_handle, &aggreg_host_callback);
+  btx_register_callbacks_aggreg_device(btx_handle, &aggreg_device_callback);
+  btx_register_callbacks_aggreg_traffic(btx_handle, &aggreg_traffic_callback);
 }