omp: fix race conditions with elastic scaling

src/planner/Planner.cpp

@@ -27,9 +27,50 @@ namespace faabric::planner {
 // Utility Functions
 // ----------------------
 
-static int availableSlots(std::shared_ptr<Host> host)
+// Helper method to calculate how many available slots in the current host we
+// can scale-up to
+int availableOpenMpSlots(
+  int appId,
+  const std::string& mainHost,
+  const std::map<std::string, std::shared_ptr<Host>>& hostMap,
+  const faabric::batch_scheduler::InFlightReqs& inFlightReqs)
 {
-    int availableSlots = host->slots() - host->usedslots();
+    // At most, we can scale-up to the host size minus one (the caller thread)
+    int availableSlots =
+      hostMap.at(mainHost)->slots() - hostMap.at(mainHost)->usedslots();
+    assert(availableSlots <= hostMap.at(mainHost)->slots() - 1);
+
+    // However, we need to discard any in-flight apps that are also running
+    // in this host. This is to prevent a situation where a co-located app
+    // elastically scales beyond another's app minimum level of parallelism
+    for (const auto& [thisAppId, inFlightPair] : inFlightReqs) {
+        if (appId == thisAppId) {
+            continue;
+        }
+
+        // Check if the first message in the decision is scheduled to the
+        // same host we are
+        if (inFlightPair.second->hosts.at(0) == mainHost) {
+            // If so, check if the total OMP num threads is more than the
+            // current number of messages in flight, and if so subtract the
+            // difference from the available slots list
+            int requestedButNotOccupiedSlots =
+              inFlightPair.first->messages(0).ompnumthreads() -
+              inFlightPair.first->messages_size();
+
+            // This value could be smaller than zero if elastically scaled-up
+            if (requestedButNotOccupiedSlots > 0) {
+                availableSlots -= requestedButNotOccupiedSlots;
+
+                SPDLOG_DEBUG("Subtracting {} possible slots for app {}'s "
+                             "elastic scale from app {}!",
+                             requestedButNotOccupiedSlots,
+                             appId,
+                             thisAppId);
+            }
+        }
+    }
+
     assert(availableSlots >= 0);
 
     return availableSlots;
@@ -442,9 +483,10 @@ void Planner::setMessageResult(std::shared_ptr<faabric::Message> msg)
         if (it == req->messages().end()) {
             // Ditto as before. We want to allow setting the message result
             // more than once without breaking
-            SPDLOG_DEBUG(
-              "Setting result for non-existant (or finished) message: {}",
-              appId);
+            SPDLOG_DEBUG("Setting result for non-existant (or finished) "
+                         "message: {} (app: {})",
+                         msg->id(),
+                         appId);
         } else {
             SPDLOG_DEBUG("Removing message {} from app {}", msg->id(), appId);
 
@@ -777,6 +819,7 @@ Planner::callBatch(std::shared_ptr<BatchExecuteRequest> req)
     // does not modify it
     auto hostMapCopy =
       convertToBatchSchedHostMap(state.hostMap, state.nextEvictedHostIps);
+
     bool isScaleChange =
       decisionType == faabric::batch_scheduler::DecisionType::SCALE_CHANGE;
     bool isDistChange =
@@ -792,7 +835,10 @@ Planner::callBatch(std::shared_ptr<BatchExecuteRequest> req)
         auto oldDec = state.inFlightReqs.at(appId).second;
         auto mainHost = oldDec->hosts.at(0);
 
-        int numAvail = availableSlots(state.hostMap.at(mainHost));
+        // If there are co-located OpenMP apps, we should never use
+        // their `ompNumThreads`' slots
+        int numAvail = availableOpenMpSlots(
+          appId, mainHost, state.hostMap, state.inFlightReqs);
         int numRequested = req->messages_size();
         int lastMsgIdx =
           numRequested == 0 ? 0 : req->messages(numRequested - 1).groupidx();
@@ -864,6 +910,37 @@ Planner::callBatch(std::shared_ptr<BatchExecuteRequest> req)
     bool isOmp = req->messages_size() > 0 && req->messages(0).isomp();
     std::shared_ptr<BatchExecuteRequest> knownSizeReq = nullptr;
 
+    // For an OpenMP decision, we want to make sure that no in-flight
+    // tasks are currently in a join phase (from a repeated fork-join)
+    if (isOmp) {
+        for (const auto& [thisAppId, inFlightPair] : state.inFlightReqs) {
+            if (thisAppId == appId) {
+                continue;
+            }
+
+            int requestedButNotOccupiedSlots =
+              inFlightPair.first->messages(0).ompnumthreads() -
+              inFlightPair.first->messages_size();
+
+            // TODO: this only works for single host OpenMP requests
+            if (requestedButNotOccupiedSlots > 0) {
+                auto mainHost = inFlightPair.second->hosts.at(0);
+
+                SPDLOG_DEBUG("Tried to schedule OpenMP app (appid: {})"
+                             " in host {} while another in-flight OpenMP app"
+                             " (appid: {}) had too few messages in flight "
+                             " ({} < {})",
+                             appId,
+                             mainHost,
+                             thisAppId,
+                             inFlightPair.first->messages_size(),
+                             inFlightPair.first->messages(0).ompnumthreads());
+                hostMapCopy.at(mainHost)->usedSlots +=
+                  requestedButNotOccupiedSlots;
+            }
+        }
+    }
+
     // Check if there exists a pre-loaded scheduling decision for this app
     // (e.g. if we want to force a migration). Note that we don't want to check
     // pre-loaded decisions for dist-change requests
@@ -934,6 +1011,26 @@ Planner::callBatch(std::shared_ptr<BatchExecuteRequest> req)
         return decision;
     }
 
+    bool isSingleHost = decision->isSingleHost();
+    if (!isSingleHost && req->singlehosthint()) {
+        // In an elastic OpenMP execution, it may happen that we try to
+        // schedule an app, but another one has been elastically scaled
+        if (isNew && isOmp && req->elasticscalehint()) {
+            // Let the caller handle that there are not enough slots
+            return std::make_shared<
+              faabric::batch_scheduler::SchedulingDecision>(
+              NOT_ENOUGH_SLOTS_DECISION);
+        }
+
+        // This is likely a fatal error and a sign that something has gone
+        // very wrong. We still do not crash the planner
+        SPDLOG_ERROR(
+          "User provided single-host hint in BER, but decision is not!");
+
+        return std::make_shared<faabric::batch_scheduler::SchedulingDecision>(
+          NOT_ENOUGH_SLOTS_DECISION);
+    }
+
     // If we have managed to schedule a frozen request, un-freeze it by
     // removing it from the evicted request map
     if (state.evictedRequests.contains(appId)) {
@@ -1016,7 +1113,7 @@ Planner::callBatch(std::shared_ptr<BatchExecuteRequest> req)
             decision->print();
 #endif
 
-            // For a NEW MPI decision that was not preloaded we have
+            // For a NEW MPI/OpenMP decision that was not preloaded we have
             // preemptively scheduled all MPI messages but now we just need to
             // return the first one, and preload the rest
             if ((isMpi || isOmp) && knownSizeReq != nullptr) {
@@ -1231,10 +1328,6 @@ void Planner::dispatchSchedulingDecision(
     }
 
     bool isThreads = req->type() == faabric::BatchExecuteRequest::THREADS;
-    if (!isSingleHost && req->singlehosthint()) {
-        SPDLOG_ERROR(
-          "User provided single-host hint in BER, but decision is not!");
-    }
 
     for (const auto& [hostIp, hostReq] : hostRequests) {
         SPDLOG_DEBUG("Dispatching {} messages to host {} for execution",

src/planner/PlannerEndpointHandler.cpp

@@ -188,8 +188,19 @@ void PlannerEndpointHandler::onRequest(
                 }
 
                 if (inFlightPair.first->messages(0).isomp()) {
-                    inFlightAppResp->set_size(
-                      inFlightPair.first->messages(0).ompnumthreads());
+                    // What if we told here the scaled-up size (?)
+                    int numOmpThreads =
+                      inFlightPair.first->messages(0).ompnumthreads();
+
+                    if (inFlightPair.first->elasticscalehint() &&
+                        numOmpThreads < inFlightPair.first->messages_size()) {
+
+                        inFlightAppResp->set_size(
+                          inFlightPair.first->messages_size());
+                    } else {
+                        inFlightAppResp->set_size(
+                          inFlightPair.first->messages(0).ompnumthreads());
+                    }
                 }
 
                 for (const auto& hostIp : decision->hosts) {