Add bm_runner 'trialrun' subcommand. (SciTools#5957)

* Add bm_runner 'trialrun' subcommand.

* Remove debug.

* Review changes.

* Reinstate use of argparse.REMAINDER

* Make trialrun 'runpath' arg non-optional.

* Better help documentation.

* Review changes.

---------

Co-authored-by: Martin Yeo <[email protected]>

benchmarks/bm_runner.py

@@ -544,6 +544,63 @@ def func(args: argparse.Namespace) -> None:
         _subprocess_runner([args.asv_sub_command, *args.asv_args], asv=True)
 
 
+class TrialRun(_SubParserGenerator):
+    name = "trialrun"
+    description = (
+        "Fast trial-run a given benchmark, to check it works : "
+        "in a provided or latest-lockfile environment, "
+        "with no repeats for accuracy of measurement."
+    )
+    epilog = (
+        "e.g. python bm_runner.py trialrun "
+        "MyBenchmarks.time_calc ${DATA_GEN_PYTHON}"
+        "\n\nNOTE: 'runpath' also replaces $DATA_GEN_PYTHON during the run."
+    )
+
+    def add_arguments(self) -> None:
+        self.subparser.add_argument(
+            "benchmark",
+            type=str,
+            help=(
+                "A benchmark name, possibly including wildcards, "
+                "as supported by the ASV '--bench' argument."
+            ),
+        )
+        self.subparser.add_argument(
+            "runpath",
+            type=str,
+            help=(
+                "A path to an existing python executable, "
+                "to completely bypass environment building."
+            ),
+        )
+
+    @staticmethod
+    def func(args: argparse.Namespace) -> None:
+        if args.runpath:
+            # Shortcut creation of a data-gen environment
+            # - which is also the trial-run env.
+            python_path = Path(args.runpath).resolve()
+            environ["DATA_GEN_PYTHON"] = str(python_path)
+        _setup_common()
+        # get path of data-gen environment, setup by previous call
+        python_path = environ["DATA_GEN_PYTHON"]
+        # allow 'on-demand' benchmarks
+        environ["ON_DEMAND_BENCHMARKS"] = "1"
+        asv_command = [
+            "run",
+            "--bench",
+            args.benchmark,
+            # no repeats for timing accuracy
+            "--quick",
+            "--show-stderr",
+            # do not build a unique env : run test in data-gen environment
+            "--environment",
+            f"existing:{python_path}",
+        ] + args.asv_args
+        _subprocess_runner(asv_command, asv=True)
+
+
 class GhPost(_SubParserGenerator):
     name = "_gh_post"
     description = (
@@ -569,11 +626,24 @@ def add_asv_arguments(self) -> None:
 def main():
     parser = ArgumentParser(
         description="Run the Iris performance benchmarks (using Airspeed Velocity).",
-        epilog="More help is available within each sub-command.",
+        epilog=(
+            "More help is available within each sub-command."
+            "\n\nNOTE(1): a separate python environment is created to "
+            "construct test files.\n   Set $DATA_GEN_PYTHON to avoid the cost "
+            "of this."
+            "\nNOTE(2): iris-test-data is downloaded and cached within the "
+            "data generation environment.\n   Set "
+            "$OVERRIDE_TEST_DATA_REPOSITORY to avoid the cost of this."
+            "\nNOTE(3): test data is cached within the "
+            "benchmarks code directory, and uses a lot of disk space "
+            "of disk space (Gb).\n   Set $BENCHMARK_DATA to specify where this "
+            "space can be safely allocated."
+        ),
+        formatter_class=argparse.RawTextHelpFormatter,
     )
     subparsers = parser.add_subparsers(required=True)
 
-    for gen in (Overnight, Branch, CPerf, SPerf, Custom, GhPost):
+    for gen in (Overnight, Branch, CPerf, SPerf, Custom, TrialRun, GhPost):
         _ = gen(subparsers).subparser
 
     parsed = parser.parse_args()

noxfile.py

@@ -313,5 +313,5 @@ def benchmarks(session: nox.sessions.Session):
         )
         session.error(message)
     session.install("asv", "nox")
-    with session.chdir(Path(__file__).parent / "benchmarks"):
-        session.run("python", "bm_runner.py", *session.posargs)
+    bm_runner_path = Path(__file__).parent / "benchmarks" / "bm_runner.py"
+    session.run("python", bm_runner_path, *session.posargs)