Periodically Check For Migration Opportunities

include/faabric/scheduler/FunctionCallApi.h

@@ -7,6 +7,7 @@ enum FunctionCalls
     ExecuteFunctions = 1,
     Flush = 2,
     Unregister = 3,
-    GetResources = 4
+    GetResources = 4,
+    AddPendingMigration = 5
 };
 }

include/faabric/scheduler/FunctionCallClient.h

@@ -22,6 +22,9 @@ getBatchRequests();
 std::vector<std::pair<std::string, faabric::EmptyRequest>>
 getResourceRequests();
 
+std::vector<std::pair<std::string, std::shared_ptr<faabric::PendingMigrations>>>
+getAddPendingMigrationRequests();
+
 std::vector<std::pair<std::string, faabric::UnregisterRequest>>
 getUnregisterRequests();
 
@@ -42,6 +45,9 @@ class FunctionCallClient : public faabric::transport::MessageEndpointClient
 
     faabric::HostResources getResources();
 
+    void sendAddPendingMigration(
+      std::shared_ptr<faabric::PendingMigrations> req);
+
     void executeFunctions(std::shared_ptr<faabric::BatchExecuteRequest> req);
 
     void unregister(faabric::UnregisterRequest& req);

include/faabric/scheduler/FunctionCallServer.h

@@ -29,6 +29,10 @@ class FunctionCallServer final
       const uint8_t* buffer,
       size_t bufferSize);
 
+    std::unique_ptr<google::protobuf::Message> recvAddPendingMigration(
+      const uint8_t* buffer,
+      size_t bufferSize);
+
     void recvExecuteFunctions(const uint8_t* buffer, size_t bufferSize);
 
     void recvUnregister(const uint8_t* buffer, size_t bufferSize);

include/faabric/scheduler/FunctionMigrationThread.h

@@ -0,0 +1,26 @@
+#pragma once
+
+#include <condition_variable>
+#include <mutex>
+#include <thread>
+
+namespace faabric::scheduler {
+class FunctionMigrationThread
+{
+  public:
+    // Start a background thread that, every wake up period, will check if there
+    // are migration opportunities for in-flight apps that have opted in to
+    // being checked for migrations.
+    void start(int wakeUpPeriodSecondsIn);
+
+    void stop();
+
+    int wakeUpPeriodSeconds;
+
+  private:
+    std::unique_ptr<std::thread> workThread = nullptr;
+    std::mutex mx;
+    std::condition_variable mustStopCv;
+    std::atomic<bool> isShutdown;
+};
+}

include/faabric/scheduler/MpiWorld.h

@@ -31,6 +31,9 @@ std::vector<faabric::MpiHostsToRanksMessage> getMpiHostsToRanksMessages();
 std::vector<std::shared_ptr<faabric::MPIMessage>> getMpiMockedMessages(
   int sendRank);
 
+std::vector<std::shared_ptr<faabric::PendingMigrations>>
+getMpiMockedPendingMigrations();
+
 typedef faabric::util::FixedCapacityQueue<std::shared_ptr<faabric::MPIMessage>>
   InMemoryMpiQueue;
 
@@ -271,6 +274,16 @@ class MpiWorld
                                                              int recvRank,
                                                              int batchSize = 0);
 
+    /* Function Migration */
+
+    void tryMigrate(int thisRank);
+
+    void prepareMigration(
+      int thisRank,
+      std::shared_ptr<faabric::PendingMigrations> pendingMigrations);
+
+    void finishMigration(int thisRank);
+
     /* Helper methods */
 
     void checkRanksRange(int sendRank, int recvRank);

include/faabric/scheduler/Scheduler.h

@@ -3,6 +3,7 @@
 #include <faabric/proto/faabric.pb.h>
 #include <faabric/scheduler/ExecGraph.h>
 #include <faabric/scheduler/FunctionCallClient.h>
+#include <faabric/scheduler/FunctionMigrationThread.h>
 #include <faabric/scheduler/InMemoryMessageQueue.h>
 #include <faabric/snapshot/SnapshotClient.h>
 #include <faabric/snapshot/SnapshotRegistry.h>
