feat: Add pending reason: Reservation

PKUHPC · Dec 19, 2024 · fc14a99 · fc14a99
1 parent b05d58a
commit fc14a99
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/src/CraneCtld/TaskScheduler.cpp b/src/CraneCtld/TaskScheduler.cpp
@@ -18,6 +18,8 @@
 
 #include "TaskScheduler.h"
 
+#include <absl/time/time.h>
+
 #include "AccountManager.h"
 #include "CranedKeeper.h"
 #include "CranedMetaContainer.h"
@@ -2109,6 +2111,9 @@ void MinLoadFirst::CalculateNodeSelectionInfoOfPartition_(
       }
     }
 
+    auto& first_resv_time =
+        node_selection_info_ref.first_resv_time_map[craned_id];
+    first_resv_time = absl::InfiniteFuture();
     for (const auto& [ReservationId, res] :
          craned_meta->reservation_resource_map) {
       const auto& reservation =
@@ -2118,11 +2123,20 @@ void MinLoadFirst::CalculateNodeSelectionInfoOfPartition_(
                     ReservationId);
         continue;
       }
-      // TODO: expired reservation should be removed from the map.
-      time_res_vec.emplace_back(std::max(now, reservation->start_time),
+      absl::Time start_time = reservation->start_time;
+      absl::Time end_time = reservation->end_time;
+      if (end_time < now) {
+        // TODO: expired reservation should be removed from the map.
+        continue;
+      }
+      if (start_time < now) {
+        start_time = now;
+      }
+      time_res_vec.emplace_back(std::max(now, start_time),
                                 std::make_pair(false, res));
-      time_res_vec.emplace_back(std::max(now, reservation->end_time),
+      time_res_vec.emplace_back(std::max(now, end_time),
                                 std::make_pair(true, res));
+      first_resv_time = std::min(first_resv_time, start_time);
     }
 
     std::sort(
@@ -2761,6 +2775,11 @@ void MinLoadFirst::NodeSelect(
       // The task can't be started now. Set pending reason and move to the next
       // pending task.
       for (auto& craned_id : craned_ids) {
+        auto& first_resv_time = node_info.first_resv_time_map.at(craned_id);
+        if (first_resv_time < now + task->time_limit) {
+          task->pending_reason = "Resource Reserved";
+          break;
+        }
         auto& res_avail =
             node_info.node_time_avail_res_map.at(craned_id).at(now);
         if (!(task->Resources().EachNodeResMap().at(craned_id) <= res_avail)) {

diff --git a/src/CraneCtld/TaskScheduler.h b/src/CraneCtld/TaskScheduler.h
@@ -170,6 +170,7 @@ class MinLoadFirst : public INodeSelectionAlgo {
     std::multimap<uint32_t /* # of running tasks */, CranedId>
         task_num_node_id_map;
     std::unordered_map<CranedId, TimeAvailResMap> node_time_avail_res_map;
+    std::unordered_map<CranedId, absl::Time> first_resv_time_map;
   };
 
   static void CalculateNodeSelectionInfoOfPartition_(