From 6993387a2a0d9c178d32137b2824ffed285ae80b Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Thu, 15 Oct 2026 23:48:11 +0400 Subject: [PATCH 1/2] Add an incremental event-batch reducer equivalent to full materialization The collector re-materializes a run from every stored event whenever a batch arrives, so ingestion cost grows with the run's history. This adds the reducer that lets the store fold a batch onto the run's stored columns instead. The fold is exact because of how the materializer's columns behave. updated_at is a maximum over every event, and started_at is a minimum over the final producer's events. Every other column is written by the latest event that sets it, from that event alone. So new events that follow the active producer's stored events fold directly onto the stored columns. Two cases need targeted reads: - an event delivered after later events of its producer replays only that producer's stored events from it onward; - an earlier producer's first event reads the active producer's earliest timestamp, once per run. materialize_run now shares the fold loop, so both paths apply events and producer resets identically. The store is unchanged in this commit. Tests replay scripted and seeded delivery schedules covering every event kind, out-of-order and duplicate deliveries, and restarted producers. After each batch they compare the folded columns with full materialization, and they count the stored events read. Co-Authored-By: Claude Opus 5.5 --- .../telemetry_collector/storage.py | 169 +++++- telemetry-service/tests/event_schedules.py | 538 ++++++++++++++++++ .../tests/test_incremental_batch_apply.py | 307 ++++++++++ 3 files changed, 994 insertions(+), 20 deletions(-) create mode 100644 telemetry-service/tests/event_schedules.py create mode 100644 telemetry-service/tests/test_incremental_batch_apply.py diff --git a/telemetry-service/telemetry_collector/storage.py b/telemetry-service/telemetry_collector/storage.py index c0bd62a..e8ffd84 100644 --- a/telemetry-service/telemetry_collector/storage.py +++ b/telemetry-service/telemetry_collector/storage.py @@ -4,6 +4,7 @@ from collections.abc import Iterable from copy import deepcopy +from dataclasses import dataclass from datetime import UTC, datetime from typing import Any, Protocol @@ -109,6 +110,14 @@ def canonical_events(events: Iterable[dict[str, Any]]) -> list[dict[str, Any]]: def _apply_event(run: dict[str, Any], event: dict[str, Any]) -> None: + """Apply one event to the run state. + + ``apply_event_batch`` folds batches onto stored columns, which relies on + ``updated_at`` staying a maximum, ``started_at`` a minimum, and every other + write depending only on the event. A rule change does not revise the + stored columns of runs already in progress. + """ + timestamp = event["timestamp"] run["started_at"] = min(run["started_at"], timestamp) run["updated_at"] = max( @@ -184,27 +193,14 @@ def _apply_event(run: dict[str, Any], event: dict[str, Any]) -> None: run["block"] = None -def materialize_run( - registered_run: dict[str, Any], events: Iterable[dict[str, Any]] -) -> dict[str, Any]: - """Rebuild mutable run state from a stable canonical event order.""" +def _fold_events( + run: dict[str, Any], + active_producer_id: str | None, + ordered_events: Iterable[dict[str, Any]], +) -> None: + """Apply ordered events that follow ``active_producer_id``'s events.""" - run = { - **registered_run, - "status": "running", - "current_stage": "created", - "started_at": registered_run["registered_at"], - "updated_at": registered_run["registered_at"], - "ended_at": None, - "heartbeat_at": None, - "resources": None, - "work": None, - "failure": None, - # Rebuilt from the events on every read; not a stored column. - "block": None, - } - active_producer_id: str | None = None - for event in canonical_events(events): + for event in ordered_events: producer_id = event["producer_id"] if active_producer_id is not None and producer_id != active_producer_id: # A producer id identifies one execution attempt. Starting a later @@ -230,9 +226,142 @@ def materialize_run( ) active_producer_id = producer_id _apply_event(run, event) + + +def materialize_run( + registered_run: dict[str, Any], events: Iterable[dict[str, Any]] +) -> dict[str, Any]: + """Rebuild mutable run state from a stable canonical event order.""" + + run = { + **registered_run, + "status": "running", + "current_stage": "created", + "started_at": registered_run["registered_at"], + "updated_at": registered_run["registered_at"], + "ended_at": None, + "heartbeat_at": None, + "resources": None, + "work": None, + "failure": None, + # Rebuilt from the events on every read; not a stored column. + "block": None, + } + _fold_events(run, None, canonical_events(events)) return run +MATERIALIZED_RUN_COLUMNS = ( + "status", + "current_stage", + "started_at", + "updated_at", + "ended_at", + "heartbeat_at", + "resources", + "work", + "failure", +) + + +@dataclass(frozen=True) +class StoredProducer: + """A registered producer and its highest sequence stored before a batch.""" + + producer_id: str + registered_at: datetime + last_sequence: int | None + + +class StoredEventReader(Protocol): + """Targeted reads of one producer's stored events, including the batch's.""" + + def events_from(self, producer_id: str, sequence: int) -> list[dict[str, Any]]: + """Return the producer's events from ``sequence`` onward, in order.""" + ... + + def earliest_timestamp(self, producer_id: str) -> datetime: + """Return the earliest timestamp among the producer's events.""" + ... + + +def _materialized_columns(run: dict[str, Any]) -> dict[str, Any]: + return {column: run[column] for column in MATERIALIZED_RUN_COLUMNS} + + +def apply_event_batch( + run: dict[str, Any], + producers: Iterable[StoredProducer], + accepted_events: Iterable[dict[str, Any]], + stored_events: StoredEventReader, +) -> dict[str, Any]: + """Return the materialized columns after a batch, without replaying the run. + + ``run`` holds the columns materialized before the batch, ``producers`` each + producer's highest sequence stored before it, and ``accepted_events`` its + newly stored events. The result equals ``materialize_run`` over every stored event. + ``updated_at`` is a maximum over every event and ``started_at`` a minimum + over the final producer's events; every other column is set by the latest + event that writes it. Folding new events onto the stored columns is + therefore exact when they follow the active producer's stored events. + Only an event delivered after later events of its producer replays that + producer's stored events from it onward. + """ + + ordered_events = canonical_events(accepted_events) + reported_producers = [ + producer for producer in producers if producer.last_sequence is not None + ] + if not reported_producers: + return _materialized_columns(materialize_run(run, ordered_events)) + active_producer = max( + reported_producers, + key=lambda producer: (producer.registered_at, producer.producer_id), + ) + active_position = ( + active_producer.registered_at, + active_producer.producer_id, + active_producer.last_sequence, + ) + state = _materialized_columns(run) + for event in ordered_events: + state["updated_at"] = max( + state["updated_at"], event["timestamp"], event["producer_registered_at"] + ) + following_events = [ + event for event in ordered_events if _event_order(event)[:3] > active_position + ] + if any( + event["producer_id"] != active_producer.producer_id + for event in following_events + ): + # A later producer's first event resets the run, so the events it + # follows contribute only to updated_at. + _fold_events(state, active_producer.producer_id, following_events) + return _materialized_columns(state) + if len(reported_producers) == 1 and any( + event["producer_id"] != active_producer.producer_id for event in ordered_events + ): + # An earlier producer's first events now precede the active producer's, + # whose attempt therefore starts at its own earliest event rather than + # at the run's registration. + state["started_at"] = stored_events.earliest_timestamp( + active_producer.producer_id + ) + late_sequences = [ + event["sequence"] + for event in ordered_events + if event["producer_id"] == active_producer.producer_id + and event["sequence"] < active_producer.last_sequence + ] + if late_sequences: + following_events = stored_events.events_from( + active_producer.producer_id, min(late_sequences) + ) + _fold_events(state, active_producer.producer_id, following_events) + return _materialized_columns(state) + + def run_summary(run: dict[str, Any]) -> dict[str, Any]: """Return the bounded run representation used by list responses.""" diff --git a/telemetry-service/tests/event_schedules.py b/telemetry-service/tests/event_schedules.py new file mode 100644 index 0000000..86e4b0b --- /dev/null +++ b/telemetry-service/tests/event_schedules.py @@ -0,0 +1,538 @@ +"""Telemetry event kinds and delivery schedules for the batch-application tests.""" + +from __future__ import annotations + +import random +from collections.abc import Callable, Sequence +from datetime import UTC, datetime, timedelta +from typing import Any + +from telemetry_collector.models import TelemetryEvent + +RUN_REGISTERED_AT = datetime(2026, 10, 1, 12, tzinfo=UTC) + +Event = dict[str, Any] +EventKind = Callable[[str, int, float], Event] + + +def telemetry_event( + producer_id: str, + sequence: int, + minutes_after_registration: float, + *, + event_type: str, + status: str, + stage_id: str | None = None, + details: dict[str, Any] | None = None, + resources: dict[str, Any] | None = None, + message: str | None = None, + event_id: str | None = None, +) -> Event: + """Return an event as a producer sends it, validated by the collector model.""" + + event = TelemetryEvent( + schema_version=1, + event_id=event_id or f"{producer_id}-{sequence}", + run_id="run-1", + producer_id=producer_id, + sequence=sequence, + timestamp=RUN_REGISTERED_AT + timedelta(minutes=minutes_after_registration), + event_type=event_type, + stage_id=stage_id, + status=status, + message=message, + details=details or {}, + resources=resources, + ).model_dump(mode="python") + del event["schema_version"] + return event + + +def resource_snapshot(sequence: int) -> dict[str, Any]: + return { + "cpu_user_seconds": sequence * 1.5, + "rss_bytes": sequence * 1024, + "peak_rss_bytes": sequence * 2048, + } + + +def run_started(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="run", + status="started", + stage_id="created", + message="Run started", + ) + + +def stage_started(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="stage", + status="started", + stage_id="calibrate", + ) + + +def stage_completed(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="stage", + status="completed", + stage_id="export", + ) + + +def stage_failed_without_failing_the_run( + producer_id: str, sequence: int, minutes: float +) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="stage", + status="failed", + stage_id="validate", + ) + + +def progress_with_totals(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="progress", + status="progress", + stage_id="calibrate", + details={ + "done": sequence, + "total": 500, + "unit": "epochs", + "elapsed_seconds": sequence * 2.5, + "loss": 0.125, + }, + ) + + +def progress_without_totals(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="progress", + status="progress", + stage_id="calibrate", + details={"done": sequence}, + ) + + +def calibration_epoch(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="calibration", + status="progress", + stage_id="calibrate", + details={"epoch": sequence, "loss": 0.5 / sequence}, + ) + + +def heartbeat(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="heartbeat", + status="progress", + resources=resource_snapshot(sequence), + ) + + +def stage_resource_snapshot(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="stage", + status="progress", + stage_id="export", + resources=resource_snapshot(sequence), + ) + + +def run_failed(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="run", + status="failed", + stage_id="calibrate", + message="Calibration diverged", + details={ + "error_type": "FloatingPointError", + "error_code": "DIVERGED", + "failure_class": "build", + "failed_during": "calibrating", + }, + ) + + +def run_failed_without_details( + producer_id: str, sequence: int, minutes: float +) -> Event: + return telemetry_event( + producer_id, sequence, minutes, event_type="run", status="failed" + ) + + +def failed_stage(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="stage", + status="failed", + stage_id="failed", + message="Interrupted", + details={"error_type": "KeyboardInterrupt", "error_code": "INTERRUPTED"}, + ) + + +def run_blocked(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="run", + status="blocked", + details={ + "phase": "terminal", + "blocking_failure_count": 1, + "blocking_gate_ids": ["uk_target_fit"], + "gate_statuses": {"uk_target_fit": "failed"}, + }, + ) + + +def blocked_stage(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="stage", + status="blocked", + stage_id="blocked", + details={ + "phase": "preflight", + "blocking_failure_count": 2, + "blocking_gate_ids": [], + }, + ) + + +def blocked_status_on_another_stage( + producer_id: str, sequence: int, minutes: float +) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="stage", + status="blocked", + stage_id="gate_battery", + details={ + "phase": "terminal", + "blocking_failure_count": 3, + "blocking_gate_ids": ["us_target_fit"], + }, + ) + + +def run_completed(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, sequence, minutes, event_type="run", status="completed" + ) + + +def complete_stage(producer_id: str, sequence: int, minutes: float) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="stage", + status="completed", + stage_id="complete", + ) + + +def complete_stage_id_without_completion( + producer_id: str, sequence: int, minutes: float +) -> Event: + return telemetry_event( + producer_id, + sequence, + minutes, + event_type="stage", + status="started", + stage_id="complete", + ) + + +# Every branch of the materializer: reopening, stage tracking, work, heartbeats, +# resources, failure, block and completion, and terminal-looking events that +# don't end the run. +EVENT_KINDS: dict[str, EventKind] = { + kind.__name__: kind + for kind in ( + run_started, + stage_started, + stage_completed, + stage_failed_without_failing_the_run, + progress_with_totals, + progress_without_totals, + calibration_epoch, + heartbeat, + stage_resource_snapshot, + run_failed, + run_failed_without_details, + failed_stage, + run_blocked, + blocked_stage, + blocked_status_on_another_stage, + run_completed, + complete_stage, + complete_stage_id_without_completion, + ) +} + + +def every_kind_stream( + producer_id: str, + kinds: Sequence[EventKind] = tuple(EVENT_KINDS.values()), + first_minute: float = 1.0, +) -> list[Event]: + return [ + kind(producer_id, sequence, first_minute + sequence) + for sequence, kind in enumerate(kinds, start=1) + ] + + +def long_stream( + producer_id: str, event_count: int, first_minute: float = 1.0 +) -> list[Event]: + kinds = list(EVENT_KINDS.values()) + return [ + kinds[(sequence - 1) % len(kinds)]( + producer_id, sequence, first_minute + sequence / 10 + ) + for sequence in range(1, event_count + 1) + ] + + +def in_batches(events: Sequence[Event], batch_size: int) -> list[list[Event]]: + return [ + list(events[start : start + batch_size]) + for start in range(0, len(events), batch_size) + ] + + +ONE_PRODUCER = {"producer-a": RUN_REGISTERED_AT} +RESTARTED_PRODUCERS = { + "producer-a": RUN_REGISTERED_AT, + "producer-b": RUN_REGISTERED_AT + timedelta(minutes=60), + "producer-c": RUN_REGISTERED_AT + timedelta(minutes=60), +} + + +def out_of_order_schedules(shuffle_count: int = 20) -> dict[str, list[list[Event]]]: + """One producer's events of every kind, delivered out of order.""" + + stream = every_kind_stream("producer-a") + every_kind_stream( + "producer-a", first_minute=100.0 + ) + stream = [ + {**event, "sequence": sequence, "event_id": f"producer-a-{sequence}"} + for sequence, event in enumerate(stream, start=1) + ] + batches = in_batches(stream, 4) + schedules = { + "batches in reverse": list(reversed(batches)), + "events in reverse": [[event] for event in reversed(stream)], + "first event last": [stream[1:], [stream[0]]], + "middle batch last": [batches[0], *batches[2:], batches[1]], + "events reversed inside each batch": [ + list(reversed(batch)) for batch in batches + ], + } + for shuffle_seed in range(shuffle_count): + shuffled = list(stream) + random.Random(shuffle_seed).shuffle(shuffled) + schedules[f"shuffled single events {shuffle_seed}"] = [ + [event] for event in shuffled + ] + schedules[f"shuffled batches of five {shuffle_seed}"] = in_batches(shuffled, 5) + + skew = random.Random(99) + skewed_stream = [ + { + **event, + "timestamp": RUN_REGISTERED_AT + timedelta(minutes=skew.uniform(-30, 300)), + } + for event in stream + ] + schedules["skewed clocks, in order"] = in_batches(skewed_stream, 3) + schedules["skewed clocks, batches in reverse"] = list( + reversed(in_batches(skewed_stream, 3)) + ) + return schedules + + +def restarted_producer_schedules() -> dict[str, list[list[Event]]]: + """Attempts registered as ``RESTARTED_PRODUCERS``, delivered in several orders.""" + + first_attempt = every_kind_stream("producer-a") + second_attempt = every_kind_stream("producer-b", first_minute=61.0) + third_attempt = every_kind_stream("producer-c", first_minute=62.0) + first_without_fifth = first_attempt[:4] + first_attempt[5:] + mixed_batch = first_attempt[:6] + second_attempt[:6] + random.Random(7).shuffle(mixed_batch) + early_clock_attempt = [ + {**event, "timestamp": RUN_REGISTERED_AT - timedelta(minutes=sequence)} + for sequence, event in enumerate(second_attempt, start=1) + ] + + schedules = { + "restart after the first attempt ends": [ + *in_batches(first_attempt, 5), + *in_batches(second_attempt, 5), + ], + "later attempt reports before the earlier one": [ + *in_batches(second_attempt, 5), + *in_batches(first_attempt, 5), + ], + "earlier attempt reports late": [ + *in_batches(first_attempt[:8], 4), + *in_batches(second_attempt, 4), + *in_batches(first_attempt[8:], 4), + ], + "one batch spans both attempts": [ + mixed_batch, + first_attempt[6:], + second_attempt[6:], + ], + "late event of the earlier attempt beside the later attempt": [ + first_without_fifth, + [first_attempt[4], *second_attempt[:3]], + second_attempt[3:], + ], + "earlier attempt's late event after the later attempt reports": [ + first_without_fifth, + second_attempt[:3], + [first_attempt[4]], + ], + "three attempts with tied registrations": [ + *in_batches(first_attempt, 6), + *[ + batch + for pair in zip( + in_batches(third_attempt, 3), + in_batches(second_attempt, 3), + strict=True, + ) + for batch in pair + ], + ], + "later attempt clocks before the registration": [ + *in_batches(first_attempt, 5), + *in_batches(early_clock_attempt, 5), + ], + "earlier attempt's last event carries the latest clock": [ + *in_batches(first_attempt, 5), + *in_batches(second_attempt, 5), + [heartbeat("producer-a", len(first_attempt) + 1, 1000.0)], + ], + } + return schedules + + +def random_delivery_schedule( + rng: random.Random, +) -> tuple[dict[str, datetime], list[list[Event]], set[str]]: + """Generate events of random kinds and deliver them out of order, with re-sends.""" + + producer_ids = ["producer-a", "producer-b", "producer-c"][: rng.randint(1, 3)] + registrations = {"producer-a": RUN_REGISTERED_AT} + for producer_id in producer_ids[1:]: + registrations[producer_id] = RUN_REGISTERED_AT + timedelta( + minutes=rng.choice([0, 30, 30, 90]) + ) + kinds_used: set[str] = set() + batches_by_producer: list[list[list[Event]]] = [] + for producer_id in producer_ids: + stream = [] + for sequence in range(1, rng.randint(1, 30) + 1): + kind_name = rng.choice(list(EVENT_KINDS)) + kinds_used.add(kind_name) + stream.append( + EVENT_KINDS[kind_name](producer_id, sequence, rng.uniform(-30, 300)) + ) + batches = [] + while stream: + batch_size = rng.randint(1, 8) + batches.append(stream[:batch_size]) + stream = stream[batch_size:] + batches_by_producer.append(batches) + + schedule: list[list[Event]] = [] + while any(batches_by_producer): + producer_batches = rng.choice( + [batches for batches in batches_by_producer if batches] + ) + schedule.append(producer_batches.pop(0)) + + for conflict_number in range(rng.randint(0, 6)): + perturbation = rng.choice( + ["swap", "merge", "shuffle", "resend", "withhold", "conflict"] + ) + position = rng.randrange(len(schedule)) + if perturbation == "swap": + other_position = rng.randrange(len(schedule)) + schedule[position], schedule[other_position] = ( + schedule[other_position], + schedule[position], + ) + elif perturbation == "merge" and position + 1 < len(schedule): + schedule[position] = schedule[position] + schedule.pop(position + 1) + elif perturbation == "shuffle": + shuffled_batch = list(schedule[position]) + rng.shuffle(shuffled_batch) + schedule[position] = shuffled_batch + elif perturbation == "resend": + schedule.insert(rng.randrange(len(schedule) + 1), list(schedule[position])) + elif perturbation == "withhold" and len(schedule[position]) > 1: + remaining_batch = list(schedule[position]) + withheld_event = remaining_batch.pop(rng.randrange(len(remaining_batch))) + schedule[position] = remaining_batch + schedule.insert( + rng.randrange(position + 1, len(schedule) + 1), [withheld_event] + ) + elif perturbation == "conflict": + original = rng.choice(schedule[position]) + kind_name = rng.choice(list(EVENT_KINDS)) + kinds_used.add(kind_name) + conflicting_event = EVENT_KINDS[kind_name]( + original["producer_id"], original["sequence"], rng.uniform(-30, 300) + ) | { + "event_id": f"{original['event_id']}-conflict-{conflict_number}", + } + schedule.insert(rng.randrange(len(schedule) + 1), [conflicting_event]) + return registrations, schedule, kinds_used diff --git a/telemetry-service/tests/test_incremental_batch_apply.py b/telemetry-service/tests/test_incremental_batch_apply.py new file mode 100644 index 0000000..ca52f7a --- /dev/null +++ b/telemetry-service/tests/test_incremental_batch_apply.py @@ -0,0 +1,307 @@ +"""Incremental event-batch application must equal full re-materialization. + +The store folds each accepted batch onto a run's materialized columns instead +of replaying the run. These tests deliver event schedules to an in-memory run +that de-duplicates events as the store does, and after every batch compare its +columns with ``materialize_run`` over every stored event. +""" + +from __future__ import annotations + +import random +from collections.abc import Sequence +from copy import deepcopy +from datetime import datetime, timedelta +from typing import Any + +from telemetry_collector.auth import HuggingFacePrincipal +from telemetry_collector.models import RunRegistration +from telemetry_collector.storage import ( + MATERIALIZED_RUN_COLUMNS, + StoredProducer, + apply_event_batch, + initial_run_state, + materialize_run, +) +from tests.event_schedules import ( + EVENT_KINDS, + ONE_PRODUCER, + RESTARTED_PRODUCERS, + RUN_REGISTERED_AT, + Event, + complete_stage, + every_kind_stream, + heartbeat, + in_batches, + long_stream, + out_of_order_schedules, + progress_with_totals, + random_delivery_schedule, + restarted_producer_schedules, + run_completed, + run_failed, +) + + +class CountingStoredEvents: + """In-memory stored-event reader that counts what the reducer reads.""" + + def __init__(self, stored_events: dict[str, Event]) -> None: + self._stored_events = stored_events + self.events_read = 0 + self.earliest_timestamp_reads = 0 + + def events_from(self, producer_id: str, sequence: int) -> list[Event]: + selected = sorted( + ( + event + for event in self._stored_events.values() + if event["producer_id"] == producer_id and event["sequence"] >= sequence + ), + key=lambda event: (event["sequence"], event["event_id"]), + ) + self.events_read += len(selected) + return deepcopy(selected) + + def earliest_timestamp(self, producer_id: str) -> datetime: + self.earliest_timestamp_reads += 1 + return min( + event["timestamp"] + for event in self._stored_events.values() + if event["producer_id"] == producer_id + ) + + +class IncrementalRun: + """One run stored as the collector stores it, folding each accepted batch.""" + + def __init__(self, producer_registrations: dict[str, datetime]) -> None: + self.state = initial_run_state( + RunRegistration( + run_id="run-1", + producer_id="producer-a", + country_code="UK", + pipeline="uk-calibration", + ), + HuggingFacePrincipal( + user_id="owner-1", + username="builder", + organizations=("policyengine",), + ), + RUN_REGISTERED_AT, + ) + self.producer_registrations = producer_registrations + self.stored_events: dict[str, Event] = {} + self.reader = CountingStoredEvents(self.stored_events) + + def deliver(self, batch: Sequence[Event]) -> tuple[int, int]: + """Store the batch's new events and fold them; return (accepted, duplicates).""" + + producers = [ + StoredProducer( + producer_id=producer_id, + registered_at=registered_at, + last_sequence=max( + ( + event["sequence"] + for event in self.stored_events.values() + if event["producer_id"] == producer_id + ), + default=None, + ), + ) + for producer_id, registered_at in self.producer_registrations.items() + ] + stored_sequences = { + (event["producer_id"], event["sequence"]) + for event in self.stored_events.values() + } + accepted_events = [] + for event in batch: + sequence_key = (event["producer_id"], event["sequence"]) + if event["event_id"] in self.stored_events or ( + sequence_key in stored_sequences + ): + continue + stored_event = { + **deepcopy(event), + "producer_registered_at": self.producer_registrations[ + event["producer_id"] + ], + } + self.stored_events[event["event_id"]] = stored_event + stored_sequences.add(sequence_key) + accepted_events.append(deepcopy(stored_event)) + if accepted_events: + self.state.update( + apply_event_batch(self.state, producers, accepted_events, self.reader) + ) + return len(accepted_events), len(batch) - len(accepted_events) + + def stored_columns(self) -> dict[str, Any]: + return {column: self.state[column] for column in MATERIALIZED_RUN_COLUMNS} + + def fully_materialized_columns(self) -> dict[str, Any]: + materialized = materialize_run(self.state, self.stored_events.values()) + return {column: materialized[column] for column in MATERIALIZED_RUN_COLUMNS} + + +def deliver_and_compare( + producer_registrations: dict[str, datetime], + schedule: Sequence[Sequence[Event]], + description: str, +) -> IncrementalRun: + """Deliver every batch, checking the stored columns after each one.""" + + run = IncrementalRun(producer_registrations) + for batch_number, batch in enumerate(schedule, start=1): + run.deliver(batch) + assert run.stored_columns() == run.fully_materialized_columns(), ( + f"{description}: batch {batch_number} of {len(schedule)}" + ) + return run + + +def test_incremental_batch_apply_matches_full_materialization_for_every_event_kind() -> ( + None +): + kind_orders = [list(EVENT_KINDS.values())] + for order_seed in range(8): + shuffled_kinds = list(EVENT_KINDS.values()) + random.Random(order_seed).shuffle(shuffled_kinds) + kind_orders.append(shuffled_kinds) + + for order_number, kinds in enumerate(kind_orders): + stream = every_kind_stream("producer-a", kinds) + for batch_size in (1, 2, 3, 5, len(stream)): + deliver_and_compare( + ONE_PRODUCER, + in_batches(stream, batch_size), + f"kind order {order_number}, batches of {batch_size}", + ) + + +def test_incremental_batch_apply_matches_full_materialization_for_out_of_order_batches() -> ( + None +): + for description, schedule in out_of_order_schedules().items(): + deliver_and_compare(ONE_PRODUCER, schedule, description) + + +def test_incremental_batch_apply_matches_full_materialization_for_duplicate_resends() -> ( + None +): + stream = every_kind_stream("producer-a") + run = IncrementalRun(ONE_PRODUCER) + + for batch in in_batches(stream, 4): + assert run.deliver(batch) == (len(batch), 0) + assert run.stored_columns() == run.fully_materialized_columns() + before_resend = run.stored_columns() + assert run.deliver(batch) == (0, len(batch)) + assert run.stored_columns() == before_resend + + assert run.deliver(stream) == (0, len(stream)) + conflicting_sequence = run_failed("producer-a", 3, 500.0) | { + "event_id": "producer-a-3-resent-with-another-id" + } + assert run.deliver([conflicting_sequence]) == (0, 1) + assert run.stored_columns() == run.fully_materialized_columns() + + next_sequence = len(stream) + 1 + new_events = [ + progress_with_totals("producer-a", next_sequence, 600.0), + heartbeat("producer-a", next_sequence + 1, 601.0), + ] + overlapping_batch = [stream[-1], new_events[0], new_events[0], new_events[1]] + assert run.deliver(overlapping_batch) == (2, 2) + assert run.stored_columns() == run.fully_materialized_columns() + + late_resend_with_late_event = [ + *stream[:3], + complete_stage("producer-a", next_sequence + 3, 602.0), + run_completed("producer-a", next_sequence + 2, 603.0), + ] + assert run.deliver(late_resend_with_late_event) == (2, 3) + assert run.stored_columns() == run.fully_materialized_columns() + + +def test_incremental_batch_apply_matches_full_materialization_across_restarted_producers() -> ( + None +): + for description, schedule in restarted_producer_schedules().items(): + deliver_and_compare(RESTARTED_PRODUCERS, schedule, description) + + +def test_incremental_batch_apply_matches_full_materialization_for_random_delivery_schedules() -> ( + None +): + kinds_covered: set[str] = set() + for schedule_seed in range(300): + registrations, schedule, kinds_used = random_delivery_schedule( + random.Random(schedule_seed) + ) + kinds_covered |= kinds_used + deliver_and_compare(registrations, schedule, f"schedule seed {schedule_seed}") + + assert kinds_covered == set(EVENT_KINDS) + + +def test_incremental_batch_apply_reads_no_stored_events_for_in_order_batches() -> None: + first_attempt = long_stream("producer-a", 600) + second_attempt = long_stream("producer-b", 600, first_minute=61.0) + + run = deliver_and_compare( + RESTARTED_PRODUCERS, + [*in_batches(first_attempt, 20), *in_batches(second_attempt, 20)], + "two in-order attempts", + ) + + assert run.reader.events_read == 0 + assert run.reader.earliest_timestamp_reads == 0 + + +def test_incremental_batch_apply_reads_only_the_events_after_a_late_event() -> None: + events_read_by_run_length = {} + for run_length in (60, 600): + stream = long_stream("producer-a", run_length) + late_event = stream[-6] + run = deliver_and_compare( + ONE_PRODUCER, + in_batches([event for event in stream if event is not late_event], 20), + f"{run_length} events with one withheld", + ) + assert run.reader.events_read == 0 + + assert run.deliver([late_event]) == (1, 0) + + assert run.stored_columns() == run.fully_materialized_columns() + events_read_by_run_length[run_length] = run.reader.events_read + + # The late event and the five stored after it, however long the run is. + assert events_read_by_run_length == {60: 6, 600: 6} + + +def test_incremental_batch_apply_reads_the_earliest_timestamp_once_when_an_earlier_producer_first_reports() -> ( + None +): + registrations = { + "producer-a": RUN_REGISTERED_AT, + "producer-b": RUN_REGISTERED_AT + timedelta(minutes=30), + } + later_attempt = long_stream("producer-b", 300, first_minute=40.0) + earlier_attempt = long_stream("producer-a", 30) + run = deliver_and_compare( + registrations, in_batches(later_attempt, 20), "later attempt first" + ) + assert run.state["started_at"] == RUN_REGISTERED_AT + + for batch in in_batches(earlier_attempt, 5): + run.deliver(batch) + assert run.stored_columns() == run.fully_materialized_columns() + + # Once an earlier attempt has events, the later one's segment starts at its + # own earliest event rather than at the run's registration. + assert run.state["started_at"] == min(event["timestamp"] for event in later_attempt) + assert run.reader.earliest_timestamp_reads == 1 + assert run.reader.events_read == 0 From e8815473cdfe2429c66ae9f4214fd7c3befb12dd Mon Sep 17 00:00:00 2001 From: Anthony Volk <14987227+anth-volk@users.noreply.github.com> Date: Fri, 16 Oct 2026 00:04:41 +0400 Subject: [PATCH 2/2] Apply each stored event batch incrementally in the Postgres store append_events re-read every stored event of the run, joined to its producer, and re-materialized the run whenever a batch added an event. Each batch therefore cost time proportional to the run's history, which is O(n^2) over a run. It also held the run's row lock for longer as the build went on. The store now reads each producer's highest stored sequence with one query: an index-only max per producer. It inserts the batch with RETURNING, so the reducer sees the same database values (time zone, JSONB) that a full re-read would. It then folds the inserted events onto the stored columns with apply_event_batch. A producer delivers in sequence order, so a normal batch reads no stored events. Two cases read more: - a late event reads only its producer's events from it onward; - an earlier producer's first event reads one minimum timestamp, at most once per run. The stored columns, the database contents and the HTTP responses are unchanged, and there is no schema change. get_run still materializes from every event, because it returns them all. Tests count the event rows each batch reads: constant per in-order batch, and only those after a late event, whatever the run's length. They also check, after every batch of scripted and seeded schedules, that the stored row equals the full materialization over the stored events, and that the run listing matches it. Co-Authored-By: Claude Opus 5.5 --- .../telemetry_collector/storage.py | 147 ++++++++----- .../tests/test_storage_postgres.py | 204 ++++++++++++++++++ 2 files changed, 303 insertions(+), 48 deletions(-) diff --git a/telemetry-service/telemetry_collector/storage.py b/telemetry-service/telemetry_collector/storage.py index e8ffd84..2aa3b93 100644 --- a/telemetry-service/telemetry_collector/storage.py +++ b/telemetry-service/telemetry_collector/storage.py @@ -8,7 +8,7 @@ from datetime import UTC, datetime from typing import Any, Protocol -from sqlalchemy import and_, select, tuple_ +from sqlalchemy import Select, and_, func, select, tuple_ from sqlalchemy.dialects.postgresql import insert from sqlalchemy.engine import Engine from sqlalchemy.exc import SQLAlchemyError @@ -497,6 +497,52 @@ def _event_mapping( } +def _stored_events_statement( + run_id: str, +) -> Select[tuple[StoredTelemetryEvent, datetime]]: + return ( + select( + StoredTelemetryEvent, + TelemetryProducer.registered_at, + ) + .join( + TelemetryProducer, + and_( + StoredTelemetryEvent.run_id == TelemetryProducer.run_id, + StoredTelemetryEvent.producer_id == TelemetryProducer.producer_id, + ), + ) + .where(StoredTelemetryEvent.run_id == run_id) + ) + + +class _PostgresStoredEvents: + """Reads one producer's stored events inside the ingestion transaction.""" + + def __init__(self, session: Session, run_id: str) -> None: + self._session = session + self._run_id = run_id + + def events_from(self, producer_id: str, sequence: int) -> list[dict[str, Any]]: + rows = self._session.execute( + _stored_events_statement(self._run_id) + .where( + StoredTelemetryEvent.producer_id == producer_id, + StoredTelemetryEvent.sequence >= sequence, + ) + .order_by(StoredTelemetryEvent.sequence, StoredTelemetryEvent.event_id) + ).all() + return [_event_mapping(event, registered_at) for event, registered_at in rows] + + def earliest_timestamp(self, producer_id: str) -> datetime: + return self._session.execute( + select(func.min(StoredTelemetryEvent.emitted_at)).where( + StoredTelemetryEvent.run_id == self._run_id, + StoredTelemetryEvent.producer_id == producer_id, + ) + ).scalar_one() + + def _registered_metadata( run: TelemetryRun, ) -> tuple[str, str, str | None, str | None, str]: @@ -585,19 +631,7 @@ def register_run( def _events_for_run(self, session: Session, run_id: str) -> list[dict[str, Any]]: rows = session.execute( - select( - StoredTelemetryEvent, - TelemetryProducer.registered_at, - ) - .join( - TelemetryProducer, - and_( - StoredTelemetryEvent.run_id == TelemetryProducer.run_id, - StoredTelemetryEvent.producer_id == TelemetryProducer.producer_id, - ), - ) - .where(StoredTelemetryEvent.run_id == run_id) - .order_by( + _stored_events_statement(run_id).order_by( TelemetryProducer.registered_at, TelemetryProducer.producer_id, StoredTelemetryEvent.sequence, @@ -606,6 +640,29 @@ def _events_for_run(self, session: Session, run_id: str) -> list[dict[str, Any]] ).all() return [_event_mapping(event, registered_at) for event, registered_at in rows] + def _stored_producers( + self, session: Session, run_id: str + ) -> dict[str, StoredProducer]: + last_sequence = ( + select(func.max(StoredTelemetryEvent.sequence)) + .where( + StoredTelemetryEvent.run_id == TelemetryProducer.run_id, + StoredTelemetryEvent.producer_id == TelemetryProducer.producer_id, + ) + .scalar_subquery() + ) + rows = session.execute( + select( + TelemetryProducer.producer_id, + TelemetryProducer.registered_at, + last_sequence, + ).where(TelemetryProducer.run_id == run_id) + ).all() + return { + producer_id: StoredProducer(producer_id, registered_at, sequence) + for producer_id, registered_at, sequence in rows + } + def append_events( self, run_id: str, @@ -625,16 +682,8 @@ def append_events( raise PermissionError("Run belongs to another Hugging Face user.") if not models: return 0, 0 - producer_ids = {model.producer_id for model in models} - registered_producer_ids = set( - session.execute( - select(TelemetryProducer.producer_id).where( - TelemetryProducer.run_id == run_id, - TelemetryProducer.producer_id.in_(producer_ids), - ) - ).scalars() - ) - if registered_producer_ids != producer_ids: + producers = self._stored_producers(session, run_id) + if not {model.producer_id for model in models} <= producers.keys(): raise PermissionError( "Telemetry producer is not registered for this run." ) @@ -656,31 +705,33 @@ def append_events( "resources": event["resources"], } ) - accepted_ids = session.execute( - insert(StoredTelemetryEvent) - .values(values) - .on_conflict_do_nothing() - .returning(StoredTelemetryEvent.event_id) - ).scalars() - accepted = len(accepted_ids.all()) - if accepted: - session.flush() - materialized = materialize_run( - _run_mapping(run), self._events_for_run(session, run_id) + # Returning the stored rows gives the fold the same database values + # (time zone, JSON) that a full re-read would. + inserted_events = ( + session.execute( + insert(StoredTelemetryEvent) + .values(values) + .on_conflict_do_nothing() + .returning(StoredTelemetryEvent) + ) + .scalars() + .all() + ) + if inserted_events: + materialized = apply_event_batch( + _run_mapping(run), + producers.values(), + [ + _event_mapping( + event, producers[event.producer_id].registered_at + ) + for event in inserted_events + ], + _PostgresStoredEvents(session, run_id), ) - for field in ( - "status", - "current_stage", - "started_at", - "updated_at", - "ended_at", - "heartbeat_at", - "resources", - "work", - "failure", - ): - setattr(run, field, materialized[field]) - return accepted, len(models) - accepted + for column, value in materialized.items(): + setattr(run, column, value) + return len(inserted_events), len(models) - len(inserted_events) def list_runs( self, diff --git a/telemetry-service/tests/test_storage_postgres.py b/telemetry-service/tests/test_storage_postgres.py index 5281fef..d7e73ef 100644 --- a/telemetry-service/tests/test_storage_postgres.py +++ b/telemetry-service/tests/test_storage_postgres.py @@ -1,15 +1,37 @@ from __future__ import annotations +import random +from collections.abc import Iterator, Sequence +from contextlib import contextmanager from datetime import UTC, datetime, timedelta +from typing import Any import pytest +from sqlalchemy import event as sqlalchemy_event +from telemetry_collector import storage from telemetry_collector.auth import HuggingFacePrincipal +from telemetry_collector.database import TelemetryRun, create_session_factory from telemetry_collector.migrate import upgrade_database from telemetry_collector.models import RunRegistration, TelemetryEvent from telemetry_collector.storage import ( + MATERIALIZED_RUN_COLUMNS, PostgresTelemetryStore, RunRegistrationConflictError, + materialize_run, + run_summary, +) +from tests.event_schedules import ( + ONE_PRODUCER, + RESTARTED_PRODUCERS, + Event, + every_kind_stream, + in_batches, + long_stream, + out_of_order_schedules, + random_delivery_schedule, + restarted_producer_schedules, + run_failed, ) @@ -146,3 +168,185 @@ def test_store_readiness_requires_alembic_head(postgres_url: str) -> None: assert store.is_ready() is True store.close() + + +def _register_producers( + store: PostgresTelemetryStore, + run_id: str, + registrations: dict[str, datetime], + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Register each producer at its stated time; the first registers the run.""" + + for producer_id, registered_at in registrations.items(): + monkeypatch.setattr( + storage, "_now", lambda registered_at=registered_at: registered_at + ) + store.register_run( + _registration(run_id=run_id, producer_id=producer_id), _principal() + ) + + +def _for_run(batch: Sequence[Event], run_id: str) -> list[TelemetryEvent]: + """Address scheduled events to one run, with event ids unique across runs.""" + + return [ + TelemetryEvent( + schema_version=1, + **{ + **event, + "run_id": run_id, + "event_id": f"{run_id}.{event['event_id']}", + }, + ) + for event in batch + ] + + +def _assert_stored_state_is_fully_materialized( + store: PostgresTelemetryStore, run_id: str, context: str +) -> dict[str, Any]: + """Compare the stored run with a re-materialization from every stored event.""" + + with create_session_factory(store.engine)() as session: + run = session.get(TelemetryRun, run_id) + stored_columns = { + column: getattr(run, column) for column in MATERIALIZED_RUN_COLUMNS + } + # What ingestion stored before batches were applied incrementally. + expected = materialize_run( + storage._run_mapping(run), store._events_for_run(session, run_id) + ) + assert stored_columns == { + column: expected[column] for column in MATERIALIZED_RUN_COLUMNS + }, context + return expected + + +def _assert_listed_as_fully_materialized( + store: PostgresTelemetryStore, run_id: str, expected: dict[str, Any] +) -> None: + listed = store.list_runs(country="US", limit=200, before=None) + assert [summary for summary in listed if summary["run_id"] == run_id] == [ + run_summary(expected) + ], run_id + + +@contextmanager +def _event_rows_read(store: PostgresTelemetryStore) -> Iterator[list[int]]: + """Record how many rows each query reading telemetry_events returns.""" + + row_counts: list[int] = [] + + def record_rows(connection, cursor, statement, parameters, context, executemany): + if statement.lstrip().upper().startswith("SELECT") and ( + "telemetry_events" in statement + ): + row_counts.append(cursor.rowcount) + + sqlalchemy_event.listen(store.engine, "after_cursor_execute", record_rows) + try: + yield row_counts + finally: + sqlalchemy_event.remove(store.engine, "after_cursor_execute", record_rows) + + +def test_incremental_batch_apply_reads_a_constant_number_of_event_rows_per_in_order_batch( + postgres_url: str, monkeypatch: pytest.MonkeyPatch +) -> None: + upgrade_database(postgres_url) + store = PostgresTelemetryStore(postgres_url) + _register_producers(store, "run-1", ONE_PRODUCER, monkeypatch) + batches = in_batches(long_stream("producer-a", 400), 20) + + rows_read_per_batch = [] + for batch in batches: + with _event_rows_read(store) as row_counts: + assert store.append_events( + "run-1", "owner-1", _for_run(batch, "run-1") + ) == (len(batch), 0) + rows_read_per_batch.append(sum(row_counts)) + expected = _assert_stored_state_is_fully_materialized(store, "run-1", "400 events") + _assert_listed_as_fully_materialized(store, "run-1", expected) + store.close() + + # Re-reading the run would read 20 rows for the first batch and 400 for the last. + assert len(set(rows_read_per_batch)) == 1, rows_read_per_batch + assert rows_read_per_batch[0] <= len(batches[0]) + + +def test_incremental_batch_apply_reads_only_the_event_rows_after_a_late_event( + postgres_url: str, monkeypatch: pytest.MonkeyPatch +) -> None: + upgrade_database(postgres_url) + store = PostgresTelemetryStore(postgres_url) + + rows_read_by_run_length = {} + for run_length in (60, 600): + run_id = f"run-{run_length}" + _register_producers(store, run_id, ONE_PRODUCER, monkeypatch) + stream = long_stream("producer-a", run_length) + late_event = stream[-6] + for batch in in_batches( + [event for event in stream if event is not late_event], 20 + ): + store.append_events(run_id, "owner-1", _for_run(batch, run_id)) + + with _event_rows_read(store) as row_counts: + assert store.append_events( + run_id, "owner-1", _for_run([late_event], run_id) + ) == (1, 0) + + rows_read_by_run_length[run_length] = sum(row_counts) + expected = _assert_stored_state_is_fully_materialized( + store, run_id, f"{run_length} events with a late one" + ) + _assert_listed_as_fully_materialized(store, run_id, expected) + store.close() + + # One producer row, then the late event and the five stored after it. + assert rows_read_by_run_length == {60: 7, 600: 7} + + +def test_incremental_batch_apply_stores_the_fully_materialized_state( + postgres_url: str, monkeypatch: pytest.MonkeyPatch +) -> None: + upgrade_database(postgres_url) + store = PostgresTelemetryStore(postgres_url) + every_kind = every_kind_stream("producer-a") + deliveries = [ + ("every kind in batches of three", ONE_PRODUCER, in_batches(every_kind, 3)), + ( + "each batch sent twice, a conflicting sequence, then everything again", + ONE_PRODUCER, + [ + *[batch for batch in in_batches(every_kind, 4) for _ in range(2)], + [run_failed("producer-a", 3, 500.0) | {"event_id": "resent-3"}], + every_kind, + ], + ), + *[ + (description, ONE_PRODUCER, schedule) + for description, schedule in out_of_order_schedules(shuffle_count=2).items() + ], + *[ + (description, RESTARTED_PRODUCERS, schedule) + for description, schedule in restarted_producer_schedules().items() + ], + ] + for schedule_seed in range(20): + registrations, schedule, _ = random_delivery_schedule( + random.Random(schedule_seed) + ) + deliveries.append((f"schedule seed {schedule_seed}", registrations, schedule)) + + for run_number, (description, registrations, schedule) in enumerate(deliveries): + run_id = f"run-{run_number}" + _register_producers(store, run_id, registrations, monkeypatch) + for batch_number, batch in enumerate(schedule, start=1): + store.append_events(run_id, "owner-1", _for_run(batch, run_id)) + expected = _assert_stored_state_is_fully_materialized( + store, run_id, f"{description}: batch {batch_number}" + ) + _assert_listed_as_fully_materialized(store, run_id, expected) + store.close()