@@ -22,6 +23,10 @@
 
 namespace faabric::scheduler {
 
+typedef std::pair<std::shared_ptr<BatchExecuteRequest>,
+                  std::shared_ptr<faabric::util::SchedulingDecision>>
+  InFlightPair;
+
 class Scheduler;
 
 Scheduler& getScheduler();
@@ -77,6 +82,9 @@ class Executor
       faabric::Message& msg,
       bool createIfNotExists = false);
 
+    void doMigration(
+      std::shared_ptr<faabric::PendingMigrations> pendingMigrations);
+
   protected:
     virtual void restore(const std::string& snapshotKey);
 
@@ -96,6 +104,8 @@ class Executor
 
     uint32_t threadPoolSize = 0;
 
+    void migrateFunction(const faabric::Message& msg, const std::string& host);
+
   private:
     std::atomic<bool> claimed = false;
 
@@ -220,6 +230,18 @@ class Scheduler
 
     ExecGraph getFunctionExecGraph(unsigned int msgId);
 
+    // ----------------------------------
+    // Function Migration
+    // ----------------------------------
+    void checkForMigrationOpportunities();
+
+    std::shared_ptr<faabric::PendingMigrations> canAppBeMigrated(
+      uint32_t appId);
+
+    void addPendingMigration(std::shared_ptr<faabric::PendingMigrations> msg);
+
+    void removePendingMigration(uint32_t appId);
+
   private:
     std::string thisHost;
 
@@ -290,6 +312,20 @@ class Scheduler
 
     // ---- Point-to-point ----
     faabric::transport::PointToPointBroker& broker;
+
+    // ---- Function migration ----
+    FunctionMigrationThread functionMigrationThread;
+    std::unordered_map<uint32_t, InFlightPair> inFlightRequests;
+    std::unordered_map<uint32_t, std::shared_ptr<faabric::PendingMigrations>>
+      pendingMigrations;
+
+    std::vector<std::shared_ptr<faabric::PendingMigrations>>
+    doCheckForMigrationOpportunities(
+      faabric::util::MigrationStrategy migrationStrategy =
+        faabric::util::MigrationStrategy::BIN_PACK);
+
+    void broadcastAddPendingMigration(
+      std::shared_ptr<faabric::PendingMigrations> pendingMigrations);
 };
 
 }

include/faabric/util/message.h

@@ -0,0 +1,7 @@
+#pragma once
+
+#include <faabric/proto/faabric.pb.h>
+
+namespace faabric::util {
+void copyMessage(const faabric::Message* src, faabric::Message* dst);
+}

include/faabric/util/scheduling.h

@@ -54,4 +54,16 @@ enum SchedulingTopologyHint
     FORCE_LOCAL,
     NEVER_ALONE
 };
+
+// Migration strategies help the scheduler decide wether the scheduling decision
+// for a batch request could be changed with the new set of available resources.
+// - BIN_PACK: sort hosts by the number of functions from the batch they are
+//             running. Bin-pack batches in increasing order to hosts in
+//             decreasing order.
+// - EMPTY_HOSTS: pack batches in increasing order to empty hosts.
+enum MigrationStrategy
+{
+    BIN_PACK,
+    EMPTY_HOSTS
+};
 }

src/proto/faabric.proto

@@ -164,6 +164,9 @@ message Message {
     bool recordExecGraph = 41;
     map<string, int32> intExecGraphDetails = 42;
     map<string, string> execGraphDetails = 43;
+
+    // Function migration
+    int32 migrationCheckPeriod = 44;
 }
 
 // ---------------------------------------------
@@ -242,3 +245,20 @@ message PointToPointMappings {
 
     repeated PointToPointMapping mappings = 3;
 }
+
+// ---------------------------------------------
+// FUNCTION MIGRATIONS
+// ---------------------------------------------
+
+message PendingMigrations {
+    int32 appId = 1;
+    int32 groupId = 2;
+
+    message PendingMigration {
+        Message msg = 1;
+        string srcHost = 2;
+        string dstHost = 3;
+    }
+
+    repeated PendingMigration migrations = 3;
+}

src/scheduler/CMakeLists.txt

@@ -4,6 +4,7 @@ faabric_lib(scheduler
     Executor.cpp
     FunctionCallClient.cpp
     FunctionCallServer.cpp
+    FunctionMigrationThread.cpp
     MpiContext.cpp
     MpiMessageBuffer.cpp
     MpiWorld.cpp

src/scheduler/Executor.cpp

@@ -735,6 +735,29 @@ void Executor::releaseClaim()
     claimed.store(false);
 }
 
