From e6237338619ae303ff0ce7718fb143db98cbd10c Mon Sep 17 00:00:00 2001 From: Nicolas Rabault Date: Wed, 1 Jul 2026 17:05:43 +0200 Subject: [PATCH] perf(tests): cache draccus docstring extraction (#3903) draccus re-parses each config class's source on every parse() to extract field help text (~2.5s for TrainPipelineConfig). Memoize it for the test session; the source is constant within a run. Fast Tests test time: 664s -> 404s (-39%). --- tests/conftest.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/tests/conftest.py b/tests/conftest.py index cadeaf0d3..31ff07b8e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,14 +14,23 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import traceback +import draccus.wrappers.docstring as _draccus_docstring import pytest from lerobot.configs.types import FeatureType, PipelineFeatureType, PolicyFeature from lerobot.utils.import_utils import is_package_available from tests.utils import DEVICE +# On every `draccus.parse()`, draccus rebuilds each dataclass field's help text by +# re-reading and re-parsing the class source (draccus.wrappers.docstring). For a config +# as large as TrainPipelineConfig this costs ~2.5s per parse — negligible for the single +# parse a CLI does, but tests parse configs hundreds of times. The source can't change +# within a run, so memoize it for the whole test session. +_draccus_docstring.get_attribute_docstring = functools.cache(_draccus_docstring.get_attribute_docstring) + # Import fixture modules as plugins. # Fixtures that depend on optional packages are only registered when those packages are available, # so that tests can be collected and run even with a minimal install.