Fix KV SLO failure propagation and data checks

.github/workflows/slo.yml

@@ -118,11 +118,14 @@ jobs:
             "$GITHUB_WORKSPACE/sdk-current"
 
           # Build baseline: baseline SDK with the current workload runner
-          # (Dockerfile + tests/slo/), so the runner-side contract changes
-          # (entrypoint, metrics format) apply uniformly to both images.
+          # Keep its Dockerfile while copying runner-side contract changes.
+          cp "$GITHUB_WORKSPACE/sdk-baseline/tests/slo/Dockerfile" \
+             "$RUNNER_TEMP/slo-baseline.Dockerfile"
           rm -rf "$GITHUB_WORKSPACE/sdk-baseline/tests/slo"
           cp -r  "$GITHUB_WORKSPACE/sdk-current/tests/slo" \
                  "$GITHUB_WORKSPACE/sdk-baseline/tests/slo"
+          cp "$RUNNER_TEMP/slo-baseline.Dockerfile" \
+             "$GITHUB_WORKSPACE/sdk-baseline/tests/slo/Dockerfile"
           docker build \
             -f "$GITHUB_WORKSPACE/sdk-baseline/tests/slo/Dockerfile" \
             -t "ydb-app-baseline" \
@@ -136,6 +139,7 @@ jobs:
           github_token: ${{ secrets.GITHUB_TOKEN }}
           workload_name: ${{ matrix.sdk.name }}
           workload_duration: "600"
+          fail_on_workload_error: true
           workload_current_ref: ${{ github.head_ref || github.ref_name }}
           workload_current_image: ydb-app-current
           workload_current_command: ${{ matrix.sdk.command }}

test-requirements.txt

@@ -1,4 +1,5 @@
 attrs==21.2.0
+aiolimiter==1.1.0
 bcrypt==3.2.0
 black==25.11.0
 build

tests/slo/src/jobs/async_table_jobs.py

@@ -2,6 +2,7 @@
 import logging
 import time
 from random import randint
+from uuid import uuid4
 
 # `aiolimiter` is a runtime dependency (see `tests/slo/requirements.txt`).
 from aiolimiter import AsyncLimiter  # pyright: ignore[reportMissingImports]
@@ -12,7 +13,8 @@
 import ydb.aio
 
 from .base import AsyncBaseJobManager
-from .table_jobs import READ_QUERY_TEMPLATE, WRITE_QUERY_TEMPLATE
+from .table_jobs import CONTROL_READ_QUERY, CONTROL_SCHEMA, READ_QUERY_TEMPLATE, WRITE_QUERY_TEMPLATE
+from .table_jobs import check_control_result, control_row
 
 logger = logging.getLogger(__name__)
 
@@ -33,12 +35,60 @@ async def run_tests(self):
         if self.workload_type != "async-query":
             raise ValueError(f"Unsupported async workload type: {self.workload_type}")
 
+        await self._run_control_checks()
         tasks = [
             *self._run_query_read_jobs(),
             *self._run_query_write_jobs(),
             *self._run_metric_job(),
         ]
-        await asyncio.gather(*tasks)
+        await self._wait_jobs(tasks)
+
+    async def _run_control_checks(self):
+        table = f"{self.table_name}_control_{uuid4().hex}"
+        write = WRITE_QUERY_TEMPLATE.format(table).replace("Digest::NumericHash($object_id)", "0")
+        types = (
+            ydb.PrimitiveType.Uint64,
+            ydb.PrimitiveType.Utf8,
+            ydb.PrimitiveType.Double,
+            ydb.PrimitiveType.Timestamp,
+        )
+        failure = None
+        async with ydb.aio.QuerySessionPool(self.aio_driver) as pool:
+            await pool.execute_with_retries(CONTROL_SCHEMA.format(table))
+            try:
+                for generation in (1, 2):
+                    expected = control_row(generation)
+                    values = {f"${key}": value for key, value in expected.items() if key != "object_hash"}
+                    values = {key: (value, kind) for (key, value), kind in zip(values.items(), types)}
+                    for query, params, operation, timeout, checker in (
+                        (write, values, OP_TYPE_WRITE, self.args.write_timeout, None),
+                        (
+                            CONTROL_READ_QUERY.format(table),
+                            {},
+                            OP_TYPE_READ,
+                            self.args.read_timeout,
+                            lambda rows: check_control_result(rows, expected),
+                        ),
+                    ):
+                        await self._execute_query(
+                            pool,
+                            query,
+                            params,
+                            ydb.BaseRequestSettings().with_timeout(timeout / 1000),
+                            ydb.RetrySettings(idempotent=True, max_session_acquire_timeout=timeout / 1000),
+                            (operation,),
+                            check_result_cb=checker,
+                        )
+            except BaseException as error:
+                failure = error
+                raise
+            finally:
+                try:
+                    await pool.execute_with_retries(f"DROP TABLE `{table}`;")
+                except BaseException:
+                    if failure is None:
+                        raise
+                    logger.exception("Cannot clean up control table %s", table)
 
     def _run_query_read_jobs(self):
         logger.info("Start async query read jobs")
