Update compare scripts for cr8 async API changes

compare_run.py

@@ -12,6 +12,7 @@
 import shutil
 import subprocess
 import tempfile
+import asyncio
 from functools import partial
 from uuid import uuid4
 from typing import Optional, Dict, Any
@@ -193,7 +194,7 @@ def perf_stat_results(proc: subprocess.Popen) -> Dict[str, Any]:
     return metrics
 
 
-def _run_spec(version, spec, result_hosts, env, settings, tmpdir, protocol):
+async def _run_spec(version, spec, result_hosts, env, settings, tmpdir, protocol):
     crate_dir = get_crate(version)
     settings.setdefault('cluster.name', str(uuid4()))
     results = []
@@ -204,7 +205,7 @@ def _run_spec(version, spec, result_hosts, env, settings, tmpdir, protocol):
             pg_address = n.addresses['psql']
             benchmark_hosts = f'asyncpg://{pg_address.host}:{pg_address.port}'
         print(f'Running benchmark using protocol={protocol}, benchmark_hosts={benchmark_hosts}')
-        do_run_spec(
+        await do_run_spec(
             spec=spec,
             benchmark_hosts=benchmark_hosts,
             log=log,
@@ -215,7 +216,7 @@ def _run_spec(version, spec, result_hosts, env, settings, tmpdir, protocol):
         jfr_file = jfr_start(n.process.pid, tmpdir)
         perf_proc = perf_stat(n.process.pid)
         log.result = results.append
-        do_run_spec(
+        await do_run_spec(
             spec=spec,
             benchmark_hosts=n.http_url,
             log=log,
@@ -224,7 +225,7 @@ def _run_spec(version, spec, result_hosts, env, settings, tmpdir, protocol):
             action=['queries', 'load_data']
         )
         jfr_stop(n.process.pid)
-        do_run_spec(
+        await do_run_spec(
             spec=spec,
             benchmark_hosts=n.http_url,
             log=log,
@@ -235,7 +236,7 @@ def _run_spec(version, spec, result_hosts, env, settings, tmpdir, protocol):
     return (results, jfr_extract_metrics(jfr_file), perf_proc and perf_stat_results(perf_proc) or {})
 
 
-def run_compare(v1,
+async def run_compare(v1,
                 v2,
                 spec,
                 result_hosts,
@@ -251,8 +252,8 @@ def run_compare(v1,
     run_v2 = partial(_run_spec, v2, spec, result_hosts, env_v2, settings_v2, tmpdir, protocol)
     try:
         for _ in range(forks):
-            results_v1, jfr_metrics1, stat_result1 = run_v1()
-            results_v2, jfr_metrics2, stat_result2 = run_v2()
+            results_v1, jfr_metrics1, stat_result1 = await run_v1()
+            results_v2, jfr_metrics2, stat_result2 = await run_v2()
             compare_results(
                 results_v1,
                 jfr_metrics1,
@@ -309,7 +310,7 @@ def main():
     settings_v2 = settings.copy()
     settings_v2.update(dict_from_kw_args(args.setting_v2))
     try:
-        run_compare(
+        asyncio.run(run_compare(
             args.v1,
             args.v2,
             args.spec,
@@ -321,7 +322,7 @@ def main():
             settings_v2=settings_v2,
             show_plot=args.show_plot,
             protocol=args.protocol,
-        )
+        ))
     except KeyboardInterrupt:
         print('Exiting..')
 

compare_run_disk_usage.py

@@ -6,6 +6,7 @@
 """
 
 import argparse
+import asyncio
 from uuid import uuid4
 from cr8.run_crate import get_crate, CrateNode
 from cr8.run_spec import do_run_spec
@@ -32,12 +33,12 @@ def optimize_tables(cursor):
         cursor.execute(f'optimize table "{schema}"."{table}" with (flush = true, max_num_segments = 1)')
 
 
-def run(version, spec, env, settings):
+async def run(version, spec, env, settings):
     crate_dir = get_crate(version)
     settings.setdefault('cluster.name', str(uuid4()))
     with Logger() as log, CrateNode(crate_dir=crate_dir, settings=settings, env=env) as n:
         n.start()
-        do_run_spec(
+        await do_run_spec(
             spec=spec,
             log=log,
             sample_mode='reservoir',
@@ -49,16 +50,16 @@ def run(version, spec, env, settings):
         return gather_sizes(n.data_path)
 
 
-def run_comparison(version1,
-                   version2,
-                   spec,
-                   result_hosts,
-                   env_v1,
-                   env_v2,
-                   settings_v1,
-                   settings_v2):
-    v1 = run(version1, spec, env_v1, settings_v1)
-    v2 = run(version2, spec, env_v2, settings_v2)
+async def run_comparison(version1,
+                         version2,
+                         spec,
+                         result_hosts,
+                         env_v1,
+                         env_v2,
+                         settings_v1,
+                         settings_v2):
+    v1 = await run(version1, spec, env_v1, settings_v1)
+    v2 = await run(version2, spec, env_v2, settings_v2)
     print(f'Version1: {version1}')
     print(f'Version2: {version2}')
     headers = ('Description', 'Version 1', 'Unit', 'Version 2', 'Unit', 'Diff')
@@ -143,7 +144,7 @@ def main():
     settings_v2 = settings.copy()
     settings_v2.update(dict_from_kw_args(args.setting_v2))
 
-    run_comparison(
+    asyncio.run(run_comparison(
         args.v1,
         args.v2,
         args.spec,
@@ -152,7 +153,7 @@ def main():
         env_v2=env_v2,
         settings_v1=settings_v1,
         settings_v2=settings_v2
-    )
+    ))
 
 
 if __name__ == '__main__':

util.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 
-def dict_from_kw_args(args):
+def dict_from_kw_args(args: list[str]) -> dict[str, str]:
     """ Return a dictionary based on ['key=val'] entries
 
     >>> dict_from_kw_args(['x=10', 'y=foo'])