Merge pull request #2780 from tdadela/refactor_benchmark

Increase functionality of benchmark tests

benchmarks/dispatch.py

@@ -1,113 +1,154 @@
 """
 This file contains a benchmark to validate the performance of Locust itself.
 More precisely, the performance of the `UsersDispatcher` class which is responsible
-for calculating the distribution of users on each workers. This benchmark is to be used
+for calculating the distribution of users on each worker. This benchmark is to be used
 by people working on Locust's development.
 """
 
 from locust import User
 from locust.dispatch import UsersDispatcher
 from locust.runners import WorkerNode
 
+import argparse
+import gc
 import itertools
 import statistics
 import time
 
 from prettytable import PrettyTable
 
-# fmt: off
-WEIGHTS = [
-     5, 55, 37,  2, 97, 41, 33, 19, 19, 34,
-    78, 76, 28, 62, 69,  5, 55, 37,  2, 97,
-    41, 33, 19, 19, 34, 78, 76, 28, 62, 69,
-    41, 33, 19, 19, 34, 78, 76, 28, 62, 69,
-    41, 33, 19, 19, 34, 78, 76, 28, 62, 69,
-     5, 55, 37,  2, 97, 41, 33, 19, 19, 34,
-    78, 76, 28, 62, 69,  5, 55, 37,  2, 97,
-    41, 33, 19, 19, 34, 78, 76, 28, 62, 69,
-    41, 33, 19, 19, 34, 78, 76, 28, 62, 69,
-    41, 33, 19, 19, 34, 78, 76, 28, 62, 69,
-]
-# fmt: on
+NUMBER_OF_USER_CLASSES: int = 1000
+USER_CLASSES: list[type[User]] = []
+WEIGHTS = list(range(1, NUMBER_OF_USER_CLASSES + 1))
 
 for i, x in enumerate(WEIGHTS):
-    exec(f"class User{i+1}(User): weight = {x}")
+    exec(f"class User{i}(User): weight = {x}")
 
 # Equivalent to:
 #
-# class User1(User):
+# class User0(User):
 #     weight = 5
 #
-# class User2(User):
+# class User1(User):
 #     weight = 55
 # .
 # .
 # .
-# class User100(User):
-#     weight = 69
 
-exec("USER_CLASSES = [" + ",".join(f"User{i+1}" for i in range(len(WEIGHTS))) + "]")
+exec("USER_CLASSES = [" + ",".join(f"User{i}" for i in range(len(WEIGHTS))) + "]")
 # Equivalent to:
 #
 # USER_CLASSES = [
+#     User0,
 #     User1,
-#     User2,
 #     .
 #     .
 #     .
-#     User100,
 # ]
 
 
 if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-f", "--full-benchmark", action="store_true", help="run benchmark on full test matrix")
+    parser.add_argument(
+        "-i",
+        "--include-fixed-users",
+        action="store_true",
+        help="add test cases when 50 percent of users use User.fixed_count instead of User.weight",
+    )
+    parser.add_argument("-r", "--repeat", default=1, type=int, help="number of test cases with the same parameters")
+    parser.add_argument("-s", "--save-output", action="store_true", help="save test results to files")
+    args = parser.parse_args()
+
     now = time.time()
 
-    worker_count_cases = [10, 100, 500, 1000, 5000, 10_000, 15_000, 20_000]
-    user_count_cases = [10, 100, 1000, 10_000, 50_000, 100_000, 500_000]
-    number_of_user_classes_cases = [1, 10, 40, 60, 80, 100]
-    spawn_rate_cases = [1, 10, 100, 500, 1000, 2500, 5000, 10_000, 20_000, 25_000]
+    worker_count_cases = [10, 100, 1000]
+    user_count_cases = [10_000, 100_000, 1_000_000]
+    number_of_user_classes_cases = [1, 30, 1000]
+    spawn_rate_cases = [100, 10_000]
+    fixed_count_cases = [False, True] if args.include_fixed_users else [False]
+    # [0% fixed_count users, 50% fixed_count users] if args.mixed_user_types else [0% fixed_count users]
+    repeat_cases = list(range(1, args.repeat + 1))
+
+    if not args.full_benchmark:
+        worker_count_cases = [max(worker_count_cases)]
+        user_count_cases = [max(user_count_cases)]
+        number_of_user_classes_cases = [max(number_of_user_classes_cases)]
+        spawn_rate_cases = [max(spawn_rate_cases)]
 
     case_count = (
-        len(worker_count_cases) * len(user_count_cases) * len(number_of_user_classes_cases) * len(spawn_rate_cases)
+        len(worker_count_cases)
+        * len(user_count_cases)
+        * len(number_of_user_classes_cases)
+        * len(spawn_rate_cases)
+        * len(fixed_count_cases)
+        * len(repeat_cases)
     )
 
     results = {}
 
     try:
