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,
+    PendingMigrations = 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>>>
+getPendingMigrationsRequests();
+
 std::vector<std::pair<std::string, faabric::UnregisterRequest>>
 getUnregisterRequests();
 
@@ -42,6 +45,8 @@ class FunctionCallClient : public faabric::transport::MessageEndpointClient
 
     faabric::HostResources getResources();
 
+    void sendPendingMigrations(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> recvPendingMigrations(
+      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 {
+// 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.
+class FunctionMigrationThread
+{
+  public:
+    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

@@ -207,6 +207,12 @@ class MpiWorld
 
     void setMsgForRank(faabric::Message& msg);
 
+    /* Function Migration */
+
+    void prepareMigration(
+      int thisRank,
+      std::shared_ptr<faabric::PendingMigrations> pendingMigrations);
+
   private:
     int id = -1;
     int size = -1;

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,18 @@
 
 namespace faabric::scheduler {
 
+class ExecutorMigratedException : public faabric::util::FaabricException
+{
+  public:
+    explicit ExecutorMigratedException(std::string message)
+      : FaabricException(std::move(message))
+    {}
+};
+
+typedef std::pair<std::shared_ptr<BatchExecuteRequest>,
+                  std::shared_ptr<faabric::util::SchedulingDecision>>
+  InFlightPair;
+
 class Scheduler;
 
 Scheduler& getScheduler();
@@ -77,13 +90,13 @@ class Executor
       faabric::Message& msg,
       bool createIfNotExists = false);
 
+    virtual std::span<uint8_t> getMemoryView();
+
   protected:
     virtual void restore(const std::string& snapshotKey);
 
     virtual void postFinish();
 
-    virtual std::span<uint8_t> getMemoryView();
-
     virtual void setMemorySize(size_t newSize);
 
     faabric::Message boundMessage;
@@ -96,6 +109,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 +235,27 @@ class Scheduler
 
     ExecGraph getFunctionExecGraph(unsigned int msgId);
 
+    // ----------------------------------
+    // Function Migration
+    // ----------------------------------
+    void checkForMigrationOpportunities();
+
+    std::shared_ptr<faabric::PendingMigrations> getPendingAppMigrations(
+      uint32_t appId);
+
+    void addPendingMigration(std::shared_ptr<faabric::PendingMigrations> msg);
+
+    void removePendingMigration(uint32_t appId);
+
+    // ----------------------------------
+    // Clients
+    // ----------------------------------
+    faabric::scheduler::FunctionCallClient& getFunctionCallClient(
+      const std::string& otherHost);
+
+    faabric::snapshot::SnapshotClient& getSnapshotClient(
+      const std::string& otherHost);
+
   private:
     std::string thisHost;
 
@@ -244,13 +280,6 @@ class Scheduler
 
     std::mutex localResultsMutex;
 
-    // ---- Clients ----
-    faabric::scheduler::FunctionCallClient& getFunctionCallClient(
-      const std::string& otherHost);
-
-    faabric::snapshot::SnapshotClient& getSnapshotClient(
-      const std::string& otherHost);
-
     // ---- Host resources and hosts ----
     faabric::HostResources thisHostResources;
     std::atomic<int32_t> thisHostUsedSlots = 0;
@@ -290,6 +319,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 broadcastPendingMigrations(
+      std::shared_ptr<faabric::PendingMigrations> pendingMigrations);
 };
 
 }

include/faabric/util/func.h

@@ -7,6 +7,14 @@
 
 namespace faabric::util {
 
+class FunctionMigratedException : public faabric::util::FaabricException
+{
+  public:
+    explicit FunctionMigratedException(std::string message)
+      : FaabricException(std::move(message))
+    {}
+};
+
 std::string funcToString(const faabric::Message& msg, bool includeId);
 
 std::string funcToString(

include/faabric/util/scheduling.h

@@ -46,12 +46,29 @@ class SchedulingDecision
 // requests in a batch.
 //  - NORMAL: bin-packs requests to slots in hosts starting from the master
 //            host, and overloadds the master if it runs out of resources.
+//  - FORCE_LOCAL: force local execution irrespective of the available
+//                 resources.
 //  - NEVER_ALONE: never allocates a single (non-master) request to a host
 //                 without other requests of the batch.
+//  - UNDERFULL: schedule up to 50% of the master hosts' capacity to force
+//               migration opportunities to appear.
 enum SchedulingTopologyHint
 {
     NORMAL,
     FORCE_LOCAL,
-    NEVER_ALONE
+    NEVER_ALONE,
+    UNDERFULL,
+};
+
+// 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

@@ -25,6 +25,7 @@ message BatchExecuteRequest {
         FUNCTIONS = 0;
         THREADS = 1;
         PROCESSES = 2;
+        MIGRATION = 3;
     }
 
     BatchExecuteType type = 2;
@@ -164,6 +165,9 @@ message Message {
     bool recordExecGraph = 41;
     map<string, int32> intExecGraphDetails = 42;
     map<string, string> execGraphDetails = 43;
+
+    // Function migration
+    int32 migrationCheckPeriod = 44;
 }
 
 // ---------------------------------------------
@@ -242,3 +246,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

@@ -1,3 +1,4 @@
+#include "faabric/scheduler/MpiWorldRegistry.h"
 #include <faabric/proto/faabric.pb.h>
 #include <faabric/scheduler/Scheduler.h>
 #include <faabric/snapshot/SnapshotRegistry.h>
@@ -528,9 +529,29 @@ void Executor::threadPoolThread(int threadPoolIdx)
 
         // Execute the task
         int32_t returnValue;
+        bool migrated = false;
         try {
             returnValue =
               executeTask(threadPoolIdx, task.messageIndex, task.req);
+        } catch (const faabric::util::FunctionMigratedException& ex) {
+            SPDLOG_DEBUG(
+              "Task {} migrated, shutting down executor {}", msg.id(), id);
+
+            // Note that when a task has been migrated, we need to perform all
+            // the normal executor shutdown, but we must NOT set the result for
+            // the call.
+            migrated = true;
+            selfShutdown = true;
+            returnValue = -99;
+
+            // MPI migration
+            if (msg.ismpi()) {
+                // TODO - delete the pending migration
+                auto& mpiWorld =
+                  faabric::scheduler::getMpiWorldRegistry().getWorld(
+                    msg.mpiworldid());
+                mpiWorld.destroy();
+            }
         } catch (const std::exception& ex) {
             returnValue = 1;
 
@@ -667,6 +688,12 @@ void Executor::threadPoolThread(int threadPoolIdx)
         // executor.
         sch.vacateSlot();
 
+        // If the function has been migrated, we drop out here and shut down the
+        // executor
+        if (migrated) {
+            break;
+        }
+
         // Finally set the result of the task, this will allow anything
         // waiting on its result to continue execution, therefore must be
         // done once the executor has been reset, otherwise the executor may

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>>>
+  pendingMigrationsRequests;
+
 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>>>
+getPendingMigrationsRequests()
+{
+    faabric::util::UniqueLock lock(mockMutex);
+    return pendingMigrationsRequests;
+}
+
 std::vector<std::pair<std::string, faabric::UnregisterRequest>>
 getUnregisterRequests()
 {
@@ -76,6 +87,7 @@ void clearMockRequests()
     functionCalls.clear();
     batchMessages.clear();
     resourceRequests.clear();
+    pendingMigrationsRequests.clear();
     unregisterRequests.clear();
 
     for (auto& p : queuedResourceResponses) {
@@ -128,6 +140,25 @@ faabric::HostResources FunctionCallClient::getResources()
     return response;
 }
 
+// This function call is used by the master host of an application to let know
+// other hosts running functions of the same application that a migration
+// opportunity has been found.
+void FunctionCallClient::sendPendingMigrations(
+  std::shared_ptr<faabric::PendingMigrations> req)
+{
+    faabric::PendingMigrations request;
+    faabric::EmptyResponse response;
+
+    if (faabric::util::isMockMode()) {
+        faabric::util::UniqueLock lock(mockMutex);
+        pendingMigrationsRequests.emplace_back(host, req);
+    } else {
+        syncSend(faabric::scheduler::FunctionCalls::PendingMigrations,
+                 req.get(),
+                 &response);
+    }
+}
+
 void FunctionCallClient::executeFunctions(
   const std::shared_ptr<faabric::BatchExecuteRequest> req)
 {