perf: Auto-register bm.Scenario sub-classes (#2791)

* perf: Auto-register bm.Scenario sub-classes

* fix syntax error

* do not shadow dict

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: brettlangdon

benchmarks/README.rst

@@ -14,7 +14,7 @@ A scenario requires:
 * ``config.yaml``: specifies one or more sets of configuration variables for the benchmark
 * ``requirements_scenario.txt``: any additional dependencies
 
-The scenario class inherits from ``bm.Scenario`` and includes the configurable variables using ``bm.var``. The execution of the benchmark uses the ``run()`` generator function to yield a function that will handle the execution of a specified number of loops. The scenario class must also be decorated with ``@bm.register`` to ensure that it is run.
+The scenario class inherits from ``bm.Scenario`` and includes the configurable variables using ``bm.var``. The execution of the benchmark uses the ``run()`` generator function to yield a function that will handle the execution of a specified number of loops.
 
 Example
 ~~~~~~~
@@ -27,7 +27,6 @@ Example
   import bm
 
 
-  @bm.register
   class MyScenario(bm.Scenario):
       size = bm.var(type=int)
 

benchmarks/bm/__init__.py

@@ -1,10 +1,8 @@
 from ._scenario import Scenario
-from ._scenario import register
 from ._scenario import var
 
 
 __all__ = [
-    "register",
     "var",
     "Scenario",
 ]

benchmarks/bm/_scenario.py

@@ -6,18 +6,49 @@
 import six
 
 
-def register(cls):
-    """Decorator for scenario class that registers the scenario to be run."""
-    wrapped_cls = attr.s(cls)
-    _register(wrapped_cls)
-    return wrapped_cls
+var = attr.ib
 
 
-var = attr.ib
+def _register(scenario_cls):
+    """Registers a scenario for benchmarking."""
+    # This extends pyperf's runner by registering arguments for the scenario config
+    def add_cmdline_args(cmd, args):
+        for field in attr.fields(scenario_cls):
+            if hasattr(args, field.name):
+                cmd.extend(("--{}".format(field.name), str(getattr(args, field.name))))
+
+    runner = pyperf.Runner(add_cmdline_args=add_cmdline_args)
+    cmd = runner.argparser
+
+    for field in attr.fields(scenario_cls):
+        cmd.add_argument("--{}".format(field.name), type=field.type)
+
+    parsed_args = runner.parse_args()
+
+    config_dict = {
+        field.name: getattr(parsed_args, field.name)
+        for field in attr.fields(scenario_cls)
+        if hasattr(parsed_args, field.name)
+    }
+    scenario = scenario_cls(**config_dict)
+
+    runner.bench_time_func(scenario.scenario_name, scenario._pyperf)
+
+
+class ScenarioMeta(abc.ABCMeta):
+    def __init__(cls, name, bases, _dict):
+        super(ScenarioMeta, cls).__init__(name, bases, _dict)
 
+        # Make sure every sub-class is wrapped by `attr.s`
+        cls = attr.s()(cls)
 
-@attr.s
-class Scenario(six.with_metaclass(abc.ABCMeta)):
+        # Do not register the base Scenario class
+        # DEV: We cannot compare `cls` to `Scenario` since it doesn't exist yet
+        if cls.__module__ != __name__:
+            _register(cls)
+
+
+class Scenario(six.with_metaclass(ScenarioMeta)):
     """The base class for specifying a benchmark."""
 
     name = attr.ib(type=str)
@@ -44,29 +75,3 @@ def _pyperf(self, loops):
             pass
         finally:
             return dt
-
-
-def _register(scenario_cls):
-    """Registers a scenario for benchmarking."""
-    # This extends pyperf's runner by registering arguments for the scenario config
-    def add_cmdline_args(cmd, args):
-        for field in attr.fields(scenario_cls):
-            if hasattr(args, field.name):
-                cmd.extend(("--{}".format(field.name), str(getattr(args, field.name))))
-
-    runner = pyperf.Runner(add_cmdline_args=add_cmdline_args)
-    cmd = runner.argparser
-
-    for field in attr.fields(scenario_cls):
-        cmd.add_argument("--{}".format(field.name), type=field.type)
-
-    parsed_args = runner.parse_args()
-
-    config_dict = {
-        field.name: getattr(parsed_args, field.name)
-        for field in attr.fields(scenario_cls)
-        if hasattr(parsed_args, field.name)
-    }
-    scenario = scenario_cls(**config_dict)
-
-    runner.bench_time_func(scenario.scenario_name, scenario._pyperf)

benchmarks/django_simple/scenario.py

@@ -2,7 +2,6 @@
 import utils
 
 
-@bm.register
 class DjangoSimple(bm.Scenario):
     tracer_enabled = bm.var(type=bool)
     profiler_enabled = bm.var(type=bool)

benchmarks/encoder/scenario.py

@@ -2,7 +2,6 @@
 import utils
 
 
-@bm.register
 class Encoder(bm.Scenario):
     ntraces = bm.var(type=int)
     nspans = bm.var(type=int)

benchmarks/flask_simple/scenario.py

@@ -2,7 +2,6 @@
 import utils
 
 
-@bm.register
 class FlaskSimple(bm.Scenario):
     tracer_enabled = bm.var(type=bool)
     profiler_enabled = bm.var(type=bool)

benchmarks/span/scenario.py

@@ -4,7 +4,6 @@
 from ddtrace import Span as dd_Span
 
 
-@bm.register
 class Span(bm.Scenario):
     nspans = bm.var(type=int)
     ntags = bm.var(type=int)

benchmarks/tracer/scenario.py

@@ -8,7 +8,6 @@ def process_trace(self, trace):
         return
 
 
-@bm.register
 class Tracer(bm.Scenario):
     depth = bm.var(type=int)