@@ -118,7 +168,7 @@ async def _run_query_writes(self, query, limiter):
 
         logger.info("Stop async query write workload")
 
-    async def _execute_query(self, pool, query, params, request_settings, retry_setting, labels):
+    async def _execute_query(self, pool, query, params, request_settings, retry_setting, labels, check_result_cb=None):
         attempt = 0
         error = None
 
@@ -127,18 +177,29 @@ async def callee(session):
             attempt += 1
 
             result_stream = await session.execute(query, parameters=params, settings=request_settings)
-            async for _ in result_stream:
-                pass
+            rows = []
+            async for part in result_stream:
+                rows.extend(part.rows)
+            return rows
 
         ts = self.metrics.start(labels)
 
         try:
-            await pool.retry_operation_async(callee, retry_setting)
+            rows = await pool.retry_operation_async(callee, retry_setting)
+            if check_result_cb:
+                check_result_cb(rows)
+            return rows
+        except asyncio.CancelledError:
+            raise
         except ydb.Error as err:
             error = err
+            self.record_failure(err)
             logger.exception("[labels: %s] Cannot retry error:", labels)
+            raise
         except BaseException as err:
             error = err
+            self.record_failure(err)
             logger.exception("[labels: %s] Unexpected error:", labels)
-
-        self.metrics.stop(labels, ts, attempts=attempt, error=error)
+            raise
+        finally:
+            self.metrics.stop(labels, ts, attempts=attempt, error=error)

tests/slo/src/jobs/base.py

@@ -17,8 +17,9 @@ class SyncRateLimiter:
     threads sharing the same instance.
     """
 
-    def __init__(self, min_interval_s: float) -> None:
+    def __init__(self, min_interval_s: float, stop_event: threading.Event | None = None) -> None:
         self._min_interval_s = max(0.0, float(min_interval_s))
+        self._stop_event = stop_event
         self._lock = threading.Lock()
         self._next_allowed_ts = 0.0  # monotonic timestamp
 
@@ -37,7 +38,10 @@ def __enter__(self):
 
             # Sleep outside the lock so other threads can observe the schedule.
             if sleep_for > 0:
-                time.sleep(sleep_for)
+                if self._stop_event is None:
+                    time.sleep(sleep_for)
+                elif self._stop_event.wait(sleep_for):
+                    return self
 
     def __exit__(self, exc_type, exc, tb) -> bool:
         # Do not suppress exceptions.
@@ -49,6 +53,26 @@ def __init__(self, driver, args, metrics):
         self.driver: ydb.Driver = driver
         self.args = args
         self.metrics = metrics
+        self._failure = None
+        self._failure_lock = threading.Lock()
+        self._stop = threading.Event()
+
+    def record_failure(self, error):
+        with self._failure_lock:
+            if self._failure is None:
+                self._failure = error
+        self._stop.set()
+
+    def raise_if_failed(self):
+        if self._failure is not None:
+            raise self._failure
+
+    def _run_worker(self, worker, *args):
+        try:
+            worker(*args)
+        except BaseException as error:
+            self.record_failure(error)
+            logger.exception("Workload worker failed")
 
     @abstractmethod
     def run_tests(self):
@@ -59,12 +83,12 @@ def _run_metric_job(self):
             return []
 
         report_period_ms = max(1, int(self.args.report_period))
-        limiter = SyncRateLimiter(min_interval_s=report_period_ms / 1000.0)
+        limiter = SyncRateLimiter(min_interval_s=report_period_ms / 1000.0, stop_event=self._stop)
 
         future = threading.Thread(
             name="slo_metrics_sender",
-            target=self._metric_sender,
-            args=(limiter, self.args.time),
+            target=self._run_worker,
+            args=(self._metric_sender, limiter, self.args.time),
         )
         future.start()
         return [future]
@@ -73,8 +97,10 @@ def _metric_sender(self, limiter, runtime):
         start_time = time.time()
         logger.info("Start push metrics")
 
-        while time.time() - start_time < runtime:
+        while not self._stop.is_set() and time.time() - start_time < runtime:
             with limiter:
+                if self._stop.is_set():
+                    break
                 self.metrics.push()
 
         logger.info("Stop push metrics")
@@ -91,6 +117,16 @@ class AsyncBaseJobManager(BaseJobManager):
     async def run_tests(self):
         pass
 
+    async def _wait_jobs(self, tasks):
+        try:
+            await asyncio.gather(*tasks)
+        except BaseException as error:
+            self.record_failure(error)
+            for task in tasks:
+                task.cancel()
+            await asyncio.gather(*tasks, return_exceptions=True)
+            self.raise_if_failed()
+
     def _run_metric_job(self):
         if not getattr(self.args, "otlp_endpoint", None):
             return []