Save current stats in benchmarks separate from final one

iddqdex · Jan 29, 2025 · 9827af5 · 9827af5
1 parent a89d4c9
commit 9827af5
Show file tree

Hide file tree

Showing 4 changed files with 23 additions and 13 deletions.
diff --git a/ydb/apps/ydb/CHANGELOG.md b/ydb/apps/ydb/CHANGELOG.md
@@ -1,3 +1,4 @@
+* Save current stats in `ydb workload run`.
 * Added message if global timeout expiried in `ydb workload run` comamnd.
 * Fixed return code of `ydb workload run` comamnd.
 * Added statistics output on the current progress of the query in `ydb workload` command

diff --git a/ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp b/ydb/public/lib/ydb_cli/commands/ydb_benchmark.cpp
@@ -375,7 +375,7 @@ int TWorkloadCommandBenchmark::RunBench(TClient* client, NYdbWorkload::IWorkload
             settings.WithProgress = true;
 
             if (PlanFileName) {
-                settings.PlanFileName = TStringBuilder() << PlanFileName << "." << queryN << "." << ToString(i);
+                settings.PlanFileName = TStringBuilder() << PlanFileName << "." << queryN << "." << ToString(i) << ".current";
             }
 
             try {

diff --git a/ydb/tests/olap/lib/ydb_cli.py b/ydb/tests/olap/lib/ydb_cli.py
@@ -39,15 +39,17 @@ def get_cli_command() -> list[str]:
             return [cli] + args
 
     class QueryPlan:
-        def __init__(self, plan: dict | None = None, table: str | None = None, ast: str | None = None, svg: str | None = None) -> None:
-            self.plan = plan
-            self.table = table
-            self.ast = ast
-            self.svg = svg
+        def __init__(self) -> None:
+            self.plan: dict = None
+            self.table: str = None
+            self.ast: str = None
+            self.svg: str = None
+            self.stats: str = None
 
     class Iteration:
         def __init__(self):
             self.plan: Optional[YdbCliHelper.QueryPlan] = None
+            self.current_plan: Optional[YdbCliHelper.QueryPlan] = None
             self.error_message: Optional[str] = None
             self.time: Optional[float] = None
 
@@ -183,12 +185,16 @@ def _load_plan(self, name: str) -> YdbCliHelper.QueryPlan:
             if (os.path.exists(f'{pp}.svg')):
                 with open(f'{pp}.svg') as f:
                     result.svg = f.read()
+            if (os.path.exists(f'{pp}.stats')):
+                with open(f'{pp}.stats') as f:
+                    result.stats = f.read()
             return result
 
         def _load_plans(self) -> None:
             for i in range(self.iterations):
                 self._init_iter(i)
                 self.result.iterations[i].plan = self._load_plan(str(i))
+                self.result.iterations[i].current_plan = self._load_plan(f'{i}.current')
             self.result.explain.plan = self._load_plan('explain')
 
         def _load_stats(self):

diff --git a/ydb/tests/olap/load/lib/conftest.py b/ydb/tests/olap/load/lib/conftest.py
@@ -157,15 +157,17 @@ def _duration_text(duration: float | int):
             s = f'{int(duration)}s ' if duration >= 1 else ''
             return f'{s}{int(duration * 1000) % 1000}ms'
 
-        def _attach_plans(plan: YdbCliHelper.QueryPlan) -> None:
+        def _attach_plans(plan: YdbCliHelper.QueryPlan, name: str) -> None:
             if plan.plan is not None:
-                allure.attach(json.dumps(plan.plan), 'Plan json', attachment_type=allure.attachment_type.JSON)
+                allure.attach(json.dumps(plan.plan), f'{name} json', attachment_type=allure.attachment_type.JSON)
             if plan.table is not None:
-                allure.attach(plan.table, 'Plan table', attachment_type=allure.attachment_type.TEXT)
+                allure.attach(plan.table, f'{name} table', attachment_type=allure.attachment_type.TEXT)
             if plan.ast is not None:
-                allure.attach(plan.ast, 'Plan ast', attachment_type=allure.attachment_type.TEXT)
+                allure.attach(plan.ast, f'{name} ast', attachment_type=allure.attachment_type.TEXT)
             if plan.svg is not None:
-                allure.attach(plan.svg, 'Plan svg', attachment_type=allure.attachment_type.SVG)
+                allure.attach(plan.svg, f'{name} svg', attachment_type=allure.attachment_type.SVG)
+            if plan.stats is not None:
+                allure.attach(plan.stats, f'{name} stats', attachment_type=allure.attachment_type.TEXT)
 
         test = cls._test_name(query_num)
         stats = result.stats.get(test)
@@ -176,7 +178,7 @@ def _attach_plans(plan: YdbCliHelper.QueryPlan) -> None:
 
         if result.explain.plan is not None:
             with allure.step('Explain'):
-                _attach_plans(result.explain.plan)
+                _attach_plans(result.explain.plan, 'Plan')
 
         for iter_num in sorted(result.iterations.keys()):
             iter_res = result.iterations[iter_num]
@@ -185,7 +187,8 @@ def _attach_plans(plan: YdbCliHelper.QueryPlan) -> None:
                 s.params['duration'] = _duration_text(iter_res.time)
             try:
                 with s:
-                    _attach_plans(iter_res.plan)
+                    _attach_plans(iter_res.plan, 'Plan')
+                    _attach_plans(iter_res.current_plan, 'Current plan')
                     if iter_res.error_message:
                         pytest.fail(iter_res.error_message)
             except BaseException: