Queue-less remote send/recv

include/faabric/scheduler/FunctionCallApi.h

@@ -4,11 +4,10 @@ namespace faabric::scheduler {
 enum FunctionCalls
 {
     NoFunctionCall = 0,
-    MpiMessage = 1,
-    ExecuteFunctions = 2,
-    Flush = 3,
-    Unregister = 4,
-    GetResources = 5,
-    SetThreadResult = 6,
+    ExecuteFunctions = 1,
+    Flush = 2,
+    Unregister = 3,
+    GetResources = 4,
+    SetThreadResult = 5,
 };
 }

include/faabric/scheduler/FunctionCallClient.h

@@ -19,8 +19,6 @@ std::vector<
   std::pair<std::string, std::shared_ptr<faabric::BatchExecuteRequest>>>
 getBatchRequests();
 
-std::vector<std::pair<std::string, faabric::MPIMessage>> getMPIMessages();
-
 std::vector<std::pair<std::string, faabric::ResponseRequest>>
 getResourceRequests();
 
@@ -44,8 +42,6 @@ class FunctionCallClient : public faabric::transport::MessageEndpointClient
 
     void sendFlush();
 
-    void sendMPIMessage(const std::shared_ptr<faabric::MPIMessage> msg);
-
     faabric::HostResources getResources();
 
     void executeFunctions(

include/faabric/scheduler/FunctionCallServer.h

@@ -2,8 +2,6 @@
 
 #include <faabric/proto/faabric.pb.h>
 #include <faabric/scheduler/FunctionCallApi.h>
-#include <faabric/scheduler/MpiWorld.h>
-#include <faabric/scheduler/MpiWorldRegistry.h>
 #include <faabric/scheduler/Scheduler.h>
 #include <faabric/transport/MessageEndpointClient.h>
 #include <faabric/transport/MessageEndpointServer.h>
@@ -25,8 +23,6 @@ class FunctionCallServer final
 
     /* Function call server API */
 
-    void recvMpiMessage(faabric::transport::Message& body);
-
     void recvFlush(faabric::transport::Message& body);
 
     void recvExecuteFunctions(faabric::transport::Message& body);

include/faabric/scheduler/MpiWorld.h

@@ -6,7 +6,9 @@
 #include <faabric/scheduler/FunctionCallClient.h>
 #include <faabric/scheduler/InMemoryMessageQueue.h>
 #include <faabric/scheduler/MpiThreadPool.h>
-#include <faabric/state/StateKeyValue.h>
+#include <faabric/transport/MpiMessageEndpoint.h>
+#include <faabric/util/logging.h>
+#include <faabric/util/timing.h>
 
 #include <atomic>
 #include <thread>
@@ -15,8 +17,6 @@ namespace faabric::scheduler {
 typedef faabric::util::Queue<std::shared_ptr<faabric::MPIMessage>>
   InMemoryMpiQueue;
 
-std::string getWorldStateKey(int worldId);
-
 class MpiWorld
 {
   public:
@@ -43,8 +43,6 @@ class MpiWorld
 
     void shutdownThreadPool();
 
-    void enqueueMessage(faabric::MPIMessage& msg);
-
     void getCartesianRank(int rank,
                           int maxDims,
                           const int* dims,
@@ -198,23 +196,31 @@ class MpiWorld
     std::string user;
     std::string function;
 
-    std::shared_ptr<state::StateKeyValue> stateKV;
-    std::vector<std::string> rankHosts;
-
-    std::vector<std::shared_ptr<InMemoryMpiQueue>> localQueues;
-
     std::shared_ptr<faabric::scheduler::MpiAsyncThreadPool> threadPool;
     int getMpiThreadPoolSize();
 
     std::vector<int> cartProcsPerDim;
 
-    faabric::scheduler::FunctionCallClient& getFunctionCallClient(
-      const std::string& otherHost);
-
-    void closeThreadLocalClients();
+    /* MPI internal messaging layer */
 
+    // Track at which host each rank lives
+    std::vector<std::string> rankHosts;
     int getIndexForRanks(int sendRank, int recvRank);
 
+    // In-memory queues for local messaging
+    std::vector<std::shared_ptr<InMemoryMpiQueue>> localQueues;
     void initLocalQueues();
+
+    // Rank-to-rank sockets for remote messaging
+    void initRemoteMpiEndpoint(int sendRank, int recvRank);
+    int getMpiPort(int sendRank, int recvRank);
+    void sendRemoteMpiMessage(int sendRank,
+                              int recvRank,
+                              const std::shared_ptr<faabric::MPIMessage>& msg);
+    std::shared_ptr<faabric::MPIMessage> recvRemoteMpiMessage(int sendRank,
+                                                              int recvRank);
+    void closeMpiMessageEndpoints();
+
+    void checkRanksRange(int sendRank, int recvRank);
 };
 }

include/faabric/transport/MessageEndpoint.h

@@ -83,6 +83,8 @@ class RecvMessageEndpoint : public MessageEndpoint
   public:
     RecvMessageEndpoint(int portIn);
 
+    RecvMessageEndpoint(int portIn, const std::string& overrideHost);
+
     void open(MessageContext& context);
 
     void close();

include/faabric/transport/MpiMessageEndpoint.h

@@ -6,8 +6,35 @@
 #include <faabric/transport/macros.h>
 
 namespace faabric::transport {
+/* These two abstract methods are used to broadcast the host-rank mapping at
+ * initialisation time.
+ */
 faabric::MpiHostsToRanksMessage recvMpiHostRankMsg();
 
 void sendMpiHostRankMsg(const std::string& hostIn,
                         const faabric::MpiHostsToRanksMessage msg);
+
+/* This class abstracts the notion of a communication channel between two MPI
+ * ranks. There will always be one rank local to this host, and one remote.
+ * Note that the port is unique per (user, function, sendRank, recvRank) tuple.
+ */
+class MpiMessageEndpoint
+{
+  public:
+    MpiMessageEndpoint(const std::string& hostIn, int portIn);
+
+    MpiMessageEndpoint(const std::string& hostIn,
+                       int portIn,
+                       const std::string& overrideRecvHost);
+
+    void sendMpiMessage(const std::shared_ptr<faabric::MPIMessage>& msg);
+
+    std::shared_ptr<faabric::MPIMessage> recvMpiMessage();
+
+    void close();
+
+  private:
+    SendMessageEndpoint sendMessageEndpoint;
+    RecvMessageEndpoint recvMessageEndpoint;
+};
 }

include/faabric/transport/common.h

@@ -5,8 +5,7 @@
 #define DEFAULT_SNAPSHOT_HOST "0.0.0.0"
 #define STATE_PORT 8003
 #define FUNCTION_CALL_PORT 8004
-#define MPI_MESSAGE_PORT 8005
-#define SNAPSHOT_PORT 8006
+#define SNAPSHOT_PORT 8005
 #define REPLY_PORT_OFFSET 100
 
 #define MPI_PORT 8800

src/scheduler/FunctionCallClient.cpp

@@ -19,8 +19,6 @@ static std::vector<
   std::pair<std::string, std::shared_ptr<faabric::BatchExecuteRequest>>>
   batchMessages;
 
-static std::vector<std::pair<std::string, faabric::MPIMessage>> mpiMessages;
-
 static std::vector<std::pair<std::string, faabric::ResponseRequest>>
   resourceRequests;
 
@@ -48,11 +46,6 @@ getBatchRequests()
     return batchMessages;
 }
 
-std::vector<std::pair<std::string, faabric::MPIMessage>> getMPIMessages()
-{
-    return mpiMessages;
-}
-
 std::vector<std::pair<std::string, faabric::ResponseRequest>>
 getResourceRequests()
 {
@@ -74,7 +67,6 @@ void clearMockRequests()
 {
     functionCalls.clear();
     batchMessages.clear();
-    mpiMessages.clear();
     resourceRequests.clear();
     unregisterRequests.clear();
 
@@ -113,17 +105,6 @@ void FunctionCallClient::sendFlush()
     }
 }
 
-void FunctionCallClient::sendMPIMessage(
-  const std::shared_ptr<faabric::MPIMessage> msg)
-{
-    if (faabric::util::isMockMode()) {
-        faabric::util::UniqueLock lock(mockMutex);
-        mpiMessages.emplace_back(host, *msg);
-    } else {
-        SEND_MESSAGE_PTR(faabric::scheduler::FunctionCalls::MpiMessage, msg);
-    }
-}
-
 faabric::HostResources FunctionCallClient::getResources()
 {
     faabric::ResponseRequest request;

src/scheduler/FunctionCallServer.cpp

@@ -27,9 +27,6 @@ void FunctionCallServer::doRecv(faabric::transport::Message& header,
     assert(header.size() == sizeof(uint8_t));
     uint8_t call = static_cast<uint8_t>(*header.data());
     switch (call) {
-        case faabric::scheduler::FunctionCalls::MpiMessage:
-            this->recvMpiMessage(body);
-            break;
         case faabric::scheduler::FunctionCalls::Flush:
             this->recvFlush(body);
             break;
@@ -48,15 +45,6 @@ void FunctionCallServer::doRecv(faabric::transport::Message& header,
     }
 }
 
-void FunctionCallServer::recvMpiMessage(faabric::transport::Message& body)
-{
-    PARSE_MSG(faabric::MPIMessage, body.data(), body.size())
-
-    MpiWorldRegistry& registry = getMpiWorldRegistry();
-    MpiWorld& world = registry.getWorld(msg.worldid());
-    world.enqueueMessage(msg);
-}
-
 void FunctionCallServer::recvFlush(faabric::transport::Message& body)
 {
     PARSE_MSG(faabric::ResponseRequest, body.data(), body.size());