diff --git a/telemetry-service/telemetry_collector/storage.py b/telemetry-service/telemetry_collector/storage.py index c0bd62a..2aa3b93 100644 --- a/telemetry-service/telemetry_collector/storage.py +++ b/telemetry-service/telemetry_collector/storage.py @@ -4,10 +4,11 @@ from collections.abc import Iterable from copy import deepcopy +from dataclasses import dataclass 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 @@ -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.""" @@ -368,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]: @@ -456,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, @@ -477,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, @@ -496,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." ) @@ -527,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/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 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()