+void Executor::doMigration(
+  std::shared_ptr<faabric::PendingMigrations> pendingMigrations)
+{
+    for (int i = 0; i < pendingMigrations->migrations_size(); i++) {
+        auto m = pendingMigrations->mutable_migrations()->at(i);
+        if (m.msg().id() == boundMessage.id()) {
+            migrateFunction(m.msg(), m.dsthost());
+            // TODO: terminate current executing thread
+        }
+    }
+}
+
+void Executor::migrateFunction(const faabric::Message& msg,
+                               const std::string& host)
+{
+    SPDLOG_INFO("Executor received request to migrate message {} from host {}"
+                " to host {}",
+                msg.id(),
+                msg.executedhost(),
+                host);
+    SPDLOG_ERROR("Executor::migrate() not implemented");
+}
+
 // ------------------------------------------
 // HOOKS
 // ------------------------------------------

src/scheduler/FunctionCallClient.cpp

@@ -28,6 +28,10 @@ static std::unordered_map<std::string,
                           faabric::util::Queue<faabric::HostResources>>
   queuedResourceResponses;
 
+static std::vector<
+  std::pair<std::string, std::shared_ptr<faabric::PendingMigrations>>>
+  addPendingMigrationRequests;
+
 static std::vector<std::pair<std::string, faabric::UnregisterRequest>>
   unregisterRequests;
 
@@ -57,6 +61,13 @@ std::vector<std::pair<std::string, faabric::EmptyRequest>> getResourceRequests()
     return resourceRequests;
 }
 
+std::vector<std::pair<std::string, std::shared_ptr<faabric::PendingMigrations>>>
+getAddPendingMigrationRequests()
+{
+    faabric::util::UniqueLock lock(mockMutex);
+    return addPendingMigrationRequests;
+}
+
 std::vector<std::pair<std::string, faabric::UnregisterRequest>>
 getUnregisterRequests()
 {
@@ -76,6 +87,7 @@ void clearMockRequests()
     functionCalls.clear();
     batchMessages.clear();
     resourceRequests.clear();
+    addPendingMigrationRequests.clear();
     unregisterRequests.clear();
 
     for (auto& p : queuedResourceResponses) {
@@ -128,6 +140,22 @@ faabric::HostResources FunctionCallClient::getResources()
     return response;
 }
 
+void FunctionCallClient::sendAddPendingMigration(
+  std::shared_ptr<PendingMigrations> req)
+{
+    faabric::PendingMigrations request;
+    faabric::EmptyResponse response;
+
+    if (faabric::util::isMockMode()) {
+        faabric::util::UniqueLock lock(mockMutex);
+        addPendingMigrationRequests.emplace_back(host, req);
+    } else {
+        syncSend(faabric::scheduler::FunctionCalls::AddPendingMigration,
+                 req.get(),
+                 &response);
+    }
+}
+
 void FunctionCallClient::executeFunctions(
   const std::shared_ptr<faabric::BatchExecuteRequest> req)
 {

src/scheduler/FunctionCallServer.cpp

@@ -49,6 +49,9 @@ std::unique_ptr<google::protobuf::Message> FunctionCallServer::doSyncRecv(
         case faabric::scheduler::FunctionCalls::GetResources: {
             return recvGetResources(buffer, bufferSize);
         }
+        case faabric::scheduler::FunctionCalls::AddPendingMigration: {
+            return recvAddPendingMigration(buffer, bufferSize);
+        }
         default: {
             throw std::runtime_error(
               fmt::format("Unrecognized sync call header: {}", header));
@@ -100,4 +103,17 @@ std::unique_ptr<google::protobuf::Message> FunctionCallServer::recvGetResources(
       scheduler.getThisHostResources());
     return response;
 }
+
+std::unique_ptr<google::protobuf::Message>
+FunctionCallServer::recvAddPendingMigration(const uint8_t* buffer,
+                                            size_t bufferSize)
+{
+    PARSE_MSG(faabric::PendingMigrations, buffer, bufferSize);
+
+    auto msgPtr = std::make_shared<faabric::PendingMigrations>(msg);
+
+    scheduler.addPendingMigration(msgPtr);
+
+    return std::make_unique<faabric::EmptyResponse>();
+}
 }