Allow collection of performance reports for geospatial benchmarks (#1569

)

tests/conftest.py

@@ -70,6 +70,12 @@ def pytest_addoption(parser):
         "--benchmark", action="store_true", help="Collect benchmarking data for tests"
     )
 
+    parser.addoption(
+        "--performance-report",
+        action="store_true",
+        help="Collect performance report for tests",
+    )
+
     parser.addoption(
         "--memray",
         action="store",
@@ -910,3 +916,33 @@ def _memray_profile(client):
                     s3.put(archive, s3_performance_url)
 
             yield _memray_profile
+
+
+@pytest.fixture
+def performance_report(
+    pytestconfig,
+    s3,
+    s3_performance_url,
+    s3_storage_options,
+    test_run_benchmark,
+    tmp_path,
+):
+    if not test_run_benchmark:
+        yield
+    else:
+        if not pytestconfig.getoption("--performance-report"):
+            yield contextlib.nullcontext
+        else:
+
+            @contextlib.contextmanager
+            def _performance_report():
+                try:
+                    filename = f"{s3_performance_url}/performance_report.html.gz"
+                    with distributed.performance_report(
+                        filename=filename, storage_options=s3_storage_options
+                    ):
+                        yield
+                finally:
+                    test_run_benchmark.performance_report_url = filename
+
+            yield _performance_report

tests/geospatial/conftest.py

@@ -27,7 +27,9 @@ def cluster_name(request, scale):
 
 
 @pytest.fixture()
-def client_factory(cluster_name, github_cluster_tags, benchmark_all, memray_profile):
+def client_factory(
+    cluster_name, github_cluster_tags, benchmark_all, memray_profile, performance_report
+):
     import contextlib
 
     @contextlib.contextmanager
@@ -43,7 +45,9 @@ def _(n_workers, env=None, **cluster_kwargs):
             with cluster.get_client() as client:
                 # FIXME https://github.com/coiled/platform/issues/103
                 client.wait_for_workers(n_workers)
-                with memray_profile(client), benchmark_all(client):
+                with performance_report(), memray_profile(client), benchmark_all(
+                    client
+                ):
                     yield client
 
     return _