Add TaskGroups variants to all async_tree benchmarks

python · kumaraditya303 · Jun 13, 2023 · May 27, 2023 · May 27, 2023 · May 27, 2023
commit 3ab4a544ea723e0b2cae063e613faf75e1f73c61
diff --git a/pyperformance/data-files/benchmarks/MANIFEST b/pyperformance/data-files/benchmarks/MANIFEST
@@ -11,6 +11,14 @@ async_tree_eager	<local:async_tree>
 async_tree_eager_cpu_io_mixed	<local:async_tree>
 async_tree_eager_io	<local:async_tree>
 async_tree_eager_memoization	<local:async_tree>
+async_tree_tg	<local:async_tree>
+async_tree_cpu_io_mixed_tg	<local:async_tree>
+async_tree_io_tg	<local:async_tree>
+async_tree_memoization_tg	<local:async_tree>
+async_tree_eager_tg	<local:async_tree>
+async_tree_eager_cpu_io_mixed_tg	<local:async_tree>
+async_tree_eager_io_tg	<local:async_tree>
+async_tree_eager_memoization_tg	<local:async_tree>
 asyncio_tcp	<local>
 asyncio_tcp_ssl	<local:asyncio_tcp>
 concurrent_imap	<local>
@@ -82,6 +90,7 @@ xml_etree	<local>
 
 
 #[groups]
+#asyncio
 #startup
 #regex
 #serialize

diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_cpu_io_mixed_tg.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_cpu_io_mixed_tg.toml
@@ -0,0 +1,7 @@
+[project]
+requires-python = ">=3.11"
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "async_tree_cpu_io_mixed_tg"
+extra_opts = ["cpu_io_mixed", "--task-groups"]
diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_cpu_io_mixed_tg.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_cpu_io_mixed_tg.toml
@@ -0,0 +1,7 @@
+[project]
+requires-python = ">=3.12"
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "async_tree_eager_cpu_io_mixed_tg"
+extra_opts = ["eager_cpu_io_mixed", "--task-groups"]
diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_io_tg.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_io_tg.toml
@@ -0,0 +1,7 @@
+[project]
+requires-python = ">=3.12"
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "async_tree_eager_io_tg"
+extra_opts = ["eager_io", "--task-groups"]
diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_memoization_tg.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_memoization_tg.toml
@@ -0,0 +1,7 @@
+[project]
+requires-python = ">=3.12"
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "async_tree_eager_memoization_tg"
+extra_opts = ["eager_memoization", "--task-groups"]
diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_tg.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_tg.toml
@@ -0,0 +1,7 @@
+[project]
+requires-python = ">=3.12"
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "async_tree_eager_tg"
+extra_opts = ["eager", "--task-groups"]
diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_io_tg.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_io_tg.toml
@@ -0,0 +1,7 @@
+[project]
+requires-python = ">=3.11"
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "async_tree_io_tg"
+extra_opts = ["io", "--task-groups"]
diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_memoization_tg.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_memoization_tg.toml
@@ -0,0 +1,7 @@
+[project]
+requires-python = ">=3.11"
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "async_tree_memoization_tg"
+extra_opts = ["memoization", "--task-groups"]
diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_tg.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_tg.toml
@@ -0,0 +1,7 @@
+[project]
+requires-python = ">=3.11"
+dynamic = ["version"]
+
+[tool.pyperformance]
+name = "async_tree_tg"
+extra_opts = ["none", "--task-groups"]
diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/run_benchmark.py b/pyperformance/data-files/benchmarks/bm_async_tree/run_benchmark.py
@@ -12,8 +12,8 @@
                    the other half simulate the same workload as the
                    "memoization" variant.
 
-All variants also have an "eager" flavor that uses
-the asyncio eager task factory (if available).
+All variants also have an "eager" flavor that uses the asyncio eager task
+factory (if available), and a "tg" variant that uses TaskGroups.
 """
 
 
@@ -34,8 +34,9 @@
 
 
 class AsyncTree:
-    def __init__(self):
+    def __init__(self, use_task_groups=False):
         self.cache = {}
+        self.use_task_groups = use_task_groups
         # set to deterministic random, so that the results are reproducible
         random.seed(RANDOM_SEED)
 
@@ -47,17 +48,31 @@ async def workload_func(self):
             "To be implemented by each variant's derived class."
         )
 
-    async def recurse(self, recurse_level):
+    async def recurse_with_gather(self, recurse_level):
         if recurse_level == 0:
             await self.workload_func()
             return
 
         await asyncio.gather(
-            *[self.recurse(recurse_level - 1) for _ in range(NUM_RECURSE_BRANCHES)]
+            *[self.recurse_with_gather(recurse_level - 1)
+              for _ in range(NUM_RECURSE_BRANCHES)]
+        )
+
+    async def recurse_with_task_groups(self, recurse_level):
+        if recurse_level == 0:
+            await self.workload_func()
+            return
+
+        await asyncio.gather(
+            *[self.recurse_with_task_groups(recurse_level - 1)
+              for _ in range(NUM_RECURSE_BRANCHES)]
         )
 
     async def run(self):
-        await self.recurse(NUM_RECURSE_LEVELS)
+        if self.use_task_groups:
+            await self.recurse_with_task_groups(NUM_RECURSE_LEVELS)
+        else:
+            await self.recurse_with_gather(NUM_RECURSE_LEVELS)
 
 
 class EagerMixin:
@@ -132,6 +147,8 @@ def add_metadata(runner):
 
 def add_cmdline_args(cmd, args):
     cmd.append(args.benchmark)
+    if args.task_groups:
+        cmd.append("--task-groups")
 
 
 def add_parser_args(parser):
@@ -149,6 +166,12 @@ def add_parser_args(parser):
                    "memoization" variant.
 """,
     )
+    parser.add_argument(
+        "--task-groups",
+        action="store_true",
+        default=False,
+        help="Use TaskGroups instead of gather.",
+    )
 
 
 BENCHMARKS = {
@@ -171,5 +194,5 @@ def add_parser_args(parser):
     benchmark = args.benchmark
 
     async_tree_class = BENCHMARKS[benchmark]
-    async_tree = async_tree_class()
+    async_tree = async_tree_class(use_task_groups=args.task_groups)
     runner.bench_async_func(f"async_tree_{benchmark}", async_tree.run)