added ramp down feature

Signed-off-by: Rishabh Singh <sngri@amazon.com>

Author: rishabh6788

osbenchmark/worker_coordinator/worker_coordinator.py

@@ -2836,11 +2836,17 @@ def schedule_for(task_allocation, parameter_source):
 
     if requires_time_period_schedule(task, runner_for_op, params_for_op):
         warmup_time_period = task.warmup_time_period if task.warmup_time_period else 0
+        ramp_down_time_period = task.ramp_down_time_period if task.ramp_down_time_period else 0
         if client_index == 0:
             logger.info("Creating time-period based schedule with [%s] distribution for [%s] with a warmup period of [%s] "
                         "seconds and a time period of [%s] seconds.", task.schedule, task.name,
                         str(warmup_time_period), str(task.time_period))
-        loop_control = TimePeriodBased(warmup_time_period, task.time_period)
+        loop_control = TimePeriodBased(warmup_time_period, task.time_period, ramp_down_time_period,
+                                       client_index, task.clients)
+        # Log individual client duration if ramp-down is enabled
+        if ramp_down_time_period > 0 and client_index == 0:
+            logger.info("Ramp-down enabled: clients will stop in reverse order over [%s] seconds",
+                        str(ramp_down_time_period))
     else:
         warmup_iterations = task.warmup_iterations if task.warmup_iterations else 0
         if task.iterations:
@@ -2950,13 +2956,32 @@ async def __call__(self):
 
 
 class TimePeriodBased:
-    def __init__(self, warmup_time_period, time_period):
+    def __init__(self, warmup_time_period, time_period, ramp_down_time_period=None,
+                 client_index=None, total_clients=None):
         self._warmup_time_period = warmup_time_period
         self._time_period = time_period
+        self._ramp_down_time_period = ramp_down_time_period or 0
+        self._client_index = client_index
+        self._total_clients = total_clients
+        self.logger = logging.getLogger(__name__)
+
         if warmup_time_period is not None and time_period is not None:
-            self._duration = self._warmup_time_period + self._time_period
+            self._base_duration = self._warmup_time_period + self._time_period
+
+            # Calculate how early this client should stop during ramp-down
+            # Clients stop in REVERSE order: Client 0 stops first, Client (N-1) stops last
+            if self._ramp_down_time_period > 0 and client_index is not None and total_clients is not None:
+                reverse_index = (total_clients - 1) - client_index
+                client_early_stop = self._ramp_down_time_period * (reverse_index / total_clients)
+                self._duration = self._base_duration - client_early_stop
+                self.logger.info("Client [%d/%d] will run for %.2f seconds (base: %.2f, early stop: %.2f due to ramp-down)",
+                                 client_index, total_clients, self._duration, self._base_duration, client_early_stop)
+            else:
+                self._duration = self._base_duration
         else:
             self._duration = None
+            self._base_duration = None
+
         self._start = None
         self._now = None
 

osbenchmark/workload/loader.py

@@ -1816,14 +1816,16 @@ def parse_parallel(self, ops_spec, ops, test_procedure_name):
         default_warmup_time_period = self._r(ops_spec, "warmup-time-period", error_ctx="parallel", mandatory=False)
         default_time_period = self._r(ops_spec, "time-period", error_ctx="parallel", mandatory=False)
         default_ramp_up_time_period = self._r(ops_spec, "ramp-up-time-period", error_ctx="parallel", mandatory=False)
+        default_ramp_down_time_period = self._r(ops_spec, "ramp-down-time-period", error_ctx="parallel", mandatory=False)
         clients = self._r(ops_spec, "clients", error_ctx="parallel", mandatory=False)
         completed_by = self._r(ops_spec, "completed-by", error_ctx="parallel", mandatory=False)
 
         # now descent to each operation
         tasks = []
         for task in self._r(ops_spec, "tasks", error_ctx="parallel"):
             tasks.append(self.parse_task(task, ops, test_procedure_name, default_warmup_iterations, default_iterations,
-                                         default_warmup_time_period, default_time_period, default_ramp_up_time_period, completed_by))
+                                         default_warmup_time_period, default_time_period, default_ramp_up_time_period,
+                                         default_ramp_down_time_period, completed_by))
 
         for task in tasks:
             if task.ramp_up_time_period != default_ramp_up_time_period:
