From 49e2682f6faaedd28b809587c68ab14c3641cd9c Mon Sep 17 00:00:00 2001
From: Sukrit Kalra <sukritkalra@gmail.com>
Date: Mon, 20 May 2024 12:46:59 -0700
Subject: [PATCH] Limit solver time in a more fine-grained manner.

---
 schedulers/tetrisched/src/GurobiSolver.cpp    | 43 ++++++++++---------
 .../tetrisched/src/OptimizationPasses.cpp     | 22 +++++-----
 2 files changed, 33 insertions(+), 32 deletions(-)

diff --git a/schedulers/tetrisched/src/GurobiSolver.cpp b/schedulers/tetrisched/src/GurobiSolver.cpp
index eb434dee..9101fa3e 100644
--- a/schedulers/tetrisched/src/GurobiSolver.cpp
+++ b/schedulers/tetrisched/src/GurobiSolver.cpp
@@ -26,28 +26,31 @@ GurobiSolver::GurobiInterruptOptimizationCallback::
 void GurobiSolver::GurobiInterruptOptimizationCallback::callback() {
   try {
     auto currentTime = std::chrono::steady_clock::now();
-    if (where == GRB_CB_POLLING) {
-      if (params.timeLimitMs.has_value()) {
-        auto elapsedTimeMs =
-            std::chrono::duration_cast<std::chrono::milliseconds>(currentTime -
-                                                                  startTime)
-                .count();
-
-        if (elapsedTimeMs > params.timeLimitMs.value()) {
-          abort();
-        }
+
+    // Check if we need to abort due to time limits.
+    if (params.timeLimitMs.has_value()) {
+      auto elapsedTimeMs =
+          std::chrono::duration_cast<std::chrono::milliseconds>(currentTime -
+                                                                startTime)
+              .count();
+
+      if (elapsedTimeMs > params.timeLimitMs.value()) {
+        abort();
       }
-      if (params.newSolutionTimeLimitMs.has_value()) {
-        auto elapsedTimeMs =
-            std::chrono::duration_cast<std::chrono::milliseconds>(
-                currentTime - lastIncumbentSolutionTime)
-                .count();
-
-        if (elapsedTimeMs > params.newSolutionTimeLimitMs.value()) {
-          abort();
-        }
+    }
+    if (params.newSolutionTimeLimitMs.has_value()) {
+      auto elapsedTimeMs =
+          std::chrono::duration_cast<std::chrono::milliseconds>(
+              currentTime - lastIncumbentSolutionTime)
+              .count();
+
+      if (elapsedTimeMs > params.newSolutionTimeLimitMs.value()) {
+        abort();
       }
-    } else if (where == GRB_CB_MIPSOL && params.utilityUpperBound.has_value()) {
+    }
+
+    // Check if we need to abort due to the upper bound of the objective.
+    if (where == GRB_CB_MIPSOL && params.utilityUpperBound.has_value()) {
       auto solutionObjectiveValue = getDoubleInfo(GRB_CB_MIPSOL_OBJ);
       if (solutionObjectiveValue >= params.utilityUpperBound.value() -
                                         TETRISCHED_SOLUTION_UPPER_BOUND_DELTA) {
diff --git a/schedulers/tetrisched/src/OptimizationPasses.cpp b/schedulers/tetrisched/src/OptimizationPasses.cpp
index 39d97dcd..888832c0 100644
--- a/schedulers/tetrisched/src/OptimizationPasses.cpp
+++ b/schedulers/tetrisched/src/OptimizationPasses.cpp
@@ -699,8 +699,7 @@ void DiscretizationSelectorOptimizationPass::runPass(
   //             - minOccupancyTime] << std::endl;
   // }
 
-  // changing the STRL expressions
-  // for (auto &[discreteTimeRange, discreteGranularity] : timeRangeToGranularities)
+  // Changing the STRL expressions
   // Find Max expressions over NCK and remove NCK expressions with redundant
   // start times
   {
@@ -714,8 +713,6 @@ void DiscretizationSelectorOptimizationPass::runPass(
       {
         auto startTime = discreteTimeRange.first;
         auto endTime = discreteTimeRange.second;
-        // find ncks within startTime and endTime for this Max expr
-        // std::vector<ExpressionPtr> ncksWithinTimeRange;
 
         ExpressionPtr minStartTimeNckExpr = nullptr;
         ExpressionPtr prevSolutionNckExpr = nullptr;
@@ -726,20 +723,19 @@ void DiscretizationSelectorOptimizationPass::runPass(
             if (prevSolutionNckExpr == nullptr &&
                 (*child)->isPreviouslySatisfied()) {
               prevSolutionNckExpr = *child;
-              // break;
             }
             if (minStartTimeNckExpr != nullptr) {
-              if (minStartTimeNckExpr->getTimeBounds().startTimeRange.first >
-                  startTimeNck) {
+              if (minStartTimeNckExpr->getTimeBounds().startTimeRange.first > startTimeNck) {
                 minStartTimeNckExpr = *child;
               }
             } else {
               minStartTimeNckExpr = *child;
             }
-          } else{
-            if (startTimeNck < startTime)
-            {
-              throw exceptions::RuntimeException("No ncks for MAX expression: NckExpr: NckExpr(" + child->get()->getId() + ", " + child->get()->getName() + ")");
+          } else {
+            if (startTimeNck < startTime) {
+              throw exceptions::RuntimeException(
+                "No nCks for MAX expression: NckExpr: NckExpr(" + child->get()->getId() +
+                ", " + child->get()->getName() + ")");
             }
             break;
           }
@@ -756,7 +752,9 @@ void DiscretizationSelectorOptimizationPass::runPass(
         }
       }
       if (savedChildren.size() == 0){
-        throw exceptions::RuntimeException("No ncks for MAX expression: MaxExprName: MaxExpr(" + maxNckExpr->getId() + ", " + maxNckExpr->getName()+")");
+        throw exceptions::RuntimeException(
+          "No ncks for MAX expression: MaxExprName: MaxExpr(" +
+          maxNckExpr->getId() + ", " + maxNckExpr->getName()+")");
       }
       maxNckExpr->replaceChildren(std::move(savedChildren));
     }