diff --git a/backend/config/batch_pipeline.py b/backend/config/batch_pipeline.py index e56ef6ba5ca..569e7020c81 100644 --- a/backend/config/batch_pipeline.py +++ b/backend/config/batch_pipeline.py @@ -892,3 +892,64 @@ def recovery_wake_lease_seconds() -> int: if value == 0: return 0 return value if 120 <= value <= RECOVERY_WAKE_LEASE_MAX else RECOVERY_WAKE_LEASE_DEFAULT + + +# Owner-hash projection shards (utils/batch_pipeline/projection_shards.py). +PROJECTION_BACKLOG_RESERVE_DEFAULT = 2 +PROJECTION_BACKLOG_RESERVE_MAX = 16 # max_owner_drains() caps slots at 16 +PROJECTION_HEAD_DEADLINE_DEFAULT = 24 * 3600 +PROJECTION_HEAD_DEADLINE_MIN = 3600 +PROJECTION_HEAD_DEADLINE_MAX = 7 * 24 * 3600 + + +def _bounded_int(raw: str, default: int, low: int, high: int) -> int: + """Malformed or out-of-range values keep the default.""" + raw = raw.strip() + if not raw: + return default + try: + value = int(raw) + except ValueError: + return default + return value if low <= value <= high else default + + +def projection_shards_enabled() -> bool: + """Default off: one shard holder drains each owner; off restores the per-instance sweep.""" + return rollout_default_off( + 'BATCH_PROJECTION_SHARDS_ENABLED', + os.getenv('BATCH_PROJECTION_SHARDS_ENABLED', 'false').strip().lower() == 'true', + ) + + +def projection_shard_member() -> bool: + """Whether this process holds shards when shards are on. ``false`` makes it serve only.""" + return os.getenv('BATCH_PROJECTION_SHARD_MEMBER', 'true').strip().lower() != 'false' + + +def projection_backlog_reserve() -> int: + """Drain slots a shard holder keeps for backlog wakes while fresh wakes wait.""" + return _bounded_int( + os.getenv('BATCH_PROJECTION_BACKLOG_RESERVE', ''), + PROJECTION_BACKLOG_RESERVE_DEFAULT, + 0, + PROJECTION_BACKLOG_RESERVE_MAX, + ) + + +def projection_head_deadline_seconds() -> int: + """How long a failing or blocked chain head may retry before it is held (default 24h).""" + return _bounded_int( + os.getenv('BATCH_PROJECTION_HEAD_DEADLINE_SECONDS', ''), + PROJECTION_HEAD_DEADLINE_DEFAULT, + PROJECTION_HEAD_DEADLINE_MIN, + PROJECTION_HEAD_DEADLINE_MAX, + ) + + +def projection_blocked_head_deadline_enabled() -> bool: + """Default off: a chain head blocked past the deadline is held (audio kept, redrivable).""" + return rollout_default_off( + 'BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED', + os.getenv('BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED', 'false').strip().lower() == 'true', + ) diff --git a/backend/config/flag_budget.json b/backend/config/flag_budget.json index b86d632e70b..c35da4e9e87 100644 --- a/backend/config/flag_budget.json +++ b/backend/config/flag_budget.json @@ -24,7 +24,7 @@ }, "seed_review_by": "2026-11-15", "budget": 169, - "count": 259, + "count": 264, "names": { "BATCH_PIPELINE_SERVING_PROJECTION_SWEEP": {"owner": "omi-pipeline / lead2", "review_by": "2026-12-05", "registry": "config/feature-flags.yaml"}, "BATCH_PIPELINE_PROJECTION_PUMP_ENABLED": {"owner": "omi-pipeline / lead2", "review_by": "2026-12-05", "registry": "config/feature-flags.yaml"}, @@ -37,6 +37,11 @@ "PROJECTION_DRAIN_CLAIM_AWARE": {"owner": "two-lane lead", "review_by": "2026-12-05", "registry": "config/feature-flags.yaml"}, "BATCH_PIPELINE_PROJECTION_STRIPING_ENABLED": {"owner": "omi-pipeline", "review_by": "2026-12-06", "registry": "config/feature-flags.yaml"}, "K_REVISION": {"owner": "omi-pipeline", "review_by": "2026-12-06"}, + "BATCH_PROJECTION_SHARDS_ENABLED": {"owner": "two-lane lead", "review_by": "2026-12-05"}, + "BATCH_PROJECTION_SHARD_MEMBER": {"owner": "two-lane lead", "review_by": "2026-12-05"}, + "BATCH_PROJECTION_BACKLOG_RESERVE": {"owner": "two-lane lead", "review_by": "2026-12-05"}, + "BATCH_PROJECTION_HEAD_DEADLINE_SECONDS": {"owner": "two-lane lead", "review_by": "2026-12-05"}, + "BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED": {"owner": "two-lane lead", "review_by": "2026-12-05"}, "PARAKEET_SHAPE_CANARY_FILE": {"owner": "two-lane lead", "review_by": "2026-12-04"}, "PARAKEET_LID_MODEL": {"owner": "planscost", "review_by": "2026-11-30"}, "PARAKEET_LID_ONLY": {"owner": "planscost", "review_by": "2026-11-30"}, diff --git a/backend/database/batch_pipeline_store.py b/backend/database/batch_pipeline_store.py index 420721b3959..6031158785e 100644 --- a/backend/database/batch_pipeline_store.py +++ b/backend/database/batch_pipeline_store.py @@ -359,6 +359,11 @@ def projection_page( def projection_successor(self, owner_key: str, previous: str | None) -> str | None: return projection_successor(self.client, self.namespace, owner_key, previous) + def projection_page_owners( + self, limit: int = 100, *, after=None, states: tuple[str, ...] = PROJECTION_DISCOVERY_STATES + ) -> tuple[list[tuple[str, str, float, float | None]], Any]: + return projection_page_owners(self.client, self.namespace, limit, after=after, states=states) + def projection_pending(client: Any, namespace: str, limit: int = 100) -> list[str]: return projection_page(client, namespace, limit)[0] @@ -396,6 +401,45 @@ def projection_page( return [doc.id for doc in docs], cursor +def projection_page_owners( + client: Any, + namespace: str, + limit: int = 100, + *, + after=None, + states: tuple[str, ...] = PROJECTION_DISCOVERY_STATES, +) -> tuple[list[tuple[str, str, float, float | None]], Any]: + """``projection_page`` returning ``(segment_id, owner_id, committed_at, capture_at)`` per row. + + ``capture_at`` is None when the row lacks a numeric one. + + The same query, index and cursor as ``projection_page``; the owner comes + from the committed result row itself, so routing a wake costs no extra read. + """ + if not states or not set(states) <= set(PROJECTION_OBLIGATION_STATES): + raise ValueError('projection discovery states must be drain states') + collection = client.collection(f'{namespace}_results') + query = collection.where('projection_state', 'in', list(states)).order_by('committed_at').order_by('__name__') + if after is not None and after[1] == '': + query = query.start_at({'committed_at': after[0]}) + elif after is not None: + query = query.start_after({'committed_at': after[0], '__name__': collection.document(after[1])}) + size = max(1, min(limit, 100)) + docs = list(query.limit(size).stream()) + rows = [] + for doc in docs: + data = doc.to_dict() or {} + owner_id = data.get('owner_id') + if isinstance(owner_id, str) and owner_id: + capture_at = data.get('capture_at') + if not isinstance(capture_at, (int, float)) or isinstance(capture_at, bool): + capture_at = None + rows.append((doc.id, owner_id, float(data.get('committed_at') or 0.0), capture_at)) + # Same wrap rule as projection_page: a short page restarts at the head. + cursor = (docs[-1].to_dict()['committed_at'], docs[-1].id) if len(docs) == size else None + return rows, cursor + + def projection_successor(client: Any, namespace: str, owner_key: str, previous: str | None) -> str | None: """The owner's result committed directly after ``previous`` in its log. diff --git a/backend/deploy/cloud_run_gmp_sync.yaml b/backend/deploy/cloud_run_gmp_sync.yaml index d5d9f474360..c5ea8297a8d 100644 --- a/backend/deploy/cloud_run_gmp_sync.yaml +++ b/backend/deploy/cloud_run_gmp_sync.yaml @@ -43,7 +43,7 @@ spec: - action: keep sourceLabels: - __name__ - regex: 'omi_(finalization_pair_.*|enrichment_summary_repair_.*|batch_record_pending_.*|record_seal_nojob_.*|enrichment_obligation_.*|enrichment_early_abort_total|enrichment_supersede_streak(_bucket|_count|_sum)?|finalization_byok_.*|listen_server_batch_suppression_total|listen_preview_batch_custody_total|delivery_outcome_total|people_scan_dedupe_total|pinecone_.*|playback_writes_total|sync_.*|journey_.*|client_journey_.*|owner_recognition_conversations_total|conversation_duration_seconds(_.*)?|conversation_relevance_decision_total|conversation_summary_vector_upserts_total|conversation_relevance_decision_total|firestore_documents_per_operation(_.*)?|capture_finalization_.*|queue_oldest_ready_.*|batch_pipeline_.*|batch_projection_drains_skipped_total|batch_projection_claim_busy_backoffs_total|batch_projection_page_collapsed_rows_total|batch_projection_drain_head_requeues_total|batch_intake_outcome_total|batch_admission_sample_window_total|batch_ramp_heavy_owner_excluded_total|batch_shadow_.*|accept_to_summary_seconds(_.*)?|summary_terminal_total|summary_debt_.*|live_session_transcript_outcome_total|fallback_total|parakeet_request_seconds(_.*)?|firestore_billed_reads_by_(caller|origin)_total|firestore_caller_label_overflow_total|firestore_billed_writes(_by_caller)?_total|firestore_write_caller_label_overflow_total|firestore_ambiguous_commit_writes_total|firestore_noop_writes_by_caller_total|memory_released_projection_total|fsr_leader_lease_total|held_content_.*|held_recovery_pacing_.*|stt_cap_recovery_reconciliation_total|stt_cap_recovery_cohort_narrowed|stt_cap_cohort_.*|stt_cap_launch_config_errors.*|stt_cap_dispatch_decisions_total|stt_cap_provider_audio_seconds_total|legacy_sync_.*|legacy_dark_.*|jobstore_.*|capture_coverage_.*|capture_dedupe_shadow_total|capture_retry_cap_total|speaker_evidence_.*|speaker_evidence_confident_owner_total|lifecycle_refactor_shadow_total|executor_task_(active|queued)|clients_refactor_shadow_total|notes_refactor_shadow_total|notes_refactor_results_total|stt_language_gate_.*|prerecorded_vendor_.*|episode_compose_.*|transcription_final_pass_total|transcription_final_pass_(total|unit_skips_total)|transcription_final_pass_(total|unit_skips_total|(agreement|latency_seconds|owner_delta_ratio)_(bucket|sum|count))|transcription_final_pass_(total|unit_skips_total|retry_terminal_(total|observed_at_seconds)|(agreement|latency_seconds|owner_delta_ratio)_(bucket|sum|count)))|conversation_llm_input_dedupe_.*|privacy_(checks|check_calls|check_reads|decision_mismatch|shadow_checks|memo_requests|memo_ineligible|ttl_cache)_total' + regex: 'omi_(finalization_pair_.*|enrichment_summary_repair_.*|batch_record_pending_.*|record_seal_nojob_.*|enrichment_obligation_.*|enrichment_early_abort_total|enrichment_supersede_streak(_bucket|_count|_sum)?|finalization_byok_.*|listen_server_batch_suppression_total|listen_preview_batch_custody_total|delivery_outcome_total|people_scan_dedupe_total|pinecone_.*|playback_writes_total|sync_.*|journey_.*|client_journey_.*|owner_recognition_conversations_total|conversation_duration_seconds(_.*)?|conversation_relevance_decision_total|conversation_summary_vector_upserts_total|conversation_relevance_decision_total|firestore_documents_per_operation(_.*)?|capture_finalization_.*|queue_oldest_ready_.*|batch_pipeline_.*|batch_projection_drains_skipped_total|batch_projection_claim_busy_backoffs_total|batch_projection_page_collapsed_rows_total|batch_projection_drain_head_requeues_total|batch_projection_blocked_head_holds_total|batch_projection_shards_(held|unheld)|batch_projection_shard_wakes_total|batch_projection_shard_wake_depth|batch_intake_outcome_total|batch_admission_sample_window_total|batch_ramp_heavy_owner_excluded_total|batch_shadow_.*|accept_to_summary_seconds(_.*)?|summary_terminal_total|summary_debt_.*|live_session_transcript_outcome_total|fallback_total|parakeet_request_seconds(_.*)?|firestore_billed_reads_by_(caller|origin)_total|firestore_caller_label_overflow_total|firestore_billed_writes(_by_caller)?_total|firestore_write_caller_label_overflow_total|firestore_ambiguous_commit_writes_total|firestore_noop_writes_by_caller_total|memory_released_projection_total|fsr_leader_lease_total|held_content_.*|held_recovery_pacing_.*|stt_cap_recovery_reconciliation_total|stt_cap_recovery_cohort_narrowed|stt_cap_cohort_.*|stt_cap_launch_config_errors.*|stt_cap_dispatch_decisions_total|stt_cap_provider_audio_seconds_total|legacy_sync_.*|legacy_dark_.*|jobstore_.*|capture_coverage_.*|capture_dedupe_shadow_total|capture_retry_cap_total|speaker_evidence_.*|speaker_evidence_confident_owner_total|lifecycle_refactor_shadow_total|executor_task_(active|queued)|clients_refactor_shadow_total|notes_refactor_shadow_total|notes_refactor_results_total|stt_language_gate_.*|prerecorded_vendor_.*|episode_compose_.*|transcription_final_pass_total|transcription_final_pass_(total|unit_skips_total)|transcription_final_pass_(total|unit_skips_total|(agreement|latency_seconds|owner_delta_ratio)_(bucket|sum|count))|transcription_final_pass_(total|unit_skips_total|retry_terminal_(total|observed_at_seconds)|(agreement|latency_seconds|owner_delta_ratio)_(bucket|sum|count)))|conversation_llm_input_dedupe_.*|privacy_(checks|check_calls|check_reads|decision_mismatch|shadow_checks|memo_requests|memo_ineligible|ttl_cache)_total' # These eagerly initialized live-STT combinations dominate empty series # on sync workers. Keep fallbacks from sync/summary/storage components. - action: drop diff --git a/backend/docs/batch-pipeline-serving.md b/backend/docs/batch-pipeline-serving.md index 3d5eacbca39..0873205ccbc 100644 --- a/backend/docs/batch-pipeline-serving.md +++ b/backend/docs/batch-pipeline-serving.md @@ -1259,6 +1259,22 @@ below have their own switches; all three off is the pre-drain behaviour): advanced= more= elapsed_ms=` per turn, and `action=owner_chain outcome=claim_busy`. +### Owner shards and push wakes (dark, 2026-11-05) + +`BATCH_PROJECTION_SHARDS_ENABLED` (default off) replaces the per-instance +sweep's owner discovery with owner-hash shard leases. Each owner has one +drain writer fleet-wide, and that writer is woken by merge commits rather than +by a polling keyset. + +Only the maintenance leader pages the durable log, discovery-only, as a +backstop. Follower instances skip the sweep and keep their async follow-up +discovery. Holders fill drain slots fresh first, keeping +`BATCH_PROJECTION_BACKLOG_RESERVE` slots for backlog. + +Any Redis failure returns that instance to the sweep. The contract, env +bounds, cutover order and rollback are in +[Batch projection owner shards](listen_pusher_pipeline.mdx#batch-projection-owner-shards). + ### Follow-up desync (2026-11-05) Every serving process runs its own async follow-up discovery over the same diff --git a/backend/docs/feature-flag-registry.md b/backend/docs/feature-flag-registry.md index 3dc476bfc8e..9e280581970 100644 --- a/backend/docs/feature-flag-registry.md +++ b/backend/docs/feature-flag-registry.md @@ -278,6 +278,8 @@ and an explicit empty literal renders as `''`. | `BATCH_PIPELINE_TERMINAL_EMPTY_DISPOSITION_ENABLED` | Default off; positive strict VAD with exhausted GPU padded and multilingual recognition completes with distinct no_recoverable_speech result | backend | env | closed | {value: 'false', category: rollout} | false (backend-listen (chart)); {value: 'false', category: rollout} (cloud_run/backend, cloud_run/backend-integration, cloud_run/backend-sync, cloud_run/backend-sync-backfill, gke/backend-listen) | false (backend-listen (chart)); true (cloud_run/backend-integration); {value: 'false', category: rollout} (cloud_run/backend, cloud_run/backend-sync, cloud_run/backend-sync-backfill, gke/backend-listen) | — | graduate | on | 2026-11-04 | dazheng | | `BATCH_PIPELINE_UID_ALLOWLIST` | Explicit durable batch admission owner cohort; default empty and requires global enable; prod David-only rollout, dev off | backend | env | closed | — | '' (backend-listen (chart), gke/backend-listen) | — | — | kill | off | 2026-11-04 | dazheng | | `BATCH_PIPELINE_WORKER_SILENCE_DISPOSITION_ENABLED` | Default off; empty-ASR worker jobs reuse integration strict local VAD to commit confirmed silence before quarantine | backend | env | closed | — | — | — | — | graduate | on | 2026-11-04 | dazheng | +| `BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED` | Default off; an owner-chain head refused by a never-passing ProjectionBlocked (closed allowlist) is held once the head deadline lapses, so later captures project | backend | env | closed | — | — | — | — | pending | on | 2026-11-19 | dazheng | +| `BATCH_PROJECTION_SHARDS_ENABLED` | Default off; owner-hash projection shards: one serving member drains each owner from Redis wake lanes (fresh first), and only the maintenance leader pages the result log, discovery-only | backend | env | closed | — | — | — | — | pending | on | 2026-11-19 | dazheng | | `BATCH_VENDOR_ROUTING_PERCENT` | Default zero; owner-hash share of batch admissions that select their STT route from the vendor route table (cheapest healthy capable route, GPU floor share) | backend | env | closed | — | — | — | — | pending | undecided | 2026-11-29 | two-lane lead | | `BATCH_VENDOR_ROUTING_SHADOW_PERCENT` | Default zero; owners outside the routing cohort whose table decision is computed and logged (event=batch_vendor_route_shadow, no ids) while their routing stays unchanged | backend | env | closed | — | — | — | — | kill | off | 2026-11-29 | two-lane lead | | `BATCH_VENDOR_ROUTING_UID_ALLOWLIST` | Comma-separated owners routed from the vendor route table regardless of percent; default empty; env only, no cohort-document membership | backend | env | closed | — | — | — | — | kill | off | 2026-11-29 | two-lane lead | @@ -815,6 +817,9 @@ and an explicit empty literal renders as `''`. | `BATCH_PIPELINE_WORKER_AUDIENCE` | Expected service OIDC audience for private durable batch worker RPC; empty denies all requests | backend | env | closed | '' (cloud_run/backend, cloud_run/backend-integration, cloud_run/backend-sync-backfill, gke/backend-listen); declared (cloud_run/backend-sync) | '' (backend-listen (chart), cloud_run/backend, cloud_run/backend-integration, cloud_run/backend-sync, cloud_run/backend-sync-backfill, gke/backend-listen) | '' (backend-listen (chart), cloud_run/backend, cloud_run/backend-integration, cloud_run/backend-sync-backfill, gke/backend-listen); https://backend-sync-hhibjajaja-uc.a.run.app (cloud_run/backend-sync) | — | keep | keep | — | dazheng | | `BATCH_PIPELINE_WORKER_IO_CONCURRENCY` | GPU worker flight artifact IO concurrency: bounded fetch/decode prefetch and encrypted result writes including pre-lock preparation; default 5, 1 = serial, max 8 | backend | env | closed | — | — | — | — | keep | keep | — | dazheng | | `BATCH_PIPELINE_WORKER_SERVICE_ACCOUNT` | Exact verified service-account email allowed for batch worker RPC; empty denies all requests | backend | env | closed | '' (cloud_run/backend, cloud_run/backend-sync-backfill, gke/backend-listen); batch-pipeline-worker@{compute_project}.iam.gserviceaccount.com (cloud_run/backend-integration); declared (cloud_run/backend-sync) | '' (backend-listen (chart), cloud_run/backend, cloud_run/backend-sync, cloud_run/backend-sync-backfill, gke/backend-listen); batch-pipeline-worker@{compute_project}.iam.gserviceaccount.com (cloud_run/backend-integration) | '' (backend-listen (chart), cloud_run/backend, cloud_run/backend-sync-backfill, gke/backend-listen); batch-pipeline-worker@based-hardware.iam.gserviceaccount.com (cloud_run/backend-sync); batch-pipeline-worker@{compute_project}.iam.gserviceaccount.com (cloud_run/backend-integration) | — | keep | keep | — | dazheng | +| `BATCH_PROJECTION_BACKLOG_RESERVE` | Default 2; drain slots a projection shard holder keeps for backlog wakes while fresh wakes wait | backend | env | closed | — | — | — | — | keep | keep | — | dazheng | +| `BATCH_PROJECTION_HEAD_DEADLINE_SECONDS` | Default 86400; how long a failing (or, with the blocked-head deadline on, blocked) owner-chain head retries before it takes the held disposition | backend | env | closed | — | — | — | — | keep | keep | — | dazheng | +| `BATCH_PROJECTION_SHARD_MEMBER` | Default true; with projection shards on, false keeps a serving process out of shard leases (it serves, routes and pushes wakes only) | backend | env | open | — | — | — | — | keep | keep | — | dazheng | | `BATCH_VENDOR_HEALTH_THRESHOLDS` | Runtime JSON for batch vendor breaker thresholds and operator-entered runway balances; unset uses the vendorscope-ops 1.4 defaults | backend | env | closed | — | — | — | — | keep | keep | — | vroute-c / batch rollout lead | | `BATCH_VENDOR_OVERFLOW_ENTER_SECONDS` | Default 1200; fresh GPU backlog age in seconds at which the GPU ring reads saturated (vendor overflow, legacy capacity fallback); it stops below half of this (60-21600, else default) | backend | env | closed | — | — | — | — | keep | keep | — | two-lane lead | | `BATCH_VENDOR_OVERFLOW_MIN_READY_UNITS` | Default 300; admissible fresh units that must be waiting in the GPU ring before an old head counts as saturation for vendor overflow and the legacy capacity fallback; saturation ends below half of it (1-100000, else default) | backend | env | closed | — | — | — | — | keep | keep | — | two-lane lead | @@ -945,6 +950,8 @@ and an explicit empty literal renders as `''`. | `BATCH_PIPELINE_SHADOW_PERCENT` | 0 at shipment → 10% after the carrying train, first hour cost observation → remove after full record-lane ramp and legacy deletion | normal GPU admission/latency unchanged; shadow paid reservations <=200 per UTC hour; vendor usage and shadow GPU busy minutes reviewed after first hour | two-lane lead | 2026-10-31 | | `BATCH_PIPELINE_SHARED_LIVE_ENABLED` | decision at the M3 100% readout → one live-pool canary → all live pools | live first-text p95 unchanged and the record-lane ramp gate clean | dazheng (two-lane lead) | 2026-11-04 | | `BATCH_PIPELINE_SHORT_SPEECH_SINGLE_ENABLED` | paired pyannote 3533c8cf run → worker canary at current admission for 24 h → all batch units | skipped_short units match a single full-diarizer speaker on the gate evidence | dazheng (two-lane lead) | 2026-11-18 | +| `BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED` | staging → prod, after David approves | omi_batch_projection_blocked_head_holds_total bounded by the pre-switch count of heads blocked past the deadline; no hold outside the allowlist; held rows visible to the release scan | two-lane lead | 2026-11-12 | +| `BATCH_PROJECTION_SHARDS_ENABLED` | staging backend-integration → prod backend-integration, every serving instance in one config train | owner_drain rows advanced/s at or above the pre-switch hour for 60 min, claim_busy turns under 10%, omi_batch_projection_shards_unheld 0 after the roll, Firestore reads/s below the pre-switch hour | two-lane lead | 2026-11-12 | | `BATCH_VENDOR_ROUTING_PERCENT` | shadow 24h at batch 5% (BATCH_VENDOR_ROUTING_SHADOW_PERCENT) → 1% with a lead-run vendor failover drill → 5% then 25% in lockstep, never above batch-pipeline → 100% of batch-admitted owners | hard-3 (T1 custody, T2 lost speech, T4 language) clean; zero paid-intent violations; provider-stt queue age p95 under 15m; no vendor failure surfaced as an intake 5xx | two-lane lead | 2026-11-29 | | `BATCH_VENDOR_ROUTING_SHADOW_PERCENT` | 100% of eligible batch owners for 24h before the 1% routing step | shadow decisions match the intended table (route mix by reason) and intake latency is unchanged | two-lane lead | 2026-11-15 | | `CALLS_RECAP_MODE` | synthetic acceptance → 1% → 10% → 100% | omi_calls_refactor_shadow_total outcome=error stays zero; no false calls, premature scheduled recaps, missing provenance or STT config rejection; recap/STT latency and vendor spend stay within baseline | dazheng | 2026-10-26 | @@ -1149,6 +1156,11 @@ their code default (`fail` tells you which way a missing value resolves). - `BATCH_PIPELINE_WORKER_PIPELINE_DEPTH` — Default 2; GPU worker pump flights overlap lease claim, compute admission and result publication with the exclusive GPU chain (fail: open) - `BATCH_PIPELINE_WORKER_SILENCE_DISPOSITION_ENABLED` — Default off; empty-ASR worker jobs reuse integration strict local VAD to commit confirmed silence before quarantine (fail: closed) - `BATCH_PIPELINE_WORKER_VALIDATION_DISPOSITION_ENABLED` — Default on; strict silence for out-of-buffer VAD frames and first-attempt terminal GPU validation handoff with retained recovery custody (fail: open) +- `BATCH_PROJECTION_BACKLOG_RESERVE` — Default 2; drain slots a projection shard holder keeps for backlog wakes while fresh wakes wait (fail: closed) +- `BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED` — Default off; an owner-chain head refused by a never-passing ProjectionBlocked (closed allowlist) is held once the head deadline lapses, so later captures project (fail: closed) +- `BATCH_PROJECTION_HEAD_DEADLINE_SECONDS` — Default 86400; how long a failing (or, with the blocked-head deadline on, blocked) owner-chain head retries before it takes the held disposition (fail: closed) +- `BATCH_PROJECTION_SHARDS_ENABLED` — Default off; owner-hash projection shards: one serving member drains each owner from Redis wake lanes (fresh first), and only the maintenance leader pages the result log, discovery-only (fail: closed) +- `BATCH_PROJECTION_SHARD_MEMBER` — Default true; with projection shards on, false keeps a serving process out of shard leases (it serves, routes and pushes wakes only) (fail: open) - `BATCH_PULL_BUSY_BORROW_ENABLED` — Default on; a pump's non-first lanes lease recent/deep backfill when every fresh shard is claimed by another puller; false restores strict fresh-first on every lane (fail: open) - `BATCH_PULL_MULTI_LEASE_ENABLED` — Default on; one claimed shard pull transaction leases up to the pull batch from the ring it already read; false restores one lease per transaction (fail: open) - `BATCH_SHARD_SPREAD_ENABLED` — Default on; token-fenced per-shard batch pull ownership with worker affinity among equal-service candidates; false restores legacy pulls (fail: open) diff --git a/backend/docs/listen_pusher_pipeline.mdx b/backend/docs/listen_pusher_pipeline.mdx index ce53355dd77..98fda31650c 100644 --- a/backend/docs/listen_pusher_pipeline.mdx +++ b/backend/docs/listen_pusher_pipeline.mdx @@ -9351,3 +9351,147 @@ owner-drain admissions use a read-only Redis claim GET and fail open on errors. Claim-busy re-requests through `again` back off 5–15 seconds; productive `more=True` continuations retain immediate FIFO progress without the peek. The advisory claim and transactional ordering fences remain unchanged. + +## Batch projection owner shards + +Contract for `utils/batch_pipeline/projection_shards.py` and +`utils/batch_pipeline/projection_shard_pump.py`. Every switch below is read by +backend-integration serving processes only. None of them is an ordering +fence. The assignment transaction's predecessor, privacy, binding and receipt +fences, and the advisory Redis owner claim around each drain step, decide +every row exactly as before. + +| Env | Default | Bounds and fallback | +| --- | --- | --- | +| `BATCH_PROJECTION_SHARDS_ENABLED` | off | Only `true`/`1`/`on` enables it. Off is the per-instance sweep. | +| `BATCH_PROJECTION_SHARD_MEMBER` | true | `false` keeps the process out of shard leases. It still routes and pushes wakes. | +| `BATCH_PROJECTION_BACKLOG_RESERVE` | 2 | 0..16, never more than the drain slots (`BATCH_PIPELINE_MAX_OWNER_DRAINS`). Malformed values keep 2. | +| `BATCH_PROJECTION_HEAD_DEADLINE_SECONDS` | 86400 | 3600..604800. Malformed or out-of-range values keep 86400, the previous constant. | +| `BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED` | off | Prod enablement needs David's approval. | + +**Shards and leases.** An owner's shard is `int(owner_key[:12], 16) % 64`, and +`owner_key` is the sha256 of the owner id. Each shard has a 30 s Redis lease, +`batch_projection_shard:{namespace}:64:{shard}`, held by one random +per-process holder id. Every process with shards on runs a lease pass every +10 s: +- It heartbeats into the members zset and prunes members idle for 30 s. +- It renews its own leases and releases everything above + `ceil(64 / live members)`. +- It takes free shards up to that share. + +A free shard means some member is below its share, so a free shard is taken +on the next pass. Shutdown releases leases and leaves the members zset at +once. The count is in every key, so a different count never shares a lease. + +**Wakes.** Two lanes per shard: `fresh` and `backlog`. +- `fresh` holds owners with a newly committed capture under 6 h old. Its score + is the newest capture time, and it is taken newest first. +- `backlog` holds everything else. Its score is the oldest commit time, and it + is taken oldest first. + +A push is one Lua call. It keeps the best score and increments a per-owner +sequence. A drain that walks to the chain end settles each wake it took, but +only if no push landed since the take. So a commit during a drain always gets +another drain. + +A drain that stops short (blocked, error, turn or step budget, hold not +landed) cools the owner for about 120 s in process memory. It also moves the +wake to `backlog` at its retry time, so stuck owners never fill the take page +ahead of due owners. A claim-busy stop cools the owner for about 15 s and +keeps its lane position. Wake keys expire 7 days after their last push. + +**Push sources.** +1. Merge commit. `Serving.drain_owner` pushes one wake per owner per commit + batch to the owner's shard lane, from whichever process served the worker + RPC. An owner on an unheld shard, or a failed push, drains locally. +2. Routed drain requests. When the owner's shard has a holder, a + `Projector.request_drain` (recovery wake, follow-up, continuation) becomes + a `backlog` wake. +3. Leader tail. The GPU worker's own merge commits run without Redis, so + they push no wake. Each 20 s supervisor tick, the maintenance leader pages + `pending` rows committed since its last pass, at most 5 pages of 100 rows. + - A caught-up pass restarts 10 s before the newest row it saw, and never + earlier than 60 s ago. A full page resumes at its cursor. + - Each owner is routed by its row's `capture_at`, so fresh GPU commits + reach the fresh lane within one tick. +4. Leader backstop. After the tail, the leader pages the whole pending log, + discovery-only. + - At most 30 pages of 100 rows in 10 s, from a resuming cursor that wraps. + - It routes each owner once per tick and runs the sweep's finish page. + +Followers skip the sweep. Every process keeps its own async follow-up +discovery. + +**Pump.** Each process's 1 s pump fills free drain slots: +- Fresh wakes come first. +- `BATCH_PROJECTION_BACKLOG_RESERVE` slots are kept for backlog whenever + backlog waits. +- Reserved slots return to fresh when no backlog waits. +- An owner present in both lanes drains once and settles both wakes. + +Redis calls run on `db_executor`, never on the event loop. + +**Fail-open.** A process is inactive, and runs the legacy per-instance sweep +with local drain requests kept local, when any of these holds: +- Its lease pass failed. +- It has had no lease pass for 30 s. +- A wake take failed. It stays inactive until the next lease pass succeeds. +- `BATCH_PIPELINE_OWNER_CHAIN_PROJECTION_ENABLED`, + `BATCH_PIPELINE_SWEEP_OWNER_CHAIN_ENABLED`, + `BATCH_PIPELINE_OWNER_DRAIN_ENABLED` or + `BATCH_PIPELINE_PROJECTION_PENDING_PAGE_ENABLED` is off, so each existing + kill switch still restores its old path. A process switched off this way also + releases its shard leases and leaves the members set, so engaged members + take its shards on their next lease pass instead of routing wakes to lanes + nobody reads. + +An owner whose shard nobody held at the last pass drains +wherever it is found. A failed push keeps the request local. + +**Blocked-head deadline** (`BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED`, +default off). This applies only to an exact `ProjectionBlocked` chain head +whose reason is one of `intake_binding_mismatch`, `invalid_recording_clock`, +`invalid_recording_identity`, `invalid_recording_interval`, +`invalid_upload_filenames` or `lineage_identity_missing`, or whose message is +one of: +- terminal empty recovery evidence required +- recovery result requires recovery handoff +- validated VAD speech accounting required +- repair requires authenticated lineage manifest +- repair requires active proven lineage fence +- repair requires proven generation disposition + +The first refusal stamps `projection_blocked_since`, with +`projection_block_type` set to `ProjectionBlocked:`. A stamp left by +any other failure restarts the clock, so the deadline measures time spent +refused the same way. Once `BATCH_PROJECTION_HEAD_DEADLINE_SECONDS` lapses, +the head takes the existing +held disposition (`quarantine_held` receipt, watermark advance, job +`recovery_state=held`) and the walk continues only if the watermark moved. + +Audio and result stay. The release scan and the reviewed redrive own the +unit. Claim, shadow, privacy, predecessor-pending, generation-ambiguous and +deployment-wiring refusals never qualify. + +Metrics: +- `omi_batch_projection_shards_held` and `omi_batch_projection_shards_unheld` + (gauges). +- `omi_batch_projection_shard_wakes_total{lane,outcome}`, where `outcome` is + `pushed`, `push_failed`, `taken`, `settled`, `repushed` or `cooled`. +- `omi_batch_projection_shard_wake_depth{lane}` (held shards, every 60 s). +- `omi_batch_projection_blocked_head_holds_total`. + +Logs: `event=batch_projection_shards action=refresh|push|take|settle|defer|backstop|tick`. + +**Cutover.** Turn it on for staging backend-integration first. Then turn it on +for every prod backend-integration serving process in one config train. +During the roll, processes still on the sweep are fenced by the claim. + +Clean metrics, all over 60 min: +- `omi_batch_projection_shards_unheld` is 0 after the roll. +- Owner-drain rows advanced per second are at or above the pre-switch hour. +- Claim-busy turns are under 10%. +- Firestore reads per second are below the pre-switch hour. + +**Rollback.** Set `BATCH_PROJECTION_SHARDS_ENABLED=false`. The sweep returns +on the next tick. Leftover wake keys are inert and expire. diff --git a/backend/docs/listen_pusher_pipeline/changes/2026-11-05-projection-owner-shards.md b/backend/docs/listen_pusher_pipeline/changes/2026-11-05-projection-owner-shards.md new file mode 100644 index 00000000000..080af238d64 --- /dev/null +++ b/backend/docs/listen_pusher_pipeline/changes/2026-11-05-projection-owner-shards.md @@ -0,0 +1,49 @@ +# Owner-hash projection shards, push wakes, fresh-first drains, blocked-head deadline + +Before this change, every backend-integration instance paged the same global +`committed_at` keyset and raced its peers for each owner's advisory claim. +Each pass reset the cursor. Most turns ended `claim_busy` or found nothing, +while the GPUs sat near 20% and the persisted backlog grew. The fix is in four +parts. + +1. **One writer per owner.** Owners hash to 64 shards. Each shard has a 30 s + Redis lease held by one serving process, and the leases balance across + members. Only the holder drains the shard's owners, so drains stop + competing for claims. The claim and the transactional predecessor fence + still guard every assignment. +2. **Push wakes instead of polling.** + - A merge commit pushes a wake for the owner's shard holder. + - The GPU worker's own commits have no Redis, so the maintenance leader + follows the newly committed tail each tick. + - The leader alone runs a discovery-only backstop over the durable log, + and followers no longer sweep. + - Firestore reads drop from N instances paging the full keyset every 20 s + to one leader paging at most 30 pages per tick, with no per-row pre-read. +3. **Fresh first.** Commits of audio captured under 6 h ago take the `fresh` + lane, newest capture first. Two drain slots per process stay reserved for + backlog while backlog waits. A stuck owner moves behind due owners. +4. **Every non-terminal head has an exit.** + - Failing heads already took the held disposition after 24 h. That bound + is now `BATCH_PROJECTION_HEAD_DEADLINE_SECONDS`. + - Behind `BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED` (default off), + never-passing `ProjectionBlocked` heads on a closed allowlist take the + same held disposition once the deadline lapses. Audio is kept and the + reviewed redrive owns the unit. + - Prod enablement needs David's approval. + +All new switches default to today's behavior. Shards apply to the serving +loops only. A dedicated projection pump process (`BATCH_PIPELINE_PROJECTION_PUMP_ENABLED`) +keeps its plain sweep. With `BATCH_PIPELINE_SERVING_PROJECTION_SWEEP=off` the +leader skips the tail and the backstop, but commit wakes still flow. + +Not included: +- Transactional multi-row assignment batching. It touches the core receipt + transaction and needs its own review. +- Provider-stt and enrichment fresh-first ordering, which is #24714. + +Contract: [Batch projection owner shards](../../listen_pusher_pipeline.mdx#batch-projection-owner-shards). +Hermetic coverage: +- `tests/unit/test_batch_projection_shards.py`: leases, wakes, pump, serving + wiring and the deadline. +- `tests/unit/test_batch_projection_shard_drain.py`: two processes with real + drains on shared fakeredis. diff --git a/backend/tests/support/firestore_batch_pipeline_drivers.py b/backend/tests/support/firestore_batch_pipeline_drivers.py index a13d7a66eeb..21520d7e0a8 100644 --- a/backend/tests/support/firestore_batch_pipeline_drivers.py +++ b/backend/tests/support/firestore_batch_pipeline_drivers.py @@ -198,6 +198,22 @@ def entries(): }, neutrals={'limit': (100, 'bounded page size only')}, ), + DriverEntry( + 'database.batch_pipeline_store.projection_page_owners', + domains={ + 'namespace': namespaces, + 'after': [None, (2.0, 'result-100'), (2.0, '')], + 'states': [('pending',), ('pending', 'persisted', 'ready')], + }, + neutrals={'limit': (100, 'bounded page size only')}, + ), + CoveredByEntry( + 'database.batch_pipeline_store.FirestoreStore.projection_page_owners', + covered_by=('database.batch_pipeline_store.projection_page_owners',), + reason='store adapter forwards namespace, page bound, cursor and states unchanged', + expect_observed=False, + body_digest='d25e6dabd5431fc92214567a023c2348de2ca5c62ff4738b3b930d42d10526ed', + ), DriverEntry( 'database.batch_pipeline_store.projection_successor', base={'owner_key': 'owner-hash'}, diff --git a/backend/tests/support/firestore_query_driver_registry.py b/backend/tests/support/firestore_query_driver_registry.py index 5362f7b709c..837b9110c14 100644 --- a/backend/tests/support/firestore_query_driver_registry.py +++ b/backend/tests/support/firestore_query_driver_registry.py @@ -159,6 +159,8 @@ def patcher(client): 'database.batch_pipeline_store.projection_pending': '52cc913cca9086557133a8f7118b9980ce5f13649897207fafcfa221513a5629', 'database.batch_pipeline_store.FirestoreStore.projection_page': 'bc9fe5cd08a05748eaa118af15031abdcb1a68d91d7f0b74bf177a27eb980a94', 'database.batch_pipeline_store.FirestoreStore.projection_successor': '96adabf372104327c5a8065ef2ff06621648ee242b3415d04b047d50ded9e6b1', + # Added 2026-11-05: sharded projection backstop; same query/index as projection_page. + 'database.batch_pipeline_store.FirestoreStore.projection_page_owners': 'd25e6dabd5431fc92214567a023c2348de2ca5c62ff4738b3b930d42d10526ed', 'database.batch_pipeline_store.FirestoreTransaction.service_frontier': '74aa1993996cdc9ca93a1141de64dfc7e05a9d831a67860ddbb8a64d55520a02', 'database.batch_pipeline_store.FirestoreStore.control_candidates': 'c1ce9e8170a7fde0e6a82c44a2db0c0c18909cba723e9cefd402cfe1f7993df9', 'database.batch_pipeline_store.FirestoreStore.control_candidate': 'd82ee8b1c03be61b104e25304d6a0a732a9b59453fd302f4e7794a75622d65ad', diff --git a/backend/tests/unit/test_attach_cloud_run_gmp_sidecar.py b/backend/tests/unit/test_attach_cloud_run_gmp_sidecar.py index e2ad91808e2..ab165caee31 100644 --- a/backend/tests/unit/test_attach_cloud_run_gmp_sidecar.py +++ b/backend/tests/unit/test_attach_cloud_run_gmp_sidecar.py @@ -1402,6 +1402,11 @@ def kept(name, component=''): assert kept(family + 'duration_seconds_bucket') assert kept('omi_batch_projection_drains_skipped_total') assert kept('omi_batch_projection_claim_busy_backoffs_total') + assert kept('omi_batch_projection_blocked_head_holds_total') + assert kept('omi_batch_projection_shards_held') + assert kept('omi_batch_projection_shards_unheld') + assert kept('omi_batch_projection_shard_wakes_total') + assert kept('omi_batch_projection_shard_wake_depth') # Direct GMP consumers: sentinel batch_metrics.py/completion_telemetry.py. # fleet_gate/david_gate use logs, durable state and native request_count. for name in ( diff --git a/backend/tests/unit/test_batch_projection_shard_drain.py b/backend/tests/unit/test_batch_projection_shard_drain.py new file mode 100644 index 00000000000..328a64da228 --- /dev/null +++ b/backend/tests/unit/test_batch_projection_shard_drain.py @@ -0,0 +1,122 @@ +"""Two serving processes, real owner drains and walks, one shard writer per owner.""" + +import asyncio +from contextlib import contextmanager + +import fakeredis +import pytest + +from tests.unit.test_batch_pipeline_owner_chain_projection import ResultLog, chain, projector_for +from utils.batch_pipeline import projection, projection_drain, projection_followup, projection_shard_pump +from utils.batch_pipeline.projection_shard_pump import ShardPump +from utils.batch_pipeline.projection_shards import BACKLOG, FRESH, ProjectionShards +from utils.batch_pipeline.scheduler import owner_key + + +@pytest.fixture(autouse=True) +def claims(monkeypatch): + state = {'held': set(), 'overlaps': 0} + + @contextmanager + def claim(namespace, key): + if key in state['held']: + state['overlaps'] += 1 + yield False + return + state['held'].add(key) + try: + yield True + finally: + state['held'].discard(key) + + monkeypatch.setattr(projection_drain, 'projection_claim_held', lambda namespace, key: key in state['held']) + monkeypatch.setattr(projection, 'projection_claim', claim) + monkeypatch.setattr(projection_followup, 'projection_claim', claim) + for name in ('BATCH_PIPELINE_OWNER_DRAIN_ENABLED', 'BATCH_PIPELINE_PROJECTION_PREFETCH'): + monkeypatch.delenv(name, raising=False) + + async def inline(executor, fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(projection_shard_pump, 'run_blocking', inline) + return state + + +async def quiesce(projectors): + for _ in range(2000): + busy = [p for p in projectors if p.drains is not None and (p.drains.tasks or p.drains.queue)] + if not busy: + return + for p in busy: + await asyncio.gather(*list(p.drains.tasks.values()), return_exceptions=True) + await asyncio.sleep(0) + raise AssertionError('drains did not settle') + + +@pytest.mark.asyncio +async def test_two_processes_drain_every_owner_once_in_chain_order(claims): + server = fakeredis.FakeServer() + log = ResultLog() + owners = [f'owner-{n}' for n in range(12)] + chains = {owner: chain(log, owner, owner, 4) for owner in owners} + walkers: dict[str, set[int]] = {owner: set() for owner in owners} + pumps = [] + for index in range(2): + client = fakeredis.FakeRedis(server=server) + shards = ProjectionShards('ns', count=8, client=lambda client=client: client, clock=lambda: 10_000.0) + projector = projector_for(log) + real = projector.drain_step + + def step(owner, budget, real=real, index=index, **kwargs): + walkers[owner].add(index) + return real(owner, budget, **kwargs) + + projector.drain_step = step + pump = ShardPump(projector, shards) + projector.router = pump.route + pumps.append(pump) + for _ in range(3): + for pump in pumps: + pump.shards.refresh() + assert sorted(s for p in pumps for s in p.shards.held) == list(range(8)) + # Commits arrive on whichever process served the worker's RPC. + for n, owner in enumerate(owners): + pumps[n % 2].committed(owner, [{'capture_at': 10_000.0 - n, 'committed_at': 10_000.0}]) + for _ in range(6): + for pump in pumps: + await pump.tick() + await quiesce([p.projector for p in pumps]) + for owner, ids in chains.items(): + assert [sid for sid in log.persist_order if log.results[sid]['owner_key'] == owner_key(owner)] == ids + assert len(walkers[owner]) == 1, 'one shard writer per owner' + holder = next(i for i, p in enumerate(pumps) if p.shards.owns(owner_key(owner))) + assert walkers[owner] == {holder} + assert claims['overlaps'] == 0 + for pump in pumps: + assert pump.shards.take(FRESH, 100) == [] and pump.shards.take(BACKLOG, 100) == [] + assert pump.inflight == {} + + +@pytest.mark.asyncio +async def test_router_sends_foreign_owners_to_their_holder(claims): + server = fakeredis.FakeServer() + log = ResultLog() + pumps = [] + for _ in range(2): + client = fakeredis.FakeRedis(server=server) + shards = ProjectionShards('ns', count=4, client=lambda client=client: client, clock=lambda: 1.0) + projector = projector_for(log) + pump = ShardPump(projector, shards) + projector.router = pump.route + pumps.append(pump) + for _ in range(3): + for pump in pumps: + pump.shards.refresh() + owner = next(f'o{n}' for n in range(1000) if pumps[1].shards.owns(owner_key(f'o{n}'))) + ids = chain(log, owner, owner, 2) + # A recovery wake on process 0 asks for a drain of process 1's owner. + pumps[0].projector.request_drain(owner) + assert pumps[0].projector.drain_requests == {} + await pumps[1].tick() + await quiesce([pumps[1].projector]) + assert log.persist_order == ids diff --git a/backend/tests/unit/test_batch_projection_shards.py b/backend/tests/unit/test_batch_projection_shards.py new file mode 100644 index 00000000000..852a7b7ef72 --- /dev/null +++ b/backend/tests/unit/test_batch_projection_shards.py @@ -0,0 +1,693 @@ +"""Owner-hash projection shards, the shard pump and the blocked-head deadline. + +Leases and wakes run against fakeredis (real Lua). The pump and serving +wiring use small fakes; chain walks use the owner-chain ``ResultLog`` model. +""" + +import asyncio +from collections import Counter +from unittest.mock import Mock + +import fakeredis +import pytest +from google.api_core import exceptions as api_exceptions # noqa: F401 (parity with chain tests) + +from database.batch_pipeline_projection import ProjectionBlocked +from tests.unit.test_batch_pipeline_owner_chain_projection import BURST, ResultLog, chain +from utils.batch_pipeline import projection, projection_shard_pump, projection_shards +from utils.batch_pipeline.projection_shard_pump import ShardPump +from utils.batch_pipeline.projection_shards import BACKLOG, FRESH, ProjectionShards, Wake, share, shard_of +from utils.batch_pipeline.scheduler import owner_key +from utils.batch_pipeline.serving import Serving + +ENV = ( + 'BATCH_PROJECTION_SHARDS_ENABLED', + 'BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED', + 'BATCH_PROJECTION_HEAD_DEADLINE_SECONDS', + 'BATCH_PROJECTION_BACKLOG_RESERVE', + 'BATCH_PIPELINE_MAX_OWNER_DRAINS', +) + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + for name in ENV: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def server(): + return fakeredis.FakeServer() + + +def make(server, *, count=8, member=True, clock=lambda: 1000.0): + client = fakeredis.FakeRedis(server=server) + return ProjectionShards('ns', count=count, member=member, client=lambda: client, clock=clock) + + +def owner_in_shard(shard: int, count: int, prefix='owner') -> str: + for n in range(10_000): + owner = f'{prefix}-{n}' + if shard_of(owner_key(owner), count) == shard: + return owner + raise AssertionError('no owner for shard') + + +# --- leases ----------------------------------------------------------------- + + +def test_shard_of_is_stable_and_bounded(): + keys = [owner_key(f'u{n}') for n in range(500)] + shards = [shard_of(key, 64) for key in keys] + assert all(0 <= s < 64 for s in shards) + assert shards == [shard_of(key, 64) for key in keys] + assert len(set(shards)) > 50 # spreads + assert share(64, 35) == 2 and share(64, 0) == 64 and share(8, 3) == 3 + + +def test_members_cover_every_shard_exactly_once_and_balance(server): + members = [make(server, count=16) for _ in range(5)] + for _ in range(3): + for member in members: + assert member.refresh() + held = [s for member in members for s in member.held] + assert sorted(held) == list(range(16)) + assert max(len(m.held) for m in members) <= share(16, 5) + assert all(m.unheld == set() for m in members[-1:]) + assert all(m.active for m in members) + + +def test_a_released_member_hands_its_shards_to_the_others(server): + members = [make(server, count=12) for _ in range(3)] + for _ in range(2): + for member in members: + member.refresh() + gone = members.pop() + released = set(gone.held) + assert released + gone.release() + assert not gone.active and gone.held == set() + for _ in range(3): + for member in members: + member.refresh() + held = [s for member in members for s in member.held] + assert sorted(held) == list(range(12)) + + +def test_a_joining_member_takes_a_share_without_gaps(server): + first = make(server, count=8) + first.refresh() + assert first.held == set(range(8)) + second = make(server, count=8) + for _ in range(6): + second.refresh() + first.refresh() + assert len(first.held | second.held) + len(first.unheld) >= 8 + assert sorted(first.held | second.held) == list(range(8)) + assert not first.held & second.held + assert len(second.held) >= 3 + + +def test_non_member_holds_nothing_and_releases_old_leases(server): + member = make(server, count=4) + member.refresh() + observer = make(server, count=4, member=False) + observer.refresh() + assert observer.held == set() and observer.unheld == set() + assert observer.covered(owner_key('anyone')) + + +def test_redis_failure_fails_open_to_inactive(): + def broken(): + raise ConnectionError('down') + + shards = ProjectionShards('ns', count=4, client=broken) + assert shards.refresh() is False + assert not shards.active and shards.held == set() and shards.unheld == set(range(4)) + assert shards.push('o', owner_key('o'), FRESH, 1.0) is False + assert shards.take(FRESH, 10, shards=[0]) == [] + + +# --- wakes ------------------------------------------------------------------ + + +def test_fresh_wakes_come_newest_first_and_backlog_oldest_first(server): + shards = make(server, count=1) + shards.refresh() + for owner, score in (('a', 10.0), ('b', 30.0), ('c', 20.0)): + shards.push(owner, owner_key(owner), FRESH, score) + shards.push(owner, owner_key(owner), BACKLOG, score) + assert [w.owner_id for w in shards.take(FRESH, 10)] == ['b', 'c', 'a'] + assert [w.owner_id for w in shards.take(BACKLOG, 10)] == ['a', 'c', 'b'] + # Fresh keeps the newest capture, backlog the oldest commit. + shards.push('a', owner_key('a'), FRESH, 40.0) + shards.push('a', owner_key('a'), BACKLOG, 50.0) + assert shards.take(FRESH, 1)[0].owner_id == 'a' + assert shards.take(BACKLOG, 1)[0].owner_id == 'a' + + +def test_settle_keeps_a_wake_pushed_after_it_was_taken(server): + shards = make(server, count=1) + shards.refresh() + shards.push('a', owner_key('a'), BACKLOG, 5.0) + taken = shards.take(BACKLOG, 10)[0] + shards.push('a', owner_key('a'), BACKLOG, 9.0) # same score kept (lt), newer push + assert shards.settle(taken) is False + assert [w.owner_id for w in shards.take(BACKLOG, 10)] == ['a'] + assert shards.settle(shards.take(BACKLOG, 10)[0]) is True + assert shards.take(BACKLOG, 10) == [] + + +def test_take_interleaves_held_shards(server): + shards = make(server, count=2) + shards.refresh() + for n in range(3): + for shard in (0, 1): + owner = owner_in_shard(shard, 2, prefix=f's{shard}n{n}') + shards.push(owner, owner_key(owner), BACKLOG, float(n)) + lanes = [w.shard for w in shards.take(BACKLOG, 3)] + assert lanes[:2] in ([0, 1], [1, 0]) + assert Counter(lanes) == {0: 3, 1: 3} + + +# --- pump ------------------------------------------------------------------- + + +class FakeDrains: + def __init__(self, limit=4): + self.limit = limit + self.tasks: dict = {} + self.queue: list = [] + self.delayed: dict = {} + self.requested: list[str] = [] + self.on_idle = None + + @property + def free_slots(self): + return max(0, self.limit - len(self.tasks) - len(self.queue)) + + def busy(self, owner): + return owner in self.tasks or owner in self.queue or owner in self.delayed + + def request(self, owners): + for owner in owners: + self.requested.append(owner) + self.tasks[owner] = True + + def finish(self, owner, stop): + self.tasks.pop(owner, None) + self.on_idle(owner, stop) + + +class FakeProjector: + def __init__(self, drains): + self.drains = drains + self.draining: set[str] = set() + self.drain_requests: dict = {} + self.local: list[str] = [] + + def ensure_drains(self): + return self.drains + + @property + def owes_drains(self): + return bool(self.drain_requests) + + def start_drains(self): + self.drains.request(list(self.drain_requests)) + self.drain_requests = {} + + def request_local_drain(self, owner): + self.local.append(owner) + + +@pytest.fixture +def blocking_inline(monkeypatch): + async def inline(executor, fn, *args, **kwargs): + return fn(*args, **kwargs) + + monkeypatch.setattr(projection_shard_pump, 'run_blocking', inline) + + +def pump_for(server, *, limit=4, count=1): + shards = make(server, count=count) + drains = FakeDrains(limit) + pump = ShardPump(FakeProjector(drains), shards) + return pump, shards, drains + + +def test_route_keeps_local_drains_when_inactive_or_unheld(server): + pump, shards, _ = pump_for(server, count=4) + assert pump.route('o') is True # never refreshed: inactive + observer = make(server, count=4, member=False) + observer.refresh() # nobody holds anything + pump.shards = observer + assert pump.route('o') is True + holder = make(server, count=4) + holder.refresh() + observer.refresh() + assert pump.route('o') is False # pushed to the holder's backlog lane + assert [w.owner_id for w in holder.take(BACKLOG, 10)] == ['o'] + + +def test_tick_takes_fresh_first_and_keeps_a_backlog_reserve(server, blocking_inline, monkeypatch): + monkeypatch.setenv('BATCH_PROJECTION_BACKLOG_RESERVE', '1') + pump, shards, drains = pump_for(server, limit=3) + for owner, score in (('f1', 3.0), ('f2', 2.0), ('f3', 1.0)): + shards.push(owner, owner_key(owner), FRESH, score) + shards.push('b1', owner_key('b1'), BACKLOG, 1.0) + assert asyncio.run(pump.tick()) == 3 + assert drains.requested == ['f1', 'f2', 'b1'] + + +def test_reserve_returns_to_fresh_when_no_backlog_waits(server, blocking_inline, monkeypatch): + monkeypatch.setenv('BATCH_PROJECTION_BACKLOG_RESERVE', '2') + pump, shards, drains = pump_for(server, limit=3) + for owner, score in (('f1', 3.0), ('f2', 2.0), ('f3', 1.0)): + shards.push(owner, owner_key(owner), FRESH, score) + assert asyncio.run(pump.tick()) == 3 + assert drains.requested == ['f1', 'f2', 'f3'] + + +def test_end_settles_the_wake_and_a_stuck_owner_cools_down(server, blocking_inline): + pump, shards, drains = pump_for(server, limit=2) + now = [100.0] + pump.monotonic = lambda: now[0] + shards.push('done', owner_key('done'), FRESH, 2.0) + shards.push('stuck', owner_key('stuck'), FRESH, 1.0) + asyncio.run(pump.tick()) + drains.finish('done', 'end') + drains.finish('stuck', 'blocked_other') + asyncio.run(pump.tick()) + assert shards.take(FRESH, 10) == [] + # The stuck wake waits in backlog at its retry time, not at the lane head. + (deferred,) = shards.take(BACKLOG, 10) + assert deferred.owner_id == 'stuck' and deferred.score > 1000.0 + assert drains.requested == ['done', 'stuck'] # stuck is cooling, not retaken + now[0] += projection_shard_pump.STUCK_RETRY_SECONDS * 1.3 + asyncio.run(pump.tick()) + assert drains.requested == ['done', 'stuck', 'stuck'] + + +def test_stuck_owners_cannot_fill_the_take_page(server, blocking_inline, monkeypatch): + monkeypatch.setattr(projection_shard_pump, 'MAX_TAKE', 4) + pump, shards, drains = pump_for(server, limit=2) + stuck = [f'stuck-{n}' for n in range(6)] + for n, owner in enumerate(stuck): + shards.push(owner, owner_key(owner), BACKLOG, float(n)) + shards.push('due', owner_key('due'), BACKLOG, 500.0) + for _ in range(3): + asyncio.run(pump.tick()) + for owner in list(drains.tasks): + drains.finish(owner, 'blocked_other') + asyncio.run(pump.tick()) + assert 'due' in drains.requested + assert drains.requested.count('due') == 1 + + +def test_claim_busy_cools_without_losing_lane_position(server, blocking_inline): + pump, shards, drains = pump_for(server, limit=1) + shards.push('a', owner_key('a'), FRESH, 5.0) + asyncio.run(pump.tick()) + drains.finish('a', 'claim_busy') + asyncio.run(pump.tick()) + assert [(w.owner_id, w.score) for w in shards.take(FRESH, 10)] == [('a', 5.0)] + + +def test_defer_keeps_a_wake_that_was_pushed_since_the_take(server): + shards = make(server, count=1) + assert shards.refresh() + shards.push('a', owner_key('a'), FRESH, 5.0) + (wake,) = shards.take(FRESH, 10) + shards.push('a', owner_key('a'), FRESH, 6.0) + assert shards.defer(wake, 9999.0) is False + assert [w.score for w in shards.take(FRESH, 10)] == [6.0] + (wake,) = shards.take(FRESH, 10) + assert shards.defer(wake, 9999.0) is True + assert shards.take(FRESH, 10) == [] + assert [(w.owner_id, w.score) for w in shards.take(BACKLOG, 10)] == [('a', 9999.0)] + + +def test_a_wake_pushed_during_the_drain_survives_the_settle(server, blocking_inline): + pump, shards, drains = pump_for(server, limit=1) + shards.push('a', owner_key('a'), FRESH, 1.0) + asyncio.run(pump.tick()) + shards.push('a', owner_key('a'), FRESH, 1.0) # another commit while draining + drains.finish('a', 'end') + asyncio.run(pump.tick()) + assert drains.requested == ['a', 'a'] + + +def test_owner_in_both_lanes_drains_once_and_settles_both(server, blocking_inline): + pump, shards, drains = pump_for(server, limit=4) + shards.push('a', owner_key('a'), FRESH, 1.0) + shards.push('a', owner_key('a'), BACKLOG, 1.0) + asyncio.run(pump.tick()) + assert drains.requested == ['a'] + drains.finish('a', 'end') + asyncio.run(pump.tick()) + assert shards.take(FRESH, 10) == [] and shards.take(BACKLOG, 10) == [] + + +def test_commit_wakes_pick_the_lane_from_capture_age(server, monkeypatch): + monkeypatch.setattr(projection_shard_pump, 'FRESH_CAPTURE_SECONDS', 3600.0) + pump, shards, _ = pump_for(server) + shards.refresh() + pump.committed('new', [{'capture_at': 1000.0 - 60, 'committed_at': 1000.0}]) + pump.committed('old', [{'capture_at': 1000.0 - 7200, 'committed_at': 999.0}]) + assert [w.owner_id for w in shards.take(FRESH, 10)] == ['new'] + assert [w.owner_id for w in shards.take(BACKLOG, 10)] == ['old'] + + +def test_backstop_routes_owners_once_and_falls_back_locally_for_unheld(server): + holder = make(server, count=2) + holder.refresh() + holder.release() # nobody holds shard leases now + member = make(server, count=2) + member.refresh() + lone = owner_in_shard(0, 2, prefix='lone') + rows = [('s1', lone, 1.0, None), ('s2', lone, 2.0, None), ('s3', 'x', 3.0, None)] + + class Store: + def projection_page_owners(self, *, limit, after, states): + assert states == ('pending',) + if after is not None: + return [], None # the tail: nothing committed lately + return rows, None + + real = projection.Projector(Mock(), Mock()) + real.store = Mock() + real.store.store = Store() + pump = ShardPump(real, member) + counts = pump.backstop() + assert counts == {'pages': 1, 'rows': 3, 'owners': 2, 'wrapped': 1, 'tail_rows': 0, 'tail_owners': 0} + assert sorted(w.owner_id for w in member.take(BACKLOG, 10)) == sorted({lone, 'x'}) + # Unheld shards drain where they were found. + observer = make(server, count=2, member=False) + member.release() + observer.refresh() + real.drain_requests.clear() + ShardPump(real, observer).backstop() + assert set(real.drain_requests) == {lone, 'x'} + + +def test_tail_discovery_wakes_gpu_commits_fresh_and_follows_the_newest_row(server): + member = make(server, count=1) + member.refresh() + calls = [] + pages = { + (940.0, ''): ([('s1', 'gpu-fresh', 995.0, 990.0), ('s2', 'gpu-old', 996.0, 1000.0 - 7 * 3600)], None), + (986.0, ''): ([], None), + } + + class Store: + def projection_page_owners(self, *, limit, after, states): + calls.append(after) + if after is None: + return [], None # the whole-log backstop + return pages[after] + + real = projection.Projector(Mock(), Mock(), clock=lambda: 1000.0) + real.store = Mock() + real.store.store = Store() + pump = ShardPump(real, member) + counts = pump.backstop() + assert counts['tail_rows'] == 2 and counts['tail_owners'] == 2 + assert [w.owner_id for w in member.take(FRESH, 10)] == ['gpu-fresh'] + assert [(w.owner_id, w.score) for w in member.take(BACKLOG, 10)] == [('gpu-old', 996.0)] + # Caught up: the next pass starts just before the newest row seen. + assert real.tail_after == (996.0 - projection_shard_pump.TAIL_SLACK_SECONDS, '') + pump.backstop() + assert calls == [(940.0, ''), None, (986.0, ''), None] + + +def test_a_failed_take_turns_the_process_back_to_the_sweep(server): + shards = make(server, count=2) + assert shards.refresh() and shards.active + shards._client = lambda: (_ for _ in ()).throw(ConnectionError('down')) + assert shards.take(FRESH, 10) == [] + assert not shards.active + + +def test_kill_switches_disengage_the_pump(server, monkeypatch): + pump, shards, _ = pump_for(server) + shards.refresh() + assert pump.engaged + monkeypatch.setenv('BATCH_PIPELINE_OWNER_DRAIN_ENABLED', 'false') + assert not pump.engaged + assert pump.route('o') is True + + +def test_a_settle_never_resets_the_push_sequence(server): + shards = make(server, count=1) + shards.refresh() + shards.push('a', owner_key('a'), FRESH, 1.0) + (first,) = shards.take(FRESH, 10) + (stale,) = shards.take(FRESH, 10) # a second take of the same push + assert shards.settle(first) + shards.push('a', owner_key('a'), FRESH, 2.0) # new work after the settle + assert not shards.settle(stale) + assert [w.owner_id for w in shards.take(FRESH, 10)] == ['a'] + + +def test_commits_on_an_unheld_shard_drain_locally(server): + pump, shards, _ = pump_for(server, count=2) + shards.refresh() + shards.held = {0} + shards.unheld = {1} + lone = owner_in_shard(1, 2, prefix='lone') + pump.committed(lone, [{'capture_at': 990.0, 'committed_at': 999.0}]) + assert pump.projector.local == [lone] + + +# --- serving wiring ----------------------------------------------------------- + + +class Pump: + def __init__(self, active): + self.engaged = active + self.backstops = 0 + self.commits = [] + + def backstop(self): + self.backstops += 1 + + def committed(self, owner, entries): + self.commits.append((owner, len(entries))) + + +def serving_with(pump, *, leads): + service = Serving(Mock(), projector=Mock(), shard_pump=pump) + service._leads = lambda periodic: leads + service._maintain = lambda **kw: None + service.merge = Mock(held_owners=set()) + return service + + +def test_sharded_reconcile_skips_the_sweep_and_only_the_leader_backstops(): + pump = Pump(active=True) + follower = serving_with(pump, leads=False) + follower.reconcile(periodic=True) + follower.projector.sweep.assert_not_called() + assert pump.backstops == 0 and follower.projector.has_async_followups is True + leader = serving_with(pump, leads=True) + leader.reconcile(periodic=True) + leader.projector.sweep.assert_not_called() + leader.projector.sweep_finish_only.assert_called_once() + assert pump.backstops == 1 + + +def test_inactive_shards_keep_the_legacy_sweep(): + service = serving_with(Pump(active=False), leads=False) + service.reconcile(periodic=True) + service.projector.sweep.assert_called_once() + service = serving_with(None, leads=True) + service.reconcile(periodic=True) + service.projector.sweep.assert_called_once() + + +def test_drain_owner_pushes_one_wake_for_its_commits(): + pump = Pump(active=True) + service = serving_with(pump, leads=False) + entries = iter([{'capture_at': 1.0}, {'capture_at': 2.0}, None]) + service.merge.commit_next = lambda owner, now, release=None: next(entries) + service.drain_owner('owner-1') + assert pump.commits == [('owner-1', 2)] + service.merge.commit_next = lambda owner, now, release=None: None + service.drain_owner('owner-2') + assert pump.commits == [('owner-1', 2)] + + +# --- blocked-head deadline ---------------------------------------------------- + + +def blocked_session(log, ids, message, *, refusals): + real = log.session + attempts = {'n': 0} + + def session(sid, *args, **kwargs): + if sid == ids[0]: + attempts['n'] += 1 + if attempts['n'] <= refusals: + raise ProjectionBlocked(message) + return real(sid, *args, **kwargs) + + log.session = session + + +def walker(log): + projector = projection.Projector(log, Mock(), clock=lambda: 1000.0) + projector.store = log + return projector + + +def test_blocked_head_waits_forever_with_the_switch_off(): + log = ResultLog() + ids = chain(log, BURST, 'a', 2) + blocked_session(log, ids, 'validated VAD speech accounting required', refusals=10**6) + projector = walker(log) + projector.project_owner(BURST) + assert log.results[ids[0]]['projection_state'] == 'pending' + assert 'projection_blocked_since' not in log.results[ids[0]] + + +def test_eligible_blocked_head_is_held_after_the_deadline(monkeypatch): + monkeypatch.setenv('BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED', 'true') + monkeypatch.setenv('BATCH_PROJECTION_HEAD_DEADLINE_SECONDS', '3600') + log = ResultLog() + ids = chain(log, BURST, 'a', 2) + # Refusals: the first pass, the in-bound second pass, the past-bound assignment. + blocked_session(log, ids, 'validated VAD speech accounting required', refusals=3) + projector = walker(log) + projector.project_owner(BURST) + assert log.results[ids[0]]['projection_state'] == 'pending' + assert log.results[ids[0]]['projection_blocked_since'] == 1000.0 + projector.project_owner(BURST) + assert log.results[ids[0]]['projection_state'] == 'pending' + assert log.results[ids[0]]['projection_block_type'] == 'ProjectionBlocked:other' + log.results[ids[0]]['projection_blocked_since'] = 1000.0 - 3601 + projector.project_owner(BURST) + assert log.results[ids[0]]['projection_state'] == 'quarantine_held' + assert log.projection_watermark(owner_key(BURST)) == ids[-1] + + +def test_a_stamp_from_another_failure_restarts_the_blocked_clock(monkeypatch): + monkeypatch.setenv('BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED', 'true') + monkeypatch.setenv('BATCH_PROJECTION_HEAD_DEADLINE_SECONDS', '3600') + log = ResultLog() + ids = chain(log, BURST, 'a', 2) + blocked_session(log, ids, 'recovery result requires recovery handoff', refusals=10**6) + # A transient error stamped this head long ago; the refusal is new. + log.results[ids[0]].update(projection_blocked_since=0.0, projection_block_type='TimeoutError') + projector = walker(log) + projector.project_owner(BURST) + row = log.results[ids[0]] + assert row['projection_state'] == 'pending' + assert row['projection_blocked_since'] == 1000.0 + assert row['projection_block_type'] == 'ProjectionBlocked:other' + + +def test_a_walk_past_a_held_head_reports_the_chain_end(monkeypatch): + monkeypatch.setenv('BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED', 'true') + log = ResultLog() + ids = chain(log, BURST, 'a', 2) + blocked_session( + log, ids, 'validated VAD speech accounting required', refusals=1 + ) # The test session refuses; the real refusal is in assignment. + log.results[ids[0]].update(projection_blocked_since=1000.0 - 86401, projection_block_type='ProjectionBlocked:other') + from contextlib import contextmanager + + @contextmanager + def claim(namespace, key): + yield True + + monkeypatch.setattr(projection, 'projection_claim', claim) + walk = walker(log).drain_step(BURST, 60.0) + assert log.results[ids[0]]['projection_state'] == 'quarantine_held' + assert walk.advanced == ids and walk.stop == 'end' + + +@pytest.mark.parametrize( + 'message', + [ + 'projection assignment claim busy', + 'shadow publication forbidden', + 'authenticated intake manifest loader required', + ], +) +def test_ineligible_blocked_heads_never_take_the_deadline(monkeypatch, message): + monkeypatch.setenv('BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED', 'true') + log = ResultLog() + ids = chain(log, BURST, 'a', 1) + blocked_session(log, ids, message, refusals=10**6) + log.results[ids[0]].update(projection_blocked_since=0.0, projection_block_type='ProjectionBlocked:other') + walker(log).project_owner(BURST) + assert log.results[ids[0]]['projection_state'] == 'pending' + + +def test_eligibility_uses_closed_reasons_and_exact_messages(): + from database.batch_pipeline_projection_protocol import ProjectionPredecessorPending + + reasoned = ProjectionBlocked('anything') + reasoned.block_reason = 'invalid_recording_clock' + assert projection.blocked_head_deadline_eligible(reasoned) + assert projection.blocked_head_deadline_eligible(ProjectionBlocked('terminal empty recovery evidence required')) + assert not projection.blocked_head_deadline_eligible(ProjectionBlocked('projection assignment claim busy')) + assert not projection.blocked_head_deadline_eligible( + ProjectionPredecessorPending('validated VAD speech accounting required') + ) + + +def test_head_deadline_is_configurable_and_bounded(monkeypatch): + entry = {'projection_blocked_since': 0.0} + assert not projection.head_retry_bound_exceeded(entry, 24 * 3600.0) + assert projection.head_retry_bound_exceeded(entry, 24 * 3600.0 + 1) + monkeypatch.setenv('BATCH_PROJECTION_HEAD_DEADLINE_SECONDS', '7200') + assert projection.head_retry_bound_exceeded(entry, 7201) + monkeypatch.setenv('BATCH_PROJECTION_HEAD_DEADLINE_SECONDS', '60') # below the floor: default + assert not projection.head_retry_bound_exceeded(entry, 7201) + assert not projection.head_retry_bound_exceeded({'projection_blocked_since': True}, 10**9) + + +def test_a_drain_task_that_raised_never_reports_the_chain_end(monkeypatch): + from utils.batch_pipeline import projection_drain + + def boom(namespace, key): + raise RuntimeError('admission failed') + + monkeypatch.setattr(projection_drain, 'projection_claim_held', boom) + projector = Mock() + projector.draining = set() + drains = projection_drain.OwnerDrains(projector) + seen = [] + drains.on_idle = lambda owner, stop: seen.append((owner, stop)) + + async def go(): + drains.request(['a']) + await asyncio.gather(*drains.tasks.values(), return_exceptions=True) + await asyncio.sleep(0) + + asyncio.run(go()) + assert seen == [('a', 'error')] + + +def test_a_kill_switched_member_hands_its_shards_to_engaged_members(server, blocking_inline, monkeypatch): + pumps = [] + for _ in range(2): + shards = make(server, count=4) + pumps.append(ShardPump(FakeProjector(FakeDrains(4)), shards)) + for _ in range(3): + for pump in pumps: + pump.shards.refresh() + assert pumps[0].shards.held and pumps[1].shards.held + monkeypatch.setenv('BATCH_PIPELINE_OWNER_DRAIN_ENABLED', 'false') + asyncio.run(pumps[0].tick()) + assert pumps[0].shards.held == set() and not pumps[0].shards.active + # The engaged member sees one member left and takes every shard. + pumps[1].shards.refresh() + assert pumps[1].shards.held == set(range(4)) and pumps[1].shards.unheld == set() + asyncio.run(pumps[0].tick()) # still switched off: no lease pass, no shards + assert pumps[0].shards.held == set() + monkeypatch.delenv('BATCH_PIPELINE_OWNER_DRAIN_ENABLED') + asyncio.run(pumps[0].tick()) + assert pumps[0].shards.active # switched back on: rejoins on the next lease pass diff --git a/backend/utils/batch_pipeline/metrics.py b/backend/utils/batch_pipeline/metrics.py index f26445899fc..64b4b314841 100644 --- a/backend/utils/batch_pipeline/metrics.py +++ b/backend/utils/batch_pipeline/metrics.py @@ -24,6 +24,20 @@ PROJECTION_DRAIN_CLAIM_BUSY_BACKOFFS = Counter( 'omi_batch_projection_claim_busy_backoffs_total', 'Owner drain turns deferred after losing the fleet claim race' ) +PROJECTION_BLOCKED_HEAD_HOLDS = Counter( + 'omi_batch_projection_blocked_head_holds_total', 'Blocked owner-chain heads held after the head deadline lapsed' +) +# Sharded projection serving (projection_shards.py). Static labels only. +PROJECTION_SHARDS_HELD = Gauge('omi_batch_projection_shards_held', 'Projection shards this process holds') +PROJECTION_SHARDS_UNHELD = Gauge( + 'omi_batch_projection_shards_unheld', 'Projection shards no member held at the last lease pass' +) +PROJECTION_SHARD_WAKES = Counter( + 'omi_batch_projection_shard_wakes_total', 'Projection owner wakes by lane and outcome', ['lane', 'outcome'] +) +PROJECTION_SHARD_WAKE_DEPTH = Gauge( + 'omi_batch_projection_shard_wake_depth', 'Owner wakes waiting in this process\'s held shards', ['lane'] +) LABELS = {Priority.FRESH: 'fresh', Priority.RECENT: 'recent', Priority.DEEP: 'deep'} AGE_TARGETS = {Priority.FRESH: 30, Priority.RECENT: 300, Priority.DEEP: 3600} diff --git a/backend/utils/batch_pipeline/projection.py b/backend/utils/batch_pipeline/projection.py index b8991784758..e1c3ba71893 100644 --- a/backend/utils/batch_pipeline/projection.py +++ b/backend/utils/batch_pipeline/projection.py @@ -20,6 +20,8 @@ from config.batch_pipeline import ( owner_drain_enabled, + projection_blocked_head_deadline_enabled, + projection_head_deadline_seconds, projection_index_deferred_enabled, projection_pending_page_enabled, projection_predecessor_replay_enabled, @@ -162,8 +164,52 @@ def projection_block_reason(error: BaseException) -> str: # fresh awaiting-writer conflicts may keep a head waiting, but not forever: # once a head has failed for this long it takes the held disposition so the # chain can advance (audio retained; a reviewed redrive owns the unit). +# The default of BATCH_PROJECTION_HEAD_DEADLINE_SECONDS (config.batch_pipeline), +# which is the live bound. PROJECTION_HEAD_RETRY_SECONDS = 24 * 3600.0 +# ProjectionBlocked heads the blocked-head deadline may hold once it lapses +# (BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED, default off). Only refusals +# that the same committed bytes and binding can never pass on retry, or that +# wait on a handoff/fence that has not appeared within the deadline. Claim, +# shadow, privacy, predecessor and deployment-wiring refusals never qualify: +# they keep waiting exactly as before. +BLOCKED_HEAD_DEADLINE_MESSAGES = frozenset( + { + 'terminal empty recovery evidence required', + 'recovery result requires recovery handoff', + 'validated VAD speech accounting required', + 'repair requires authenticated lineage manifest', + 'repair requires active proven lineage fence', + 'repair requires proven generation disposition', + } +) +BLOCKED_HEAD_DEADLINE_REASONS = frozenset( + { + 'intake_binding_mismatch', + 'invalid_recording_clock', + 'invalid_recording_identity', + 'invalid_recording_interval', + 'invalid_upload_filenames', + 'lineage_identity_missing', + } +) + + +def blocked_head_deadline_eligible(error: BaseException) -> bool: + """Whether a ProjectionBlocked chain head may be held once its deadline lapses.""" + if type(error) is not ProjectionBlocked: + return False # Predecessor-pending and generation-ambiguous keep their own rules. + if projection_block_reason(error) in BLOCKED_HEAD_DEADLINE_REASONS: + return True + return str(error) in BLOCKED_HEAD_DEADLINE_MESSAGES + + +def blocked_head_type(error: BaseException) -> str: + """Bounded, log-safe episode token for a deadline-eligible refusal.""" + return f'ProjectionBlocked:{projection_block_reason(error)}' + + # Result-row fields stamping a bounded-retry head's first failure. Written # once per failure episode; the release redrive clears them with the receipt. PROJECTION_HEAD_BLOCKED_SINCE = 'projection_blocked_since' @@ -241,7 +287,11 @@ def head_holds_without_retry(error: BaseException, job: dict | None, *, now: flo def head_retry_bound_exceeded(entry: dict, now: float) -> bool: since = entry.get(PROJECTION_HEAD_BLOCKED_SINCE) - return isinstance(since, (int, float)) and now - since > PROJECTION_HEAD_RETRY_SECONDS + return ( + isinstance(since, (int, float)) + and not isinstance(since, bool) + and now - since > projection_head_deadline_seconds() + ) @contextmanager @@ -438,6 +488,13 @@ def __init__( self.drain_requests: dict[str, None] = {} self.draining: set[str] = set() self.drains = None + # Sharded serving: ``router(owner_id) -> bool`` keeps a drain request + # local (True) or has already routed it to the shard holder (False). + self.router: Callable[[str], bool] | None = None + # Resume point of the sharded discovery backstop. + self.backstop_after = None + # Start of the sharded tail discovery (rows committed since the last pass). + self.tail_after = None # Sweeps and wakes hand off from sync-pool threads while overlapping # reconcile passes and drains swap or extend on the event loop. self.handoff = threading.Lock() @@ -450,9 +507,106 @@ def __init__( self.release_after: str | None = None def request_drain(self, owner_id: str) -> None: + # Sharded serving: an owner whose shard another member holds is routed + # there as a wake instead of draining here (see projection_shards). + router = self.router + if router is not None and not router(owner_id): + return with self.handoff: self.drain_requests[owner_id] = None + def request_local_drain(self, owner_id: str) -> None: + """Queue a drain in this process, bypassing the shard router.""" + with self.handoff: + self.drain_requests[owner_id] = None + + def ensure_drains(self): + """This process's owner drain, created on first use. Event-loop thread only.""" + from utils.batch_pipeline.projection_drain import OwnerDrains + + if self.drains is None: + self.drains = OwnerDrains(self) + return self.drains + + def discover_tail( + self, + route: Callable[[str, float, float | None], None], + *, + pages: int, + overlap_seconds: float, + slack_seconds: float, + ) -> dict: + """Discovery of newly committed rows for sharded serving (maintenance leader only). + + Every commit path pushes a wake except the GPU worker's own merge, + which has no Redis. This pages pending rows committed since the last + pass, oldest first, and hands each owner to ``route`` once per call + with the row's ``committed_at`` and ``capture_at``. A short page + restarts the next pass ``slack_seconds`` before its newest row (commit + stamps precede their transaction), never earlier than + ``overlap_seconds`` ago; a full page resumes at its cursor. + """ + now = self.clock() + floor = (now - overlap_seconds, '') + start = self.tail_after or floor + counts = {'pages': 0, 'rows': 0, 'owners': 0} + seen: set[str] = set() + newest = None + for _ in range(max(1, pages)): + rows, after = self.store.store.projection_page_owners(limit=100, after=start, states=('pending',)) + counts['pages'] += 1 + counts['rows'] += len(rows) + for _segment_id, owner_id, committed_at, capture_at in rows: + newest = committed_at if newest is None else max(newest, committed_at) + if owner_id in seen: + continue + seen.add(owner_id) + route(owner_id, committed_at, capture_at) + if after is None: + start = None + break + start = after + if start is None: + # Caught up: restart near the newest row seen, bounded below. + restart = floor[0] if newest is None else max(floor[0], newest - slack_seconds) + start = (max(restart, (self.tail_after or floor)[0]), '') + self.tail_after = start + counts['owners'] = len(seen) + return counts + + def discover_owners( + self, route: Callable[[str, float, float | None], None], *, pages: int, budget_seconds: float + ) -> dict: + """Discovery-only backstop for sharded serving (maintenance leader only). + + Pages the durable pending log and hands each owner found to ``route`` + once per call, oldest commit first. No assignment, claim or per-row + read: owners come from the page rows. The cursor resumes across calls + and wraps at a short page, so every pending row is revisited. + """ + deadline = time.monotonic() + budget_seconds + counts = {'pages': 0, 'rows': 0, 'owners': 0, 'wrapped': 0} + seen: set[str] = set() + for _ in range(max(1, pages)): + if counts['pages'] and time.monotonic() >= deadline: + break + rows, after = self.store.store.projection_page_owners( + limit=100, after=self.backstop_after, states=PROJECTION_ASSIGNMENT_STATES + ) + self.backstop_after = after + counts['pages'] += 1 + counts['rows'] += len(rows) + for _segment_id, owner_id, committed_at, capture_at in rows: + if owner_id in seen: + continue + seen.add(owner_id) + route(owner_id, committed_at, capture_at) + if after is None: + counts['wrapped'] = 1 + break + counts['owners'] = len(seen) + return counts + def owe_followups(self, segment_ids) -> None: with self.handoff: self.chain_followups.extend(segment_ids) @@ -472,14 +626,10 @@ def owes_index(self) -> bool: def start_drains(self) -> None: """Hand discovered owners to this process's drain. Event-loop thread only.""" - # Imported lazily: the synchronous worker seam remains lightweight. - from utils.batch_pipeline.projection_drain import OwnerDrains - - if self.drains is None: - self.drains = OwnerDrains(self) + drains = self.ensure_drains() with self.handoff: requests, self.drain_requests = self.drain_requests, {} - self.drains.request(requests) + drains.request(requests) def schedule_index_flush(self, owner_id: str | None = None) -> None: """Converge owed search work off-chain, optionally for one owner.""" @@ -842,20 +992,26 @@ def assign_record_silence(self, item, binding, session, *, hint=None): def _mirror_assignment(self, session) -> None: mirror_assignment(self.store, session) - def _note_head_blocked(self, segment_id: str, error: BaseException) -> None: - """Stamp a bounded-retry head's first failure time; one write per episode.""" + def _note_head_blocked(self, segment_id: str, error: BaseException, *, block_type: str | None = None) -> None: + """Stamp a bounded-retry head's first failure time; one write per episode. + + With ``block_type`` (the blocked-head deadline) the episode is that + refusal alone: a stamp left by a different failure restarts the clock, + so the deadline always measures time spent refused the same way. + """ + stamp_type = block_type or error_type(error) def commit(transaction): row = self.store.ref('results', segment_id).get(transaction=transaction).to_dict() or {} - if row.get(PROJECTION_HEAD_BLOCKED_SINCE) is not None or row.get('projection_state') not in ( - 'pending', - 'persisted', - 'ready', + if row.get('projection_state') not in ('pending', 'persisted', 'ready'): + return False + if row.get(PROJECTION_HEAD_BLOCKED_SINCE) is not None and ( + block_type is None or row.get(PROJECTION_HEAD_BLOCK_TYPE) == block_type ): return False transaction.update( self.store.ref('results', segment_id), - {PROJECTION_HEAD_BLOCKED_SINCE: self.clock(), PROJECTION_HEAD_BLOCK_TYPE: error_type(error)}, + {PROJECTION_HEAD_BLOCKED_SINCE: self.clock(), PROJECTION_HEAD_BLOCK_TYPE: stamp_type}, ) return True @@ -895,6 +1051,29 @@ def resolve_failing_head(self, segment_id, job, entry, binding, error: BaseExcep self._note_head_blocked(segment_id, error) return False + def resolve_blocked_head(self, segment_id, job, entry, binding, error: BaseException) -> bool: + """Deadline for an eligible ProjectionBlocked chain head (default off). + + The first refusal of this kind stamps the episode start (a stamp left + by any other failure or refusal restarts it); once the head deadline + lapses the unit takes the same held disposition as a poison head + (quarantine_held receipt, audio and result kept, release scan and + reviewed redrive own it) and the chain advances. True only when the + hold was attempted and reported success. + """ + block_type = blocked_head_type(error) + entry = entry or {} + if entry.get(PROJECTION_HEAD_BLOCK_TYPE) == block_type and head_retry_bound_exceeded(entry, self.clock()): + held = self._hold_failing_head(segment_id, job, binding) + if held: + from utils.batch_pipeline.metrics import PROJECTION_BLOCKED_HEAD_HOLDS + + PROJECTION_BLOCKED_HEAD_HOLDS.inc() + logger.info('event=batch_projection action=blocked_head outcome=held') + return held + self._note_head_blocked(segment_id, error, block_type=block_type) + return False + def assign_parts(self, item, entry, binding, *, finish=True, policy=None, hint=None, job=None): policy = policy or self.store.policy(item.segment_id) from database.conversations import get_closest_conversation_to_timestamps @@ -1465,7 +1644,28 @@ def owe(segment_id: str, binding: ProjectionBinding) -> bool: break except ProjectionBlocked as error: stop = walk_block_stop(error, projection_block_reason(error)) - break + if not ( + binding is not None + and job is not None + and projection_blocked_head_deadline_enabled() + and blocked_head_deadline_eligible(error) + and self.resolve_blocked_head(segment_id, job, entry, binding, error) + ): + break + # The deadline lapsed and the hold landed: continue iff the chain advanced. + try: + watermark = self.store.projection_watermark(key) + except Exception as watermark_error: + failed(watermark_error) + stop = 'watermark_error' + break + if watermark != segment_id: + stop = 'hold_not_landed' + break + previous = segment_id + advanced.append(segment_id) + stop = 'end' # The hold moved the chain on; the next row decides. + continue except Exception as error: failed(error) if binding is None or not self.resolve_failing_head(segment_id, job, entry, binding, error): @@ -1668,6 +1868,13 @@ def sweep(self, limit: int = 100) -> dict[str, int]: ) return counts + def sweep_finish_only(self, limit: int = 100) -> dict[str, int]: + """The sweep's finish page alone (sharded serving's leader backstop).""" + counts = {'completed': 0, 'waiting': 0, 'blocked': 0, 'error': 0} + if projection_pending_page_enabled(): + self._sweep_finish(max(1, min(limit, 100)), counts, set()) + return counts + def _sweep_finish(self, size: int, counts: dict[str, int], touched: set[str]) -> bool: """Advance persisted/ready rows that the pending-only assignment page skips. diff --git a/backend/utils/batch_pipeline/projection_drain.py b/backend/utils/batch_pipeline/projection_drain.py index 12ed6e87349..b3674e9dba0 100644 --- a/backend/utils/batch_pipeline/projection_drain.py +++ b/backend/utils/batch_pipeline/projection_drain.py @@ -15,7 +15,7 @@ import logging import random import time -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Callable from config.batch_pipeline import max_owner_drains, projection_drain_claim_aware_enabled from database.batch_pipeline_projection_claim import projection_claim_held @@ -27,7 +27,7 @@ from utils.batch_pipeline.projection import Projector logger = logging.getLogger(__name__) -# Owners draining at once in one process. +# Owners draining at once in one process (BATCH_PIPELINE_MAX_OWNER_DRAINS, read at import). MAX_OWNER_DRAINS = max_owner_drains() # One owner's turn; with more rows it re-queues behind every waiting owner. DRAIN_TURN_SECONDS = 120.0 @@ -47,7 +47,24 @@ def __init__(self, projector: Projector): self.continuations: set[str] = set() # Only active admissions/turns write this bounded stop state. self.stops: dict[str, str] = {} + # The last turn's stop token per owner, for the idle hook. + self.turn_stops: dict[str, str] = {} self.closed = False + # Called as ``on_idle(owner, stop)`` once an owner leaves this drain + # entirely (no task, queue entry, delay or continuation). Shard pumps + # settle or cool their wakes here. + self.on_idle: Callable[[str, str], None] | None = None + + @property + def limit(self) -> int: + return MAX_OWNER_DRAINS + + @property + def free_slots(self) -> int: + return max(0, self.limit - len(self.tasks) - len(self.queue)) + + def busy(self, owner: str) -> bool: + return owner in self.tasks or owner in self.queue or owner in self.delayed def request(self, owners) -> None: if self.closed: @@ -62,7 +79,7 @@ def request(self, owners) -> None: self.launch() def launch(self) -> None: - while self.queue and len(self.tasks) < MAX_OWNER_DRAINS and not self.closed: + while self.queue and len(self.tasks) < self.limit and not self.closed: owner = self.queue.popleft() self.projector.draining.add(owner) continuation = owner in self.continuations @@ -93,7 +110,8 @@ def defer(self, owner: str) -> None: if len(self.delayed) >= MAX_DELAYED_OWNERS: # Return overflow intent to the next serving pass; no extra timers, # immediate busy loop, or delay for productive FIFO continuations. - self.projector.request_drain(owner) + # Local on purpose: the shard router would push to Redis on the loop. + getattr(self.projector, 'request_local_drain', self.projector.request_drain)(owner) return def retry(): @@ -137,6 +155,18 @@ def finished(self, owner: str, task: asyncio.Task) -> None: else: self.queue.append(owner) self.again.discard(owner) + last = self.turn_stops.pop(owner, None) + if task.exception() is not None: + last = 'error' + if self.on_idle is not None and not self.busy(owner): + try: + # Only a turn that walked to the chain end reports 'end'. + self.on_idle(owner, stop or last or 'unknown') + except Exception as error: + logger.warning( + 'event=batch_projection action=owner_drain_idle outcome=error exception_type=%s', + type(error).__name__, + ) self.launch() async def run(self, owner: str) -> bool: @@ -178,6 +208,7 @@ async def run(self, owner: str) -> bool: finally: if prefetch is not None: prefetch.close() + self.turn_stops[owner] = stop logger.info( 'event=batch_projection action=owner_drain outcome=turn steps=%d advanced=%d more=%s stop=%s elapsed_ms=%d', steps, diff --git a/backend/utils/batch_pipeline/projection_shard_pump.py b/backend/utils/batch_pipeline/projection_shard_pump.py new file mode 100644 index 00000000000..1a9ffd45d3a --- /dev/null +++ b/backend/utils/batch_pipeline/projection_shard_pump.py @@ -0,0 +1,308 @@ +"""Feeds a shard holder's owner drains from its shards' wake lanes. + +The pump runs on the serving event loop. Each tick it renews shard leases +when due, then pulls wakes into free drain slots: fresh wakes first (newest +capture first), with ``backlog_reserve`` slots kept for backlog wakes (oldest +commit first) whenever backlog is waiting, so fresh audio reaches the user +first and the backlog still drains. A wake leaves Redis only after its owner's +drain walks to the chain end, and only if no newer push landed meanwhile. A +stuck owner cools down in process memory and its wake moves to the backlog +lane at its retry time, behind every owner due sooner. + +Every Redis call runs off the event loop and fails open. +""" + +from __future__ import annotations + +import logging +import random +import time +from typing import TYPE_CHECKING, Callable + +from config.batch_pipeline import ( + owner_chain_projection_enabled, + owner_drain_enabled, + projection_backlog_reserve, + projection_pending_page_enabled, + sweep_owner_chain_enabled, +) +from utils.batch_pipeline.metrics import ( + PROJECTION_SHARD_WAKE_DEPTH, + PROJECTION_SHARD_WAKES, + PROJECTION_SHARDS_HELD, + PROJECTION_SHARDS_UNHELD, +) +from utils.batch_pipeline.projection_shards import BACKLOG, FRESH, ProjectionShards, Wake +from utils.batch_pipeline.scheduler import owner_key +from utils.executors import db_executor, run_blocking + +if TYPE_CHECKING: + from utils.batch_pipeline.projection import Projector + +logger = logging.getLogger(__name__) +# A drain that stopped short of the chain end waits this long before its +# owner is taken again (a stuck head, an error, a lost claim race). +STUCK_RETRY_SECONDS = 120.0 +CLAIM_RETRY_SECONDS = 15.0 +CLAIM_STOPS = frozenset({'claim_busy', 'claim_held', 'blocked_projection_assignment_claim_busy'}) +MAX_TAKE = 256 +MAX_COOLDOWN = 20_000 +# The backstop spends at most this much of one supervisor tick, in pages of +# 100 rows: one pass over a 180k-row backlog is about 20 minutes of 150 reads/s. +BACKSTOP_SECONDS = 10.0 +BACKSTOP_PAGES = 30 +# Tail discovery: rows committed since the last leader pass (the GPU worker's +# merge commits push no wake). Up to 500 rows per 20 s supervisor tick; a full +# page resumes at its cursor on the next tick. +TAIL_PAGES = 5 +TAIL_OVERLAP_SECONDS = 60.0 +TAIL_SLACK_SECONDS = 10.0 +# Commits of audio captured this recently take the fresh lane (the sync fresh +# lane's age, utils/sync/lanes.py). +FRESH_CAPTURE_SECONDS = 6 * 3600.0 +# Shard pump cadence on the event loop. +TICK_SECONDS = 1.0 + + +class ShardPump: + def __init__( + self, + projector: Projector, + shards: ProjectionShards, + *, + monotonic: Callable[[], float] = time.monotonic, + ): + self.projector = projector + self.shards = shards + self.monotonic = monotonic + # Owner -> the wakes this pump handed to the drain for it. + self.inflight: dict[str, list[Wake]] = {} + # Owner -> monotonic time before which the pump does not take it again. + self.cooldown: dict[str, float] = {} + # Wakes to settle on the next tick (Redis stays off the event loop). + self.settling: list[Wake] = [] + # Cooled wakes to move behind due work on the next tick, with retry scores. + self.deferring: list[tuple[Wake, float]] = [] + self.attached = False + + # Any thread ----------------------------------------------------------- + + @property + def engaged(self) -> bool: + """Sharded serving applies: shards answer and the drain it relies on is on. + + The pump only feeds the owner drain, and the leader backstop leaves + persisted/ready rows to the sweep's finish page. When any of those kill + switches is off, this process serves exactly as it did before shards. + """ + return self.shards.active and self.switches_on + + @property + def switches_on(self) -> bool: + """Every existing drain kill switch the sharded path relies on is on.""" + return ( + owner_chain_projection_enabled() + and sweep_owner_chain_enabled() + and owner_drain_enabled() + and projection_pending_page_enabled() + ) + + def route(self, owner_id: str, score: float | None = None) -> bool: + """Projector router: True keeps the drain request in this process. + + With shards engaged, an owner whose shard some member holds becomes a + backlog wake for that holder (this process included, so its pump owns + the bookkeeping). Unheld shards and Redis failures keep the legacy + local drain, so no owner is ever left without a walker. + """ + return not self.wake(owner_id, None, self.shards.clock() if score is None else score) + + def wake(self, owner_id: str, capture_at: float | None, committed_at: float) -> bool: + """Push one wake to the owner's shard holder. True when a holder will drain it. + + Audio captured within ``FRESH_CAPTURE_SECONDS`` takes the fresh lane, + scored by capture time; anything else is backlog, scored by commit. + False (not engaged, shard unheld, push failed) means drain it locally. + """ + if not self.engaged: + return False + key = owner_key(owner_id) + if not self.shards.covered(key): + return False + now = self.shards.clock() + if capture_at is not None and now - capture_at <= FRESH_CAPTURE_SECONDS: + lane, score = FRESH, float(capture_at) + else: + lane, score = BACKLOG, float(committed_at) + pushed = self.shards.push(owner_id, key, lane, score) + PROJECTION_SHARD_WAKES.labels(lane=lane, outcome='pushed' if pushed else 'push_failed').inc() + return pushed + + def committed(self, owner_id: str, entries: list[dict]) -> None: + """Merge committed ``entries`` for ``owner_id``: wake its shard holder. Blocking.""" + if not entries or not self.engaged: + return # Not engaged: this process sweeps, as before shards. + now = self.shards.clock() + captures = [ + float(e['capture_at']) + for e in entries + if isinstance(e.get('capture_at'), (int, float)) and not isinstance(e.get('capture_at'), bool) + ] + fresh = [c for c in captures if now - c <= FRESH_CAPTURE_SECONDS] + commits = [float(e['committed_at']) for e in entries if isinstance(e.get('committed_at'), (int, float))] + capture_at = max(fresh) if fresh else None + if not self.wake(owner_id, capture_at, min(commits) if commits else now): + # Unheld shard or Redis error: followers no longer sweep, so drain here. + self.projector.request_local_drain(owner_id) + + def backstop(self) -> dict: + """Leader-only discovery that re-pushes wakes for pending owners. Blocking. + + First the tail (rows committed since the last pass, which covers the + GPU worker's own merge commits), then a bounded slice of the whole + pending log from a resuming cursor (lost wakes, pre-switch rows). + """ + + def route(owner_id: str, committed_at: float, capture_at: float | None) -> None: + if not self.wake(owner_id, capture_at, committed_at): + self.projector.request_local_drain(owner_id) + + tail = self.projector.discover_tail( + route, pages=TAIL_PAGES, overlap_seconds=TAIL_OVERLAP_SECONDS, slack_seconds=TAIL_SLACK_SECONDS + ) + counts = self.projector.discover_owners(route, pages=BACKSTOP_PAGES, budget_seconds=BACKSTOP_SECONDS) + logger.info( + 'event=batch_projection_shards action=backstop tail_rows=%d tail_owners=%d pages=%d rows=%d owners=%d' + ' wrapped=%d', + tail['rows'], + tail['owners'], + counts['pages'], + counts['rows'], + counts['owners'], + counts['wrapped'], + ) + return {**counts, 'tail_rows': tail['rows'], 'tail_owners': tail['owners']} + + # Event loop ----------------------------------------------------------- + + def attach(self) -> None: + drains = self.projector.ensure_drains() + drains.on_idle = self.idle + self.attached = True + + def idle(self, owner: str, stop: str) -> None: + """OwnerDrains hook: ``owner`` left the drain after a turn ending ``stop``.""" + wakes = self.inflight.pop(owner, None) + if wakes is None: + return + if stop == 'end': + self.settling.extend(wakes) + return + delay = (CLAIM_RETRY_SECONDS if stop in CLAIM_STOPS else STUCK_RETRY_SECONDS) * random.uniform(0.8, 1.2) + if len(self.cooldown) < MAX_COOLDOWN: + self.cooldown[owner] = self.monotonic() + delay + if stop not in CLAIM_STOPS: + # A stuck head keeps its wake but gives up its place in the lane, so + # stuck owners cannot fill the take pages ahead of due owners. + retry_at = self.shards.clock() + delay + self.deferring.extend((wake, retry_at) for wake in wakes) + PROJECTION_SHARD_WAKES.labels(lane=wakes[0].lane, outcome='cooled').inc() + + async def tick(self) -> int: + """One pump pass; returns the number of owners handed to the drain.""" + if not self.attached: + self.attach() + shards = self.shards + if not self.switches_on: + # A kill switch took this process off the sharded path. It must not + # keep shard leases: members route wakes to holders, and these lanes + # would wait unread. Hand them back so an engaged member takes them. + if shards.held or shards.healthy: + await run_blocking(db_executor, shards.release) + PROJECTION_SHARDS_HELD.set(0) + if self.projector.owes_drains: + self.projector.start_drains() + return 0 + if shards.refresh_due(): + await run_blocking(db_executor, shards.refresh) + PROJECTION_SHARDS_HELD.set(len(shards.held)) + PROJECTION_SHARDS_UNHELD.set(len(shards.unheld)) + if self.settling: + settling, self.settling = self.settling, [] + settled = await run_blocking(db_executor, lambda: [shards.settle(w) for w in settling]) + for wake, done in zip(settling, settled): + PROJECTION_SHARD_WAKES.labels(lane=wake.lane, outcome='settled' if done else 'repushed').inc() + if self.deferring: + deferring, self.deferring = self.deferring, [] + await run_blocking(db_executor, lambda: [shards.defer(w, score) for w, score in deferring]) + projector = self.projector + if projector.owes_drains: + # Local requests: unheld-shard fallback and Redis failures. + projector.start_drains() + if not self.engaged or not shards.held: + return 0 + drains = projector.ensure_drains() + free = drains.free_slots + if free <= 0: + return 0 + now = self.monotonic() + if self.cooldown: + self.cooldown = {owner: until for owner, until in self.cooldown.items() if until > now} + # Per held shard: this pass's share of the slots, plus room to skip + # owners already draining or cooling (claim-busy owners keep their place). + held = max(1, len(shards.held)) + fetch = min(MAX_TAKE, -(-free * 4 // held) + len(self.cooldown) + len(self.inflight)) + fresh, backlog = await run_blocking( + db_executor, lambda: (shards.take(FRESH, fetch), shards.take(BACKLOG, fetch)) + ) + picked = self.pick(fresh, backlog, free, drains) + for owner, wakes in picked.items(): + self.inflight[owner] = wakes + PROJECTION_SHARD_WAKES.labels(lane=wakes[0].lane, outcome='taken').inc() + if picked: + drains.request(list(picked)) + return len(picked) + + def pick(self, fresh: list[Wake], backlog: list[Wake], free: int, drains) -> dict[str, list[Wake]]: + """Fresh first, keeping ``backlog_reserve`` slots for waiting backlog.""" + now = self.monotonic() + + settling = {w.owner_id for w in self.settling} | {w.owner_id for w, _ in self.deferring} + + def eligible(wake: Wake) -> bool: + owner = wake.owner_id + return ( + owner not in self.inflight + and owner not in settling + and not drains.busy(owner) + and owner not in self.projector.draining + and self.cooldown.get(owner, 0.0) <= now + ) + + fresh_ok = [w for w in fresh if eligible(w)] + backlog_ok = [w for w in backlog if eligible(w)] + reserve = min(projection_backlog_reserve(), free) if backlog_ok else 0 + picked: dict[str, list[Wake]] = {} + + def take(wakes: list[Wake], room: int) -> None: + for wake in wakes: + if wake.owner_id in picked: + continue + if len(picked) >= room: + return + picked[wake.owner_id] = [wake] + + take(fresh_ok, free - reserve) + take(backlog_ok, free) + take(fresh_ok, free) # Backlog ran short: give fresh the reserve back. + # A picked owner waiting in the other lane settles with it when its drain ends. + for wake in fresh_ok + backlog_ok: + wakes = picked.get(wake.owner_id) + if wakes is not None and all(w.lane != wake.lane for w in wakes): + wakes.append(wake) + return picked + + async def depth(self) -> None: + counts = await run_blocking(db_executor, self.shards.depth) + for lane, count in counts.items(): + PROJECTION_SHARD_WAKE_DEPTH.labels(lane=lane).set(count) diff --git a/backend/utils/batch_pipeline/projection_shards.py b/backend/utils/batch_pipeline/projection_shards.py new file mode 100644 index 00000000000..93dd871a914 --- /dev/null +++ b/backend/utils/batch_pipeline/projection_shards.py @@ -0,0 +1,471 @@ +"""Owner-hash projection shards: one drain writer per owner across the fleet. + +Every serving instance used to page the same global ``committed_at`` keyset and +race the others for each owner's advisory Redis claim, so most drain turns +ended ``claim_busy`` or found nothing. With shards on: + +- Owners hash to one of ``count`` shards. Each shard has a short Redis lease + held by at most one member instance; members balance to + ``ceil(count / live members)`` shards each, and a free shard is taken on the + next lease pass by any member below its share. +- A merge commit pushes a wake for the owner into its shard's lane (``fresh`` + for recently captured audio, ``backlog`` otherwise). Only the shard holder + pulls wakes and drains those owners, so drains stop competing for claims. +- The maintenance leader's discovery-only backstop pages the durable result + log and re-pushes wakes, which covers Redis loss and rows committed before + the switch. + +Nothing here is a correctness fence. Assignment ordering is still the +transactional predecessor fence and the advisory owner claim still wraps each +walk step. Every Redis error fails open: the process reports itself inactive +and the caller runs the legacy per-instance sweep, exactly as before. +""" + +from __future__ import annotations + +from functools import lru_cache +import logging +import math +import random +import threading +import time +from typing import Any, Callable, Iterable +import uuid + +from database import redis_db + +logger = logging.getLogger(__name__) + +REDIS_TIMEOUT_SECONDS = 1.0 +# Every serving process must agree on the count; a change re-forms the fleet +# (keys carry the count, so old and new layouts never share a lease). +PROJECTION_SHARD_COUNT = 64 +LEASE_SECONDS = 30.0 +# Lease passes run this often; a holder that misses two passes keeps its lease. +LEASE_REFRESH_SECONDS = LEASE_SECONDS / 3 +# Idle wake keys expire; every push refreshes the TTL. +WAKE_TTL_SECONDS = 7 * 24 * 3600 +FRESH = 'fresh' +BACKLOG = 'backlog' +LANES = (FRESH, BACKLOG) + +# KEYS: lease keys. ARGV[1]: holder, ARGV[2]: TTL ms. Renews only leases this +# holder still owns; returns 1/0 per key. +_RENEW_LUA = """ +local out = {} +for i, key in ipairs(KEYS) do + if redis.call('get', key) == ARGV[1] then + redis.call('pexpire', key, ARGV[2]) + out[i] = 1 + else + out[i] = 0 + end +end +return out +""" +_RELEASE_LUA = """ +local n = 0 +for _, key in ipairs(KEYS) do + if redis.call('get', key) == ARGV[1] then + n = n + redis.call('del', key) + end +end +return n +""" +# KEYS[1]: lane zset, KEYS[2]: its push-sequence hash. ARGV: member, score, +# mode ('gt' keeps the newest score, 'lt' the oldest), ttl seconds. Every push +# bumps the member's sequence, so a drain that took an older sequence cannot +# settle away a wake pushed after it looked. Plain ZADD GT/LT needs Redis 6.2. +_PUSH_LUA = """ +local current = redis.call('zscore', KEYS[1], ARGV[1]) +local score = tonumber(ARGV[2]) +if not current + or (ARGV[3] == 'gt' and score > tonumber(current)) + or (ARGV[3] == 'lt' and score < tonumber(current)) then + redis.call('zadd', KEYS[1], score, ARGV[1]) +end +local seq = redis.call('hincrby', KEYS[2], ARGV[1], 1) +redis.call('expire', KEYS[1], ARGV[4]) +redis.call('expire', KEYS[2], ARGV[4]) +return seq +""" +# KEYS[1]: lane zset, KEYS[2]: sequence hash. ARGV[1]: 'asc' or 'desc', ARGV[2]: limit. +# Returns member, score, sequence triples. +_TAKE_LUA = """ +local page +if ARGV[1] == 'desc' then + page = redis.call('zrevrange', KEYS[1], 0, tonumber(ARGV[2]) - 1, 'WITHSCORES') +else + page = redis.call('zrange', KEYS[1], 0, tonumber(ARGV[2]) - 1, 'WITHSCORES') +end +local out = {} +for i = 1, #page, 2 do + out[#out + 1] = page[i] + out[#out + 1] = page[i + 1] + out[#out + 1] = redis.call('hget', KEYS[2], page[i]) or '0' +end +return out +""" +# Removes a wake only when no push landed since it was taken. The sequence +# itself stays (it expires with the lane): resetting it would let a stale take +# of the same owner match a later push's number and settle that push away. +_SETTLE_LUA = """ +local seq = redis.call('hget', KEYS[2], ARGV[1]) or '0' +if seq == ARGV[2] then + return redis.call('zrem', KEYS[1], ARGV[1]) +end +return 0 +""" + + +# Moves a cooled wake behind every wake that is due sooner. KEYS[1]/KEYS[2]: +# the source lane zset and sequence hash, KEYS[3]/KEYS[4]: the backlog lane's. +# ARGV: member, the sequence taken, retry score, ttl seconds. A push that +# landed since the take keeps the wake where it is (it is new work). Backlog +# keeps the earlier score when the member already waits there. +_DEFER_LUA = """ +local seq = redis.call('hget', KEYS[2], ARGV[1]) or '0' +if seq ~= ARGV[2] then + return 0 +end +local score = tonumber(ARGV[3]) +if KEYS[1] == KEYS[3] then + redis.call('zadd', KEYS[3], score, ARGV[1]) +else + redis.call('zrem', KEYS[1], ARGV[1]) + local current = redis.call('zscore', KEYS[3], ARGV[1]) + if not current or score < tonumber(current) then + redis.call('zadd', KEYS[3], score, ARGV[1]) + end +end +redis.call('hincrby', KEYS[4], ARGV[1], 1) +redis.call('expire', KEYS[3], ARGV[4]) +redis.call('expire', KEYS[4], ARGV[4]) +return 1 +""" + + +@lru_cache(maxsize=1) +def _redis_client(): + # Bounded socket I/O: an advisory coordinator must fail open promptly. + return redis_db.create_bounded_redis_client(REDIS_TIMEOUT_SECONDS) + + +def _text(value) -> str: + return value.decode() if isinstance(value, bytes) else str(value) + + +def shard_of(owner_key: str, count: int) -> int: + """Stable shard of a hashed owner key (hex sha256, see scheduler.owner_key).""" + return int(owner_key[:12], 16) % count + + +def share(count: int, members: int) -> int: + """Shards one member should hold so that the members cover every shard.""" + return math.ceil(count / max(1, members)) + + +class Wake: + __slots__ = ('owner_id', 'lane', 'shard', 'score', 'seq') + + def __init__(self, owner_id: str, lane: str, shard: int, score: float, seq: str = '0'): + self.owner_id, self.lane, self.shard, self.score, self.seq = owner_id, lane, shard, score, seq + + def __repr__(self) -> str: # never logged with ids; tests only + return f'Wake(lane={self.lane}, shard={self.shard})' + + +class ProjectionShards: + """Process-local view of shard leases plus the per-shard wake lanes.""" + + def __init__( + self, + namespace: str, + *, + count: int, + member: bool = True, + client: Callable[[], Any] = _redis_client, + clock: Callable[[], float] = time.time, + monotonic: Callable[[], float] = time.monotonic, + holder: str | None = None, + ): + if count < 1: + raise ValueError('at least one projection shard required') + self.namespace = namespace + self.count = count + self.member = member + self._client = client + self.clock = clock + self.monotonic = monotonic + # Random per-process identity; no pod, host or user data. + self.holder = holder or uuid.uuid4().hex + self.held: set[int] = set() + # Shards no member held at the last lease pass; their owners drain + # wherever they are discovered (the legacy behavior). + self.unheld: set[int] = set(range(count)) + self.members = 0 + # False until a lease pass succeeds, and after any Redis error: the + # caller then runs the legacy sweep on this instance. + self.healthy = False + self.refreshed_at: float | None = None + self._lock = threading.Lock() + + # Keys ----------------------------------------------------------------- + + def _lease_key(self, shard: int) -> str: + return f'batch_projection_shard:{self.namespace}:{self.count}:{shard}' + + def _members_key(self) -> str: + return f'batch_projection_shard_members:{self.namespace}:{self.count}' + + def _lane_key(self, lane: str, shard: int) -> str: + return f'batch_projection_wake:{self.namespace}:{self.count}:{lane}:{shard}' + + def shard(self, owner_key: str) -> int: + return shard_of(owner_key, self.count) + + @property + def active(self) -> bool: + """Sharded serving applies on this instance: Redis answered recently.""" + if not self.healthy or self.refreshed_at is None: + return False + return self.monotonic() - self.refreshed_at < LEASE_SECONDS + + def owns(self, owner_key: str) -> bool: + return self.shard(owner_key) in self.held + + def covered(self, owner_key: str) -> bool: + """Some member held this owner's shard at the last lease pass.""" + return self.shard(owner_key) not in self.unheld + + # Leases --------------------------------------------------------------- + + def refresh_due(self) -> bool: + return self.refreshed_at is None or self.monotonic() - self.refreshed_at >= LEASE_REFRESH_SECONDS + + def refresh(self) -> bool: + """One lease pass: heartbeat, renew, rebalance, acquire. Fails open.""" + with self._lock: + try: + self._refresh() + except Exception as error: + self.healthy = False + self.held = set() + self.unheld = set(range(self.count)) + logger.warning( + 'event=batch_projection_shards action=refresh outcome=unavailable exception_type=%s', + type(error).__name__, + ) + return False + self.healthy = True + self.refreshed_at = self.monotonic() + return True + + def _refresh(self) -> None: + client = self._client() + now = self.clock() + keys = [self._lease_key(shard) for shard in range(self.count)] + ttl_ms = int(LEASE_SECONDS * 1000) + pipe = client.pipeline(transaction=False) + if self.member: + pipe.zadd(self._members_key(), {self.holder: now}) + pipe.expire(self._members_key(), int(LEASE_SECONDS * 4)) + pipe.zremrangebyscore(self._members_key(), 0, now - LEASE_SECONDS) + pipe.zcard(self._members_key()) + pipe.mget(keys) + results = pipe.execute() + members, values = int(results[-2]), results[-1] + self.members = members + holders = [value.decode() if isinstance(value, bytes) else value for value in values] + mine = {shard for shard, holder in enumerate(holders) if holder == self.holder} + if not self.member: + if mine: + client.eval(_RELEASE_LUA, len(mine), *(keys[s] for s in sorted(mine)), self.holder) + self.held = set() + self.unheld = {shard for shard, holder in enumerate(holders) if not holder} + return + if mine: + renewed = client.eval(_RENEW_LUA, len(mine), *(keys[s] for s in sorted(mine)), self.holder, ttl_ms) + mine = {shard for shard, ok in zip(sorted(mine), renewed) if int(ok)} + target = share(self.count, members) + if len(mine) > target: + # Members joined: hand back the excess so they can take it on their + # next pass. In-flight drains finish their turn under the claim. + extra = random.sample(sorted(mine), len(mine) - target) + client.eval(_RELEASE_LUA, len(extra), *(keys[s] for s in extra), self.holder) + for shard in extra: + mine.discard(shard) + holders[shard] = None + free = [shard for shard, holder in enumerate(holders) if not holder and shard not in mine] + random.shuffle(free) + for shard in free: + if len(mine) >= target: + break + if client.set(keys[shard], self.holder, nx=True, px=ttl_ms): + mine.add(shard) + holders[shard] = self.holder + self.held = mine + self.unheld = {shard for shard, holder in enumerate(holders) if not holder} + + def release(self) -> None: + """Hand every held shard back and leave the members set, so others rebalance now. + + Used at shutdown and when a kill switch takes this process off the + sharded path. Fails open: unreleased leases lapse after LEASE_SECONDS. + """ + with self._lock: + held, self.held = sorted(self.held), set() + self.healthy = False + self.refreshed_at = None # Rejoin on the very next lease pass. + try: + client = self._client() + if held: + client.eval(_RELEASE_LUA, len(held), *(self._lease_key(s) for s in held), self.holder) + client.zrem(self._members_key(), self.holder) + except Exception as error: + logger.warning( + 'event=batch_projection_shards action=release outcome=unavailable exception_type=%s', + type(error).__name__, + ) + + # Wakes ---------------------------------------------------------------- + + def _seq_key(self, lane: str, shard: int) -> str: + return f'batch_projection_wake_seq:{self.namespace}:{self.count}:{lane}:{shard}' + + def push(self, owner_id: str, owner_key: str, lane: str, score: float) -> bool: + """Record that ``owner_id`` has projection work. Fails open (False).""" + if lane not in LANES: + raise ValueError('unknown projection wake lane') + shard = self.shard(owner_key) + try: + self._client().eval( + _PUSH_LUA, + 2, + self._lane_key(lane, shard), + self._seq_key(lane, shard), + owner_id, + float(score), + 'gt' if lane == FRESH else 'lt', + WAKE_TTL_SECONDS, + ) + return True + except Exception as error: + logger.warning( + 'event=batch_projection_shards action=push outcome=unavailable exception_type=%s', + type(error).__name__, + ) + return False + + def take(self, lane: str, limit: int, shards: Iterable[int] | None = None) -> list[Wake]: + """Read (never remove) up to ``limit`` wakes per held shard, best first. + + Fresh wakes come newest capture first; backlog wakes oldest commit first. + Shards interleave so one deep shard cannot crowd out the others. + """ + shards = sorted(self.held if shards is None else shards) + if not shards or limit <= 0: + return [] + try: + pipe = self._client().pipeline(transaction=False) + for shard in shards: + pipe.eval( + _TAKE_LUA, + 2, + self._lane_key(lane, shard), + self._seq_key(lane, shard), + 'desc' if lane == FRESH else 'asc', + limit, + ) + pages = pipe.execute() + except Exception as error: + # Lanes this process cannot read are lanes nobody drains: go + # inactive so this instance runs the legacy sweep until a lease + # pass succeeds again. + self.healthy = False + logger.warning( + 'event=batch_projection_shards action=take outcome=unavailable exception_type=%s', + type(error).__name__, + ) + return [] + columns = [] + for shard, page in zip(shards, pages): + column = [] + for index in range(0, len(page) - 2, 3): + member, score, seq = page[index], page[index + 1], page[index + 2] + column.append(Wake(_text(member), lane, shard, float(_text(score)), _text(seq))) + columns.append(column) + merged: list[Wake] = [] + for row in range(limit): + for column in columns: + if row < len(column): + merged.append(column[row]) + return merged + + def settle(self, wake: Wake) -> bool: + """Drop a wake whose owner drained to the end, unless it was re-pushed.""" + try: + return bool( + self._client().eval( + _SETTLE_LUA, + 2, + self._lane_key(wake.lane, wake.shard), + self._seq_key(wake.lane, wake.shard), + wake.owner_id, + wake.seq, + ) + ) + except Exception as error: + logger.warning( + 'event=batch_projection_shards action=settle outcome=unavailable exception_type=%s', + type(error).__name__, + ) + return False + + def defer(self, wake: Wake, score: float) -> bool: + """Re-queue a cooled wake in the backlog lane at retry ``score``. + + A stuck owner would otherwise keep its place at the head of its lane, + and enough of them would fill every take page. Unless a push landed + since the take, the wake moves behind every owner due sooner. Fails + open (False): the wake stays where it was. + """ + try: + return bool( + self._client().eval( + _DEFER_LUA, + 4, + self._lane_key(wake.lane, wake.shard), + self._seq_key(wake.lane, wake.shard), + self._lane_key(BACKLOG, wake.shard), + self._seq_key(BACKLOG, wake.shard), + wake.owner_id, + wake.seq, + float(score), + WAKE_TTL_SECONDS, + ) + ) + except Exception as error: + logger.warning( + 'event=batch_projection_shards action=defer outcome=unavailable exception_type=%s', + type(error).__name__, + ) + return False + + def depth(self) -> dict[str, int]: + """Wake count per lane across the shards this process holds (gauge).""" + out: dict[str, int] = {lane: 0 for lane in LANES} + shards = sorted(self.held) + if not shards: + return out + try: + pipe = self._client().pipeline(transaction=False) + for lane in LANES: + for shard in shards: + pipe.zcard(self._lane_key(lane, shard)) + counts = pipe.execute() + except Exception: + return out + for index, lane in enumerate(LANES): + out[lane] = sum(int(c) for c in counts[index * len(shards) : (index + 1) * len(shards)]) + return out diff --git a/backend/utils/batch_pipeline/serving.py b/backend/utils/batch_pipeline/serving.py index d9b594f76e9..c8663030de0 100644 --- a/backend/utils/batch_pipeline/serving.py +++ b/backend/utils/batch_pipeline/serving.py @@ -9,6 +9,8 @@ from config.batch_pipeline import ( held_partial_enabled, + projection_shard_member, + projection_shards_enabled, owner_chain_projection_enabled, recovered_merge_release_enabled, recovery_owner_only_enabled, @@ -28,11 +30,13 @@ from utils.batch_pipeline.artifact_policy import artifacts_for from utils.batch_pipeline.live import drain_recordings from utils.batch_pipeline.projection import HELD_PARTIAL_PAGE, Projector, error_type -from utils.executors import run_blocking, start_background_task, sync_executor +from utils.executors import db_executor, run_blocking, start_background_task, sync_executor from utils.batch_pipeline.contracts import PoolCapacity, WorkResult from utils.batch_pipeline.leader import SERVING_MAINTENANCE, LeaderLease from utils.batch_pipeline.merge import OrderedMerge from utils.batch_pipeline.metrics import QueueMetrics +from utils.batch_pipeline.projection_shard_pump import TICK_SECONDS as SHARD_TICK_SECONDS, ShardPump +from utils.batch_pipeline.projection_shards import PROJECTION_SHARD_COUNT, ProjectionShards from utils.batch_pipeline.scheduler import Scheduler from utils.batch_pipeline.retention import cleanup_completed_artifacts from utils.sync.stage import current_stage @@ -71,8 +75,11 @@ def __init__( enrichment_sweep=None, recovery=None, leader=None, + shard_pump=None, ): self.store = store + # Sharded projection serving (projection_shard_pump.ShardPump), or None. + self.shard_pump = shard_pump self.metrics = metrics self.projector = projector self.cleanup = cleanup @@ -119,9 +126,22 @@ def drain_owner(self, owner_id: str, *, recovered: str | None = None): release = ( self.merge.release_set(recovered) if recovered is not None and recovered_merge_release_enabled() else None ) + committed = [] for _ in range(100): - if self.merge.commit_next(owner_id, time.time(), release=release) is None: + entry = self.merge.commit_next(owner_id, time.time(), release=release) + if entry is None: break + committed.append(entry) + if committed and self.shard_pump is not None: + # Push wake: the owner's shard holder drains it without waiting + # for discovery. Fail-open; the backstop re-finds lost wakes. + try: + self.shard_pump.committed(owner_id, committed) + except Exception as error: + logging.getLogger(__name__).warning( + 'event=batch_projection_shards action=commit_wake outcome=error exception_type=%s', + type(error).__name__, + ) @with_batch_privacy_attempt async def reconcile_async( @@ -182,7 +202,8 @@ def reconcile( if owner_id is None or recovered is None: raise ValueError('owner-only recovery requires an owner and recovered unit') return self._reconcile_owner(owner_id, recovered) - if self._leads(periodic): + leads = self._leads(periodic) + if leads: self._maintain(recover_quarantine=recover_quarantine) priority = [] if owner_id is not None and self.projector and owner_chain_projection_enabled(): @@ -193,8 +214,25 @@ def reconcile( # Projection and the local gauges stay per instance: projection is # throughput (ordered and fenced in its own transactions) and every # instance exports its own queue sample. - if self.projector and serving_projection_sweep_enabled(): - self.projector.sweep() + if self.projector: + pump = self.shard_pump + if pump is not None and pump.engaged: + # Sharded: shard holders drain their owners from wakes. Only the + # maintenance leader pages the durable log, discovery-only, to + # re-push wakes and to finish synchronous-policy rows. + if leads and serving_projection_sweep_enabled(): + try: + pump.backstop() + except Exception as error: + logging.getLogger(__name__).error( + 'event=batch_projection_shards action=backstop outcome=retry exception_type=%s', + error_type(error), + ) + self.projector.sweep_finish_only() + # Follow-up discovery keeps its own per-process page as before. + self.projector.has_async_followups = True + elif serving_projection_sweep_enabled(): + self.projector.sweep() # Catch-up rows this instance held for projection depth may now fit. for held_owner in list(self.merge.held_owners): try: @@ -348,8 +386,18 @@ def serving_for_stage(stage: str) -> Serving: projector = Projector(store, artifacts.load_result, fetch_audio=artifacts.fetch_audio) from utils.batch_pipeline.quarantine_recovery import QuarantineRecovery + shard_pump = None + if projection_shards_enabled(): + shards = ProjectionShards( + getattr(store, 'namespace', f'batch_pipeline_{stage}'), + count=PROJECTION_SHARD_COUNT, + member=projection_shard_member(), + ) + shard_pump = ShardPump(projector, shards) + projector.router = shard_pump.route return Serving( store, + shard_pump=shard_pump, recovery=QuarantineRecovery(store, artifacts), metrics=_metrics(), projector=projector, @@ -471,11 +519,53 @@ async def dispatch_recovery(): except asyncio.TimeoutError: pass + async def shard_pump(): + # Sharded projection: feed this holder's owner drains from its wakes. + + service = None + while not stop.is_set() and service is None: + try: + service = await run_blocking(sync_executor, serving_for_stage, current_stage()) + except Exception as error: + logging.getLogger(__name__).error( + 'event=batch_projection_shards action=start outcome=retry exception_type=%s', type(error).__name__ + ) + try: + await asyncio.wait_for(stop.wait(), timeout=20) + except asyncio.TimeoutError: + pass + pump = service.shard_pump if service is not None else None + if pump is None: + return + depth_due = 0.0 + try: + while not stop.is_set(): + try: + await pump.tick() + if time.monotonic() >= depth_due: + depth_due = time.monotonic() + 60 + await pump.depth() + except Exception as error: + logging.getLogger(__name__).error( + 'event=batch_projection_shards action=tick outcome=retry exception_type=%s', + type(error).__name__, + ) + try: + await asyncio.wait_for(stop.wait(), timeout=SHARD_TICK_SECONDS) + except asyncio.TimeoutError: + pass + finally: + # Hand shards back so the next holder starts at once on a roll. + await run_blocking(db_executor, pump.shards.release) + async def run(): if pump_enabled: await projection_pump() else: - await asyncio.gather(supervise(), drain(), dispatch_recovery()) + loops = [supervise(), drain(), dispatch_recovery()] + if projection_shards_enabled(): + loops.append(shard_pump()) + await asyncio.gather(*loops) return stop, start_background_task( run(), name='batch_projection_pump' if pump_enabled else 'batch_pipeline_serving' diff --git a/config/feature-flags.yaml b/config/feature-flags.yaml index 5e620f7746f..f0bccb74498 100644 --- a/config/feature-flags.yaml +++ b/config/feature-flags.yaml @@ -218,6 +218,11 @@ flags: - {"key":"BATCH_PIPELINE_MAX_OWNER_DRAINS","aliases":[],"kind":"env","lifecycle":"ops_kill","surfaces":["backend"],"summary":"Default 8; concurrent owner drains per process, clamped to 1-16 at import; malformed keeps 8","fail":"closed","owner":"lead2","created":"2026-11-05","terminal":"keep","decision":"keep","notes":"Code only; deployment shape pending David. Safe start: dedicated service only, N=4 instances, pages 1, tick 2 s, MAX_OWNER_DRAINS 8. Never enable the pump on backend-integration; the selector is service-wide. Lead2 owns rollout after the deployment-shape decision. Keep serving sweeps ON until owner_drain advanced >=20 rows/s for 30 min and a dedicated-service missing/stale omi_batch_projection_pump_last_tick_seconds alert is installed and verified. Alarm on Firestore read rate above +50% of pre-cutover; owner_drain rows/s must remain at least 70% of pre-cutover. Watch sweep elapsed time and executor saturation. Evaluate claim_busy share no higher than pre-cutover only after serving sweeps are switched off, not during dual-running. Rollback via env-only new revisions: serving sweeps ON on existing service first, pump false on dedicated service, owner drains 8."} - {"key":"BATCH_PIPELINE_OWNER_DRAIN_ENABLED","aliases":[],"kind":"env","lifecycle":"ops_kill","surfaces":["backend"],"summary":"Default on; a per-process owner drain walks each discovered owner's result log continuously (up to 8 owners at once, 120 s turns, FIFO re-queue) instead of the sweep-shared 30 s walk budget","fail":"open","owner":"dazheng","created":"2026-10-20","terminal":"keep","decision":"keep","notes":"False restores the sweep owner-chain walk with every pre-check; read-ahead (BATCH_PIPELINE_PROJECTION_PREFETCH=0) and the deferred index (BATCH_PIPELINE_PROJECTION_INDEX_DEFERRED_ENABLED=false) have their own switches, and all three off is the pre-drain behaviour. On: the sweep only discovers owners; each walk holds the advisory owner claim for its 15 s step and skips the predecessor/per-part pre-check transactions for rows it reached as the watermark's successor. The assignment transaction's predecessor, privacy and receipt fences stay authoritative; one owner never runs two chains. See backend/docs/batch-pipeline-serving.md."} - {"key":"PROJECTION_DRAIN_CLAIM_AWARE","aliases":[],"kind":"env","lifecycle":"ops_kill","surfaces":["backend"],"summary":"Default on; read-only fleet claim admission and jittered backoff for redundant claim-busy owner-drain requests","fail":"open","owner":"dazheng","created":"2026-11-05","terminal":"keep","decision":"keep","notes":"Only off disables this switch, including when CAPTURE_COVERAGE_MODE=serve; false, 0 and unknown values remain on. Initial admissions GET the existing Redis claim and fail open on any error. Backoff applies only to claim_busy turns re-requested through again; productive more=True continuations retain immediate FIFO progress without the peek. See backend/docs/listen_pusher_pipeline.mdx#batch-projection-drain-admission."} + - {"key":"BATCH_PROJECTION_SHARDS_ENABLED","aliases":[],"kind":"env","lifecycle":"rollout","surfaces":["backend"],"summary":"Default off; owner-hash projection shards: one serving member drains each owner from Redis wake lanes (fresh first), and only the maintenance leader pages the result log, discovery-only","fail":"closed","owner":"dazheng","created":"2026-11-05","review_by":"2026-11-19","terminal":"on","decision":"pending","ramp":{"cohorts":["staging backend-integration","prod backend-integration, every serving instance in one config train"],"clean_metric":"owner_drain rows advanced/s at or above the pre-switch hour for 60 min, claim_busy turns under 10%, omi_batch_projection_shards_unheld 0 after the roll, Firestore reads/s below the pre-switch hour","expander":"two-lane lead","done_by":"2026-11-12"},"notes":"Rollback: false (the per-instance sweep returns on the next 20 s tick). Fails open: any Redis error makes a process inactive and it runs the legacy sweep; owners of unheld shards drain where they are found. Advisory only: the transactional predecessor fence and owner claim still order every assignment. See backend/docs/listen_pusher_pipeline.mdx#batch-projection-owner-shards."} + - {"key":"BATCH_PROJECTION_SHARD_MEMBER","aliases":[],"kind":"env","lifecycle":"config_switch","surfaces":["backend"],"summary":"Default true; with projection shards on, false keeps a serving process out of shard leases (it serves, routes and pushes wakes only)","fail":"open","owner":"dazheng","created":"2026-11-05","terminal":"keep","decision":"keep","notes":"Lets a dedicated drain pool hold every shard while request-serving instances opt out. At least one member must run or every owner falls back to local drains. See backend/docs/listen_pusher_pipeline.mdx#batch-projection-owner-shards."} + - {"key":"BATCH_PROJECTION_BACKLOG_RESERVE","aliases":[],"kind":"env","lifecycle":"config_switch","surfaces":["backend"],"summary":"Default 2; drain slots a projection shard holder keeps for backlog wakes while fresh wakes wait","fail":"closed","owner":"dazheng","created":"2026-11-05","terminal":"keep","decision":"keep","notes":"Bounds 0..16; malformed keeps 2. Fresh wakes (capture within 6 h) take every other slot newest capture first; reserved slots return to fresh whenever no backlog waits. See backend/docs/listen_pusher_pipeline.mdx#batch-projection-owner-shards."} + - {"key":"BATCH_PROJECTION_HEAD_DEADLINE_SECONDS","aliases":[],"kind":"env","lifecycle":"config_switch","surfaces":["backend"],"summary":"Default 86400; how long a failing (or, with the blocked-head deadline on, blocked) owner-chain head retries before it takes the held disposition","fail":"closed","owner":"dazheng","created":"2026-11-05","terminal":"keep","decision":"keep","notes":"Bounds 3600..604800; malformed or out-of-range keeps 86400, the previous constant. A held head keeps its audio and result, advances the chain, and is owned by the release scan and reviewed redrive. See backend/docs/listen_pusher_pipeline.mdx#batch-projection-owner-shards."} + - {"key":"BATCH_PROJECTION_BLOCKED_HEAD_DEADLINE_ENABLED","aliases":[],"kind":"env","lifecycle":"rollout","surfaces":["backend"],"summary":"Default off; an owner-chain head refused by a never-passing ProjectionBlocked (closed allowlist) is held once the head deadline lapses, so later captures project","fail":"closed","owner":"dazheng","created":"2026-11-05","review_by":"2026-11-19","terminal":"on","decision":"pending","ramp":{"cohorts":["staging","prod, after David approves"],"clean_metric":"omi_batch_projection_blocked_head_holds_total bounded by the pre-switch count of heads blocked past the deadline; no hold outside the allowlist; held rows visible to the release scan","expander":"two-lane lead","done_by":"2026-11-12"},"notes":"Needs David's yes before prod. Nothing is deleted: the hold is the existing quarantine_held disposition (audio and result kept, redrivable). Claim, shadow, privacy, predecessor and wiring refusals never qualify. See backend/docs/listen_pusher_pipeline.mdx#batch-projection-owner-shards."} - {"key":"BATCH_PIPELINE_HELD_TERMINAL_STATUS_ENABLED","aliases":[],"kind":"env","lifecycle":"ops_kill","surfaces":["backend"],"summary":"Default on; terminally held batch uploads retain phone custody as processing/batch_unit_held with a five-minute poll hint","fail":"open","owner":"dazheng","created":"2026-10-20","terminal":"keep","decision":"keep","ramp":{"cohorts":["all batch uploads at deploy (default on)"],"clean_metric":"Zero completed status responses for uploads with unresolved held units; projected siblings remain visible and held polls are paced at five minutes","expander":"dazheng (two-lane lead)","done_by":"2026-10-27"},"notes":"Read-only status projection; no ledger, job or retention write. Phones keep their independent WAL copy until actual completion; server artifact retention alone never permits acknowledgement. Open operator intents or in-flight siblings also remain processing. False restores failed/batch_worker_quarantined and still never acknowledges custody. See backend/docs/batch-pipeline-serving.md."} - {"key":"BATCH_PIPELINE_RECOVERY_POLICY_REFRESH_ENABLED","aliases":[],"kind":"env","lifecycle":"ops_kill","surfaces":["backend"],"summary":"Default on; a quarantined held unit whose recorded recovery policy is disabled and which has no recovery intent (no paid call authorized) re-resolves the recovery setting at sweep or dispatch, and adopts it only when it is now enabled","fail":"open","owner":"dazheng","created":"2026-10-20","terminal":"keep","decision":"keep","ramp":{"cohorts":["all held units at deploy (default on)"],"clean_metric":"recovery sweep outcome=disabled count falls to zero and paid recovery starts stay within the per-owner allowance","expander":"dazheng (two-lane lead)","done_by":"2026-10-27"},"notes":"Spends at most the one at-most-once paid recovery the unit would have had if admitted while recovery was on; the existing allowance/cap gates and RecoveryStore.start fence still apply. An enabled snapshot is never re-resolved or downgraded. False restores the frozen snapshot (outcome=disabled forever). See backend/docs/batch-pipeline-serving.md."} - {"key":"BATCH_PIPELINE_HELD_PARTIAL_ENABLED","aliases":[],"kind":"env","lifecycle":"ops_kill","surfaces":["backend"],"summary":"Default on; publish healthy siblings and charge their speech through a partial ledger while unresolved units retain phone custody","fail":"open","owner":"dazheng","created":"2026-10-20","terminal":"keep","decision":"keep","ramp":{"cohorts":["all batch uploads at deploy (default on)"],"clean_metric":"batch_v1 ledgers in processing over 6 h fall to the count with a sibling still recovering, with no new projection_completed regressions and zero paid recovery calls from the scan","expander":"dazheng (two-lane lead)","done_by":"2026-10-27"},"notes":"The ledger remains partial, owned by its batch job, with failed_segments=held_segments. Status polls and duplicate uploads return processing/batch_unit_held with projected IDs and a five-minute poll hint. No provider, job, intent, merge or retention mutation from the scan. Recovery upgrades to completed and charges only remaining speech. False stops new partials; existing partials remain nonterminal and can still upgrade. See backend/docs/batch-pipeline-serving.md."}