From c66ef7525ab1ac57598204a9b736756e106131cf Mon Sep 17 00:00:00 2001
From: Sukrit Kalra <sukritkalra@gmail.com>
Date: Sun, 24 Mar 2024 10:23:26 -0700
Subject: [PATCH] Implement no consideration period in STRL-based schedulers.|

---
 main.py                            |  7 +++++++
 schedulers/tetrisched_scheduler.py | 26 +++++++++++++++++++++-----
 2 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/main.py b/main.py
index 5db02992..b9804410 100644
--- a/main.py
+++ b/main.py
@@ -376,6 +376,13 @@
     "the available tasks (in µs). The default value is to pack until the maximum "
     "deadline for each batch of available tasks.",
 )
+flags.DEFINE_integer(
+    "scheduler_plan_ahead_no_consideration_gap",
+    0,
+    "The length of time gap (in µs) for which the reconsiderations are frozen. "
+    "From the current time to the consideration gap, any tasks placed will not be "
+    "reconsidered for rescheduling.",
+)
 flags.DEFINE_integer(
     "scheduler_time_discretization",
     1,
diff --git a/schedulers/tetrisched_scheduler.py b/schedulers/tetrisched_scheduler.py
index fb737126..700b4297 100644
--- a/schedulers/tetrisched_scheduler.py
+++ b/schedulers/tetrisched_scheduler.py
@@ -160,6 +160,10 @@ class TetriSchedScheduler(BaseScheduler):
             log to files with the format "gurobi_{sim_time}.log".
         _flags (`Optional[absl.flags]`): The runtime flags that are used to initialize
             a logger instance.
+        plan_ahead_no_consideration_gap (`EventTime`): The time gap after the current
+            time for which plan ahead is not considered malleable. Any tasks placed
+            during this time will not be reconsidered for scheduling in a particular
+            run.
     """
 
     def __init__(
@@ -181,6 +185,7 @@ def __init__(
         max_occupancy_threshold: float = 0.8,
         finer_discretization_at_prev_solution: bool = False,
         finer_discretization_window: EventTime = EventTime(5, EventTime.Unit.US),
+        plan_ahead_no_consideration_gap: EventTime = EventTime.zero(),
     ):
         if preemptive:
             raise ValueError("TetrischedScheduler does not support preemption.")
@@ -231,6 +236,7 @@ def __init__(
             _flags.scheduler_selective_rescheduling_sample_size if _flags else 5
         )
         self._plan_ahead_multiplier: int = 2
+        self._plan_ahead_no_consideration_gap = plan_ahead_no_consideration_gap
 
         # Scheduler configuration.
         self._scheduler_configuration = tetrisched.SchedulerConfig()
@@ -1096,11 +1102,21 @@ def construct_task_strl(
             retract_schedules,
         )
 
-        # If the task is already running or scheduled and we're not reconsidering
-        # previous schedulings, we block off all the time discretizations from the
-        # task's start until it completes.
-        if task.state == TaskState.RUNNING or (
-            task.state == TaskState.SCHEDULED and not retract_schedules
+        # We block off all the time discretizations from the task's start until it
+        # completes if any one of the following conditions hold:
+        # 1. The Task is already running (we do not consider preemption for now).
+        # 2. The Task is scheduled and the scheduler was specifically requested to
+        #    not retract schedules.
+        # 3. The intended start time of the task falls within the no consideration
+        #    period from the current time.
+        if (
+            task.state == TaskState.RUNNING  # 1.
+            or (task.state == TaskState.SCHEDULED and not retract_schedules)  # 2.
+            or (  # 3.
+                task.state == TaskState.SCHEDULED
+                and task.current_placement.placement_time
+                < current_time + self._plan_ahead_no_consideration_gap
+            )
         ):
             # Find the Partition where the task is running or to be scheduled.
             scheduled_partition = partitions.get_partition_for_worker_id(