-        for case_index, (worker_count, user_count, number_of_user_classes, spawn_rate) in enumerate(
-            itertools.product(worker_count_cases, user_count_cases, number_of_user_classes_cases, spawn_rate_cases)
+        for case_index, (
+            worker_count,
+            user_count,
+            number_of_user_classes,
+            spawn_rate,
+            fixed_users,
+            iteration,
+        ) in enumerate(
+            itertools.product(
+                worker_count_cases,
+                user_count_cases,
+                number_of_user_classes_cases,
+                spawn_rate_cases,
+                fixed_count_cases,
+                repeat_cases,
+            )
         ):
-            if user_count / spawn_rate > 1000:
-                print(f"Skipping user_count = {user_count:,} - spawn_rate = {spawn_rate:,}")
-                continue
+            workers = [WorkerNode(str(i)) for i in range(worker_count)]
+            if fixed_users:
+                sum_fixed_weight = 0
+                for j in range(0, number_of_user_classes, 2):
+                    sum_fixed_weight += USER_CLASSES[j].weight
 
-            workers = [WorkerNode(str(i + 1)) for i in range(worker_count)]
+                for j in range(0, number_of_user_classes, 2):  # set fixed_weights for 50% of users
+                    USER_CLASSES[j].fixed_count = max(1, USER_CLASSES[j].weight // sum_fixed_weight)  # type: ignore # assigned .weight is int
 
-            ts = time.perf_counter()
+            ts = time.process_time()
             users_dispatcher = UsersDispatcher(
                 worker_nodes=workers,
-                user_classes=USER_CLASSES[:number_of_user_classes],  # noqa: F821 (Undefined name `USER_CLASSES`) -> It's created inside "exec"
+                user_classes=USER_CLASSES[:number_of_user_classes],
             )
-            instantiate_duration = time.perf_counter() - ts
+            instantiate_duration = time.process_time() - ts
 
             # Ramp-up
-            ts = time.perf_counter()
+            gc.disable()
+            ts = time.process_time()
             users_dispatcher.new_dispatch(target_user_count=user_count, spawn_rate=spawn_rate)
-            new_dispatch_ramp_up_duration = time.perf_counter() - ts
+            new_dispatch_ramp_up_duration = time.process_time() - ts
+            gc.enable()
+
             assert len(users_dispatcher.dispatch_iteration_durations) == 0
             users_dispatcher._wait_between_dispatch = 0
             all_dispatched_users_ramp_up = list(users_dispatcher)
             dispatch_iteration_durations_ramp_up = users_dispatcher.dispatch_iteration_durations[:]
 
             # Ramp-down
-            ts = time.perf_counter()
+            gc.disable()
+            ts = time.process_time()
             users_dispatcher.new_dispatch(target_user_count=0, spawn_rate=spawn_rate)
-            new_dispatch_ramp_down_duration = time.perf_counter() - ts
+            new_dispatch_ramp_down_duration = time.process_time() - ts
+            gc.enable()
+
             assert len(users_dispatcher.dispatch_iteration_durations) == 0
             users_dispatcher._wait_between_dispatch = 0
             all_dispatched_users_ramp_down = list(users_dispatcher)
             dispatch_iteration_durations_ramp_down = users_dispatcher.dispatch_iteration_durations[:]
 
+            if fixed_users:
+                for j in range(0, number_of_user_classes, 2):
+                    USER_CLASSES[j].fixed_count = None
+
             cpu_ramp_up = "{:3.3f}/{:3.3f}/{:3.3f}".format(  # noqa: UP032
                 1000 * statistics.mean(dispatch_iteration_durations_ramp_up),
                 1000 * min(dispatch_iteration_durations_ramp_up),
@@ -135,7 +176,10 @@
                 )
             )
 
-            results[(worker_count, user_count, number_of_user_classes, spawn_rate)] = (cpu_ramp_up, cpu_ramp_down)
+            results[(worker_count, user_count, number_of_user_classes, spawn_rate, fixed_users, iteration)] = (
+                cpu_ramp_up,
+                cpu_ramp_down,
+            )
 
     finally:
         table = PrettyTable()
@@ -144,13 +188,17 @@
             "Users",
             "User Classes",
             "Spawn Rate",
+            "Fixed Users",
+            "Iteration",
             "Ramp-Up (avg/min/max) (ms)",
             "Ramp-Down (avg/min/max) (ms)",
         ]
         table.align["Workers"] = "l"
         table.align["Users"] = "l"
         table.align["User Classes"] = "l"
         table.align["Spawn Rate"] = "l"
+        table.align["Fixed Users"] = "l"
+        table.align["Iteration"] = "c"
         table.align["Ramp-Up (avg/min/max) (ms)"] = "c"
         table.align["Ramp-Down (avg/min/max) (ms)"] = "c"
         table.add_rows(
@@ -160,22 +208,23 @@
                     f"{user_count:,}",
                     number_of_user_classes,
                     f"{spawn_rate:,}",
+                    "50%" if fixed_users else "0%",
+                    iteration,
                     cpu_ramp_up,
                     cpu_ramp_down,
                 ]
-                for (worker_count, user_count, number_of_user_classes, spawn_rate), (
+                for (worker_count, user_count, number_of_user_classes, spawn_rate, fixed_users, iteration), (
                     cpu_ramp_up,
                     cpu_ramp_down,
                 ) in results.items()
             ]
         )
         print()
-        print()
-        print()
         print(table)
 
-        with open(f"results-dispatch-benchmarks-{int(now)}.txt", "w") as file:
-            file.write(table.get_string())
+        if args.save_output:
+            with open(f"results-dispatch-benchmarks-{int(now)}.txt", "w") as file:
+                file.write(table.get_string())
 
-        with open(f"results-dispatch-benchmarks-{int(now)}.json", "w") as file:
-            file.write(table.get_json_string())
+            with open(f"results-dispatch-benchmarks-{int(now)}.json", "w") as file:
+                file.write(table.get_json_string())