@@ -1833,6 +1835,13 @@ def parse_parallel(self, ops_spec, ops, test_procedure_name):
                 else:
                     self._error(f"task '{task.name}' specifies a different ramp-up-time-period than its enclosing "
                                 f"'parallel' element in test-procedure '{test_procedure_name}'.")
+            if task.ramp_down_time_period != default_ramp_down_time_period:
+                if default_ramp_down_time_period is None:
+                    self._error(f"task '{task.name}' in 'parallel' element of test-procedure '{test_procedure_name}' specifies "
+                                f"a ramp-down-time-period but it is only allowed on the 'parallel' element.")
+                else:
+                    self._error(f"task '{task.name}' specifies a different ramp-down-time-period than its enclosing "
+                                f"'parallel' element in test-procedure '{test_procedure_name}'.")
         if completed_by:
             completion_task = None
             for task in tasks:
@@ -1848,7 +1857,7 @@ def parse_parallel(self, ops_spec, ops, test_procedure_name):
         return workload.Parallel(tasks, clients)
 
     def parse_task(self, task_spec, ops, test_procedure_name, default_warmup_iterations=None, default_iterations=None,
-                   default_warmup_time_period=None, default_time_period=None, default_ramp_up_time_period=None,
+                   default_warmup_time_period=None, default_time_period=None, default_ramp_up_time_period=None, default_ramp_down_time_period=None,
                    completed_by_name=None):
 
         op_spec = task_spec["operation"]
@@ -1874,6 +1883,8 @@ def parse_task(self, task_spec, ops, test_procedure_name, default_warmup_iterati
                                               default_value=default_time_period),
                           ramp_up_time_period=self._r(task_spec, "ramp-up-time-period", error_ctx=op.name,
                                                          mandatory=False, default_value=default_ramp_up_time_period),
+                          ramp_down_time_period=self._r(task_spec, "ramp-down-time-period", error_ctx=op.name,
+                                                           mandatory=False, default_value=default_ramp_down_time_period),
                           clients=self._r(task_spec, "clients", error_ctx=op.name, mandatory=False, default_value=1),
                           completes_parent=(task_name == completed_by_name),
                           schedule=schedule,
@@ -1901,6 +1912,19 @@ def parse_task(self, task_spec, ops, test_procedure_name, default_warmup_iterati
                 self._error(f"The warmup-time-period of operation '{op.name}' in test_procedure '{test_procedure_name}' is "
                             f"{task.warmup_time_period} seconds but must be greater than or equal to the "
                             f"ramp-up-time-period of {task.ramp_up_time_period} seconds.")
+        if task.ramp_down_time_period is not None:
+            if task.time_period is None:
+                self._error(f"Operation '{op.name}' in test_procedure '{test_procedure_name}' defines a ramp-down time period of "
+                            f"{task.ramp_down_time_period} seconds but no time-period.")
+            elif task.time_period < task.ramp_down_time_period:
+                self._error(f"The time-period of operation '{op.name}' in test_procedure '{test_procedure_name}' is "
+                            f"{task.time_period} seconds but must be greater than or equal to the "
+                            f"ramp-down-time-period of {task.ramp_down_time_period} seconds.")
+
+        if (task.warmup_iterations is not None or task.iterations is not None) and task.ramp_down_time_period is not None:
+            self._error(f"Operation '{op.name}' in test_procedure '{test_procedure_name}' defines a ramp-down time period of "
+                        f"{task.ramp_down_time_period} seconds as well as {task.warmup_iterations} warmup iterations and "
+                        f"{task.iterations} iterations but mixing time periods and iterations is not allowed.")
 
         return task
 

osbenchmark/workload/workload.py

@@ -937,7 +937,8 @@ class Task:
     IGNORE_RESPONSE_ERROR_LEVEL_WHITELIST = ["non-fatal"]
 
     def __init__(self, name, operation, tags=None, meta_data=None, warmup_iterations=None, iterations=None,
-                 warmup_time_period=None, time_period=None, ramp_up_time_period=None, clients=1, completes_parent=False,
+                 warmup_time_period=None, time_period=None, ramp_up_time_period=None, ramp_down_time_period=None,
+                 clients=1, completes_parent=False,
                  schedule=None, params=None):
         self.name = name
         self.operation = operation
@@ -953,6 +954,7 @@ def __init__(self, name, operation, tags=None, meta_data=None, warmup_iterations
         self.warmup_time_period = warmup_time_period
         self.time_period = time_period
         self.ramp_up_time_period = ramp_up_time_period
+        self.ramp_down_time_period = ramp_down_time_period
         self.clients = clients
         self.completes_parent = completes_parent
         self.schedule = schedule
@@ -1034,16 +1036,17 @@ def __hash__(self):
         # Note that we do not include `params` in __hash__ and __eq__ (the other attributes suffice to uniquely define a task)
         return hash(self.name) ^ hash(self.operation) ^ hash(self.warmup_iterations) ^ hash(self.iterations) ^ \
                hash(self.warmup_time_period) ^ hash(self.time_period) ^ hash(self.ramp_up_time_period) ^ \
-               hash(self.clients) ^ hash(self.schedule) ^ hash(self.completes_parent)
+               hash(self.ramp_down_time_period) ^ hash(self.clients) ^ hash(self.schedule) ^ hash(self.completes_parent)
 
     def __eq__(self, other):
         # Note that we do not include `params` in __hash__ and __eq__ (the other attributes suffice to uniquely define a task)
         return isinstance(other, type(self)) and (self.name, self.operation, self.warmup_iterations, self.iterations,
                                                   self.warmup_time_period, self.time_period, self.ramp_up_time_period,
-                                                  self.clients, self.schedule,self.completes_parent) == (other.name,
-                                                                             other.operation, other.warmup_iterations,
+                                                  self.ramp_down_time_period, self.clients, self.schedule,
+                                                  self.completes_parent) == (other.name, other.operation, other.warmup_iterations,
                                                                              other.iterations, other.warmup_time_period, other.time_period,
-                                                                             self.ramp_up_time_period, other.clients, other.schedule,
+                                                                             other.ramp_up_time_period, other.ramp_down_time_period,
+                                                                             other.clients, other.schedule,
                                                                              other.completes_parent)
 
     def __iter__(self):

tests/worker_coordinator/worker_coordinator_test.py

@@ -2313,3 +2313,166 @@ def test_check_cpu_usage_drops_error_when_queue_full(self):
         # Should not raise
         self.actor._check_cpu_usage() # pylint: disable=protected-access
         assert full_queue.qsize() == 1  # Still full; error dropped
+
+class TimePeriodBasedTests(TestCase):
+    # pylint: disable=protected-access
+    def test_time_period_based_without_ramp_down(self):
+        # Test basic time-period based schedule without ramp-down
+        loop_control = worker_coordinator.TimePeriodBased(
+            warmup_time_period=10,
+            time_period=100,
+            ramp_down_time_period=None,
+            client_index=0,
+            total_clients=4
+        )
+
+        # Verify duration calculation
+        self.assertEqual(110, loop_control._duration)
+        self.assertEqual(110, loop_control._base_duration)
+        self.assertEqual(10, loop_control._warmup_time_period)
+        self.assertEqual(100, loop_control._time_period)
+        self.assertFalse(loop_control.infinite)
+
+    def test_time_period_based_with_ramp_down_client_0(self):
+        # Test ramp-down for client 0 (first client, stops earliest)
+        loop_control = worker_coordinator.TimePeriodBased(
+            warmup_time_period=10,
+            time_period=100,
+            ramp_down_time_period=20,
+            client_index=0,
+            total_clients=4
+        )
+
+        # Client 0: reverse_index = 3, early_stop = 20 * (3/4) = 15
+        # duration = 110 - 15 = 95
+        self.assertEqual(110, loop_control._base_duration)
+        self.assertEqual(95, loop_control._duration)
+        self.assertEqual(20, loop_control._ramp_down_time_period)
+
+    def test_time_period_based_with_ramp_down_client_3(self):
+        # Test ramp-down for client 3 (last client, runs full duration)
+        loop_control = worker_coordinator.TimePeriodBased(
+            warmup_time_period=10,
+            time_period=100,
+            ramp_down_time_period=20,
+            client_index=3,
+            total_clients=4
+        )
+
+        # Client 3: reverse_index = 0, early_stop = 20 * (0/4) = 0
+        # duration = 110 - 0 = 110
+        self.assertEqual(110, loop_control._base_duration)
+        self.assertEqual(110, loop_control._duration)
+
+    def test_time_period_based_with_ramp_down_all_clients(self):
+        # Test that clients stop in reverse order with correct spacing
+        warmup = 10
+        time_period = 100
+        ramp_down = 20
+        total_clients = 4
+
+        durations = []
+        for client_index in range(total_clients):
+            loop_control = worker_coordinator.TimePeriodBased(
+                warmup_time_period=warmup,
+                time_period=time_period,
+                ramp_down_time_period=ramp_down,
+                client_index=client_index,
+                total_clients=total_clients
+            )
+            durations.append(loop_control._duration)
+
+        # Expected: [95, 100, 105, 110] - clients stop 5s apart
+        self.assertEqual([95, 100, 105, 110], durations)
+
+        # Verify spacing is correct
+        for i in range(1, len(durations)):
+            spacing = durations[i] - durations[i-1]
+            expected_spacing = ramp_down / total_clients
+            self.assertAlmostEqual(expected_spacing, spacing, places=2)
+
+class TaskRampDownTests(TestCase):
+    def test_task_with_ramp_down_time_period(self):
+        # Test that Task accepts ramp_down_time_period parameter
+        op = workload.Operation("test-op", workload.OperationType.Bulk.to_hyphenated_string(), {})
+        task = workload.Task(
+            name="test-task",
+            operation=op,
+            warmup_time_period=10,
+            time_period=100,
+            ramp_up_time_period=20,
+            ramp_down_time_period=30,
+            clients=4
+        )
+
+        self.assertEqual(10, task.warmup_time_period)
+        self.assertEqual(100, task.time_period)
+        self.assertEqual(20, task.ramp_up_time_period)
+        self.assertEqual(30, task.ramp_down_time_period)
+        self.assertEqual(4, task.clients)
+
+    def test_task_without_ramp_down_defaults_to_none(self):
+        # Test that ramp_down_time_period defaults to None
+        op = workload.Operation("test-op", workload.OperationType.Bulk.to_hyphenated_string(), {})
+        task = workload.Task(
+            name="test-task",
+            operation=op,
+            time_period=100,
+            clients=4
+        )
+
+        self.assertIsNone(task.ramp_down_time_period)
+
+    def test_task_equality_with_ramp_down(self):
+        # Test that tasks with different ramp_down_time_period are not equal
+        op = workload.Operation("test-op", workload.OperationType.Bulk.to_hyphenated_string(), {})
+
+        task1 = workload.Task(
+            name="test-task",
+            operation=op,
+            time_period=100,
+            ramp_down_time_period=20,
+            clients=4
+        )
+
+        task2 = workload.Task(
+            name="test-task",
+            operation=op,
+            time_period=100,
+            ramp_down_time_period=30,
+            clients=4
+        )
+
+        task3 = workload.Task(
+            name="test-task",
+            operation=op,
+            time_period=100,
+            ramp_down_time_period=20,
+            clients=4
+        )
+
+        self.assertNotEqual(task1, task2)
+        self.assertEqual(task1, task3)
+
+    def test_task_hash_includes_ramp_down(self):
+        # Test that hash includes ramp_down_time_period
+        op = workload.Operation("test-op", workload.OperationType.Bulk.to_hyphenated_string(), {})
+
+        task1 = workload.Task(
+            name="test-task",
+            operation=op,
+            time_period=100,
+            ramp_down_time_period=20,
+            clients=4
+        )
+
+        task2 = workload.Task(
+            name="test-task",
+            operation=op,
+            time_period=100,
+            ramp_down_time_period=30,
+            clients=4
+        )
+
+        # Different ramp_down should produce different hashes
+        self.assertNotEqual(hash(task1), hash(task2))