Track Message Type in Execution Graph (#175)

* track message type in execution graph * adding test for tracing by message type * switching to define-based constants for prefixes and keep them in an MPI header * consistent trailing - in constants
faasm · Nov 12, 2021 · 38cc87f · 38cc87f
1 parent 41207d9
commit 38cc87f
Show file tree

Hide file tree

Showing 4 changed files with 88 additions and 9 deletions.
diff --git a/include/faabric/scheduler/MpiWorld.h b/include/faabric/scheduler/MpiWorld.h
@@ -12,6 +12,13 @@
 #include <atomic>
 #include <unordered_map>
 
+// Constants for profiling MPI parameters like number of messages sent or
+// message breakdown by type in the execution graph. Remember to increase the
+// counter if you add another one
+#define NUM_MPI_EXEC_GRAPH_DETAILS 2
+#define MPI_MSG_COUNT_PREFIX "mpi-msgcount-torank"
+#define MPI_MSGTYPE_COUNT_PREFIX "mpi-msgtype-torank"
+
 namespace faabric::scheduler {
 typedef faabric::util::Queue<std::shared_ptr<faabric::MPIMessage>>
   InMemoryMpiQueue;

diff --git a/include/faabric/util/exec_graph.h b/include/faabric/util/exec_graph.h
@@ -7,13 +7,12 @@
 #include <map>
 
 namespace faabric::util::exec_graph {
+
 void addDetail(faabric::Message& msg,
                const std::string& key,
                const std::string& value);
 
 void incrementCounter(faabric::Message& msg,
                       const std::string& key,
                       const int valueToIncrement = 1);
-
-static inline std::string const mpiMsgCountPrefix = "mpi-msgcount-torank-";
 }
diff --git a/src/scheduler/MpiWorld.cpp b/src/scheduler/MpiWorld.cpp
@@ -590,8 +590,15 @@ void MpiWorld::send(int sendRank,
     if (thisRankMsg != nullptr && thisRankMsg->recordexecgraph()) {
         faabric::util::exec_graph::incrementCounter(
           *thisRankMsg,
-          faabric::util::exec_graph::mpiMsgCountPrefix +
-            std::to_string(recvRank));
+          fmt::format("{}-{}", MPI_MSG_COUNT_PREFIX, std::to_string(recvRank)));
+
+        // Work out the message type breakdown
+        faabric::util::exec_graph::incrementCounter(
+          *thisRankMsg,
+          fmt::format("{}-{}-{}",
+                      MPI_MSGTYPE_COUNT_PREFIX,
+                      std::to_string(messageType),
+                      std::to_string(recvRank)));
     }
 }
 

diff --git a/tests/test/scheduler/test_mpi_exec_graph.cpp b/tests/test/scheduler/test_mpi_exec_graph.cpp
@@ -24,7 +24,7 @@ TEST_CASE_METHOD(MpiTestFixture,
 
     int numToSend = 10;
     std::string expectedKey =
-      faabric::util::exec_graph::mpiMsgCountPrefix + std::to_string(rankA2);
+      fmt::format("{}-{}", MPI_MSG_COUNT_PREFIX, std::to_string(rankA2));
 
     for (int i = 0; i < numToSend; i++) {
         world.send(rankA1,
@@ -36,7 +36,7 @@ TEST_CASE_METHOD(MpiTestFixture,
           rankA1, rankA2, BYTES(buffer), MPI_INT, messageData.size(), &status);
     }
 
-    REQUIRE(msg.intexecgraphdetails_size() == 1);
+    REQUIRE(msg.intexecgraphdetails_size() == NUM_MPI_EXEC_GRAPH_DETAILS);
     REQUIRE(msg.execgraphdetails_size() == 0);
     REQUIRE(msg.intexecgraphdetails().count(expectedKey) == 1);
     REQUIRE(msg.intexecgraphdetails().at(expectedKey) == numToSend);
@@ -59,7 +59,7 @@ TEST_CASE_METHOD(MpiTestFixture,
 
     int numToSend = 10;
     std::string expectedKey =
-      faabric::util::exec_graph::mpiMsgCountPrefix + std::to_string(rankA2);
+      fmt::format("{}-{}", MPI_MSG_COUNT_PREFIX, std::to_string(rankA2));
 
     for (int i = 0; i < numToSend; i++) {
         world.send(rankA1,
@@ -124,10 +124,76 @@ TEST_CASE_METHOD(MpiBaseTestFixture,
     }
 
     std::string expectedKey =
-      faabric::util::exec_graph::mpiMsgCountPrefix + std::to_string(rank);
+      fmt::format("{}-{}", MPI_MSG_COUNT_PREFIX, std::to_string(rank));
     REQUIRE(otherMsg.mpirank() == otherRank);
-    REQUIRE(otherMsg.intexecgraphdetails_size() == 1);
     REQUIRE(otherMsg.intexecgraphdetails().count(expectedKey) == 1);
     REQUIRE(otherMsg.intexecgraphdetails().at(expectedKey) == 1);
 }
+
+TEST_CASE_METHOD(MpiTestFixture,
+                 "Test tracing the number of MPI messages by type",
+                 "[util][exec-graph]")
+{
+    msg.set_recordexecgraph(true);
+
+    // Send one message
+    int rankA1 = 0;
+    int rankA2 = 1;
+    MPI_Status status{};
+
+    std::vector<int> messageData = { 0, 1, 2 };
+    auto buffer = new int[messageData.size()];
+
+    std::string expectedKey;
+    int msgCount;
+
+    SECTION("Normal send")
+    {
+        expectedKey = fmt::format("{}-{}-{}",
+                                  MPI_MSGTYPE_COUNT_PREFIX,
+                                  faabric::MPIMessage::NORMAL,
+                                  std::to_string(rankA2));
+        msgCount = 1;
+
+        world.send(rankA1,
+                   rankA2,
+                   BYTES(messageData.data()),
+                   MPI_INT,
+                   messageData.size());
+        world.recv(
+          rankA1, rankA2, BYTES(buffer), MPI_INT, messageData.size(), &status);
+    }
+
+    SECTION("Reduce")
+    {
+        std::vector<int> data(2, 0);
+
+        expectedKey = fmt::format("{}-{}-{}",
+                                  MPI_MSGTYPE_COUNT_PREFIX,
+                                  faabric::MPIMessage::REDUCE,
+                                  std::to_string(rankA2));
+        msgCount = worldSize - 1;
+
+        // Reduce expects to receive a message from all other ranks
+        for (int r = 0; r < worldSize; r++) {
+            if (r != rankA2) {
+                world.reduce(
+                  r, rankA2, BYTES(&data[0]), nullptr, MPI_INT, 1, MPI_SUM);
+            }
+        }
+
+        world.reduce(rankA2,
+                     rankA2,
+                     BYTES(&data[1]),
+                     BYTES(&data[1]),
+                     MPI_INT,
+                     1,
+                     MPI_SUM);
+    }
+
+    REQUIRE(msg.intexecgraphdetails_size() == NUM_MPI_EXEC_GRAPH_DETAILS);
+    REQUIRE(msg.execgraphdetails_size() == 0);
+    REQUIRE(msg.intexecgraphdetails().count(expectedKey) == 1);
+    REQUIRE(msg.intexecgraphdetails().at(expectedKey) == msgCount);
+}
 }