Fixes for MPI message graph (#173)

* fixes for exec graph to work * add regression test
faasm · Nov 9, 2021 · f7a910b · f7a910b
1 parent 832aafe
commit f7a910b
Show file tree

Hide file tree

Showing 5 changed files with 75 additions and 5 deletions.
diff --git a/include/faabric/scheduler/MpiWorld.h b/include/faabric/scheduler/MpiWorld.h
@@ -184,6 +184,10 @@ class MpiWorld
 
     std::vector<bool> getInitedUMB();
 
+    /* Profiling */
+
+    void setMsgForRank(faabric::Message& msg);
+
   private:
     int id = -1;
     int size = -1;

diff --git a/src/scheduler/MpiWorld.cpp b/src/scheduler/MpiWorld.cpp
@@ -300,7 +300,6 @@ void MpiWorld::initialiseFromMsg(faabric::Message& msg)
     user = msg.user();
     function = msg.function();
     size = msg.mpiworldsize();
-    thisRankMsg = &msg;
 
     // Block until we receive
     faabric::MpiHostsToRanksMessage hostRankMsg = recvMpiHostRankMsg();
@@ -324,6 +323,11 @@ void MpiWorld::initialiseFromMsg(faabric::Message& msg)
     initLocalQueues();
 }
 
+void MpiWorld::setMsgForRank(faabric::Message& msg)
+{
+    thisRankMsg = &msg;
+}
+
 std::string MpiWorld::getHostForRank(int rank)
 {
     assert(rankHosts.size() == size);

diff --git a/src/scheduler/MpiWorldRegistry.cpp b/src/scheduler/MpiWorldRegistry.cpp
@@ -51,7 +51,11 @@ MpiWorld& MpiWorldRegistry::getOrInitialiseWorld(faabric::Message& msg)
 
     {
         faabric::util::SharedLock lock(registryMutex);
-        return worldMap[worldId];
+        MpiWorld& world = worldMap[worldId];
+        if (msg.recordexecgraph()) {
+            world.setMsgForRank(msg);
+        }
+        return world;
     }
 }
 

diff --git a/src/util/json.cpp b/src/util/json.cpp
@@ -218,9 +218,11 @@ std::string messageToJson(const faabric::Message& msg)
             }
 
             std::string out = ss.str();
-            d.AddMember("int_exec_graph_detail",
-                        Value(out.c_str(), out.size()).Move(),
-                        a);
+
+            // Need to create a value (instead of move) as the string's scope
+            // is smaller than the document's one
+            Value value = Value(out.c_str(), out.size(), a);
+            d.AddMember("int_exec_graph_detail", value, a);
         }
     }
 

diff --git a/tests/test/scheduler/test_mpi_exec_graph.cpp b/tests/test/scheduler/test_mpi_exec_graph.cpp
@@ -4,6 +4,7 @@
 
 #include <faabric/scheduler/MpiWorld.h>
 #include <faabric/util/exec_graph.h>
+#include <faabric/util/json.h>
 #include <faabric/util/macros.h>
 
 namespace tests {
@@ -74,4 +75,59 @@ TEST_CASE_METHOD(MpiTestFixture,
     REQUIRE(msg.intexecgraphdetails_size() == 0);
     REQUIRE(msg.execgraphdetails_size() == 0);
 }
+
+TEST_CASE_METHOD(MpiBaseTestFixture,
+                 "Test different threads populate the graph",
+                 "[util][exec-graph]")
+{
+    int rank = 0;
+    int otherRank = 1;
+    int worldSize = 2;
+    int worldId = 123;
+
+    faabric::Message msg = faabric::util::messageFactory("mpi", "hellompi");
+    msg.set_ismpi(true);
+    msg.set_recordexecgraph(true);
+    msg.set_mpiworldsize(worldSize);
+    msg.set_mpiworldid(worldId);
+
+    faabric::Message otherMsg = msg;
+    otherMsg.set_mpirank(otherRank);
+    msg.set_mpirank(rank);
+
+    faabric::scheduler::MpiWorld& thisWorld =
+      faabric::scheduler::getMpiWorldRegistry().createWorld(msg, worldId);
+
+    std::vector<int> messageData = { 0, 1, 2 };
+    auto buffer = new int[messageData.size()];
+    std::thread otherWorldThread([&messageData, &otherMsg, rank, otherRank] {
+        faabric::scheduler::MpiWorld& otherWorld =
+          faabric::scheduler::getMpiWorldRegistry().getOrInitialiseWorld(
+            otherMsg);
+
+        otherWorld.send(otherRank,
+                        rank,
+                        BYTES(messageData.data()),
+                        MPI_INT,
+                        messageData.size());
+
+        otherWorld.destroy();
+    });
+
+    thisWorld.recv(
+      otherRank, rank, BYTES(buffer), MPI_INT, messageData.size(), nullptr);
+
+    thisWorld.destroy();
+
+    if (otherWorldThread.joinable()) {
+        otherWorldThread.join();
+    }
+
+    std::string expectedKey =
+      faabric::util::exec_graph::mpiMsgCountPrefix + 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);
+}
 }