gh-287: Add TaskGroups variants to async_tree benchmarks (#293)

itamaro · web-flow · commit c2e1adb82c28 · 2023-06-13T11:32:23.000+05:30
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.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager.toml
@@ -1,3 +1,7 @@
+[project]
+requires-python = ">=3.12"
+dynamic = ["version"]
+
 [tool.pyperformance]
 name = "async_tree_eager"
 extra_opts = ["eager"]
diff --git a/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_cpu_io_mixed.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_cpu_io_mixed.toml
@@ -1,3 +1,7 @@
+[project]
+requires-python = ">=3.12"
+dynamic = ["version"]
+
 [tool.pyperformance]
 name = "async_tree_eager_cpu_io_mixed"
 extra_opts = ["eager_cpu_io_mixed"]
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.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_io.toml
@@ -1,3 +1,7 @@
+[project]
+requires-python = ">=3.12"
+dynamic = ["version"]
+
 [tool.pyperformance]
 name = "async_tree_eager_io"
 extra_opts = ["eager_io"]
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.toml b/pyperformance/data-files/benchmarks/bm_async_tree/bm_async_tree_eager_memoization.toml
@@ -1,3 +1,7 @@
+[project]
+requires-python = ">=3.12"
+dynamic = ["version"]
+
 [tool.pyperformance]
 name = "async_tree_eager_memoization"
 extra_opts = ["eager_memoization"]
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/pyproject.toml b/pyperformance/data-files/benchmarks/bm_async_tree/pyproject.toml
@@ -7,4 +7,5 @@ dynamic = ["version"]
 
 [tool.pyperformance]
 name = "async_tree"
+tags = "asyncio"
 extra_opts = ["none"]
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_group(self, recurse_level):
+        if recurse_level == 0:
+            await self.workload_func()
+            return
+
+        async with asyncio.TaskGroup() as tg:
+            for _ in range(NUM_RECURSE_BRANCHES):
+                tg.create_task(
+                    self.recurse_with_task_group(recurse_level - 1))
+
     async def run(self):
-        await self.recurse(NUM_RECURSE_LEVELS)
+        if self.use_task_groups:
+            await self.recurse_with_task_group(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,8 @@ def add_parser_args(parser):
     benchmark = args.benchmark
 
     async_tree_class = BENCHMARKS[benchmark]
-    async_tree = async_tree_class()
-    runner.bench_async_func(f"async_tree_{benchmark}", async_tree.run)
+    async_tree = async_tree_class(use_task_groups=args.task_groups)
+    bench_name = f"async_tree_{benchmark}"
+    if args.task_groups:
+        bench_name += "_tg"
+    runner.bench_async_func(bench_name, async_tree.run)
