@pentatonic-ai/ai-agent-sdk 0.10.37 → 0.10.39
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- package/dist/index.cjs +1 -1
- package/dist/index.js +1 -1
- package/package.json +1 -1
- package/packages/memory-engine-v2/PLAN-async-escalation.md +497 -0
- package/packages/memory-engine-v2/RFC-async-escalation.md +218 -0
- package/packages/memory-engine-v2/extractor-async/test_async_escalation.py +314 -0
- package/packages/memory-engine-v2/extractor-async/worker.py +247 -12
- package/packages/memory-engine-v2/org-model/migrations/013_distillation_transient_attempts.sql +27 -0
- package/packages/memory-engine-v2/org-model/migrations/014_escalated_pending.sql +22 -0
- package/packages/memory-engine-v2/scripts/redistill.py +25 -9
- package/packages/memory-engine-v2/scripts/test_redistill.py +81 -0
|
@@ -116,6 +116,11 @@ def _envflag(name: str, default: str = "false") -> bool:
|
|
|
116
116
|
# event_distillations ledger (migration 010) so escalation can be monitored and
|
|
117
117
|
# student rows can be re-distilled later ("mop up").
|
|
118
118
|
CASCADE_ENABLED = _envflag("CASCADE_ENABLED")
|
|
119
|
+
# Async escalation (RFC-async-escalation.md). When true, the student-primary
|
|
120
|
+
# pass does NOT call the teacher inline; escalated events become
|
|
121
|
+
# 'escalated_pending' for the teacher-mop-up consumer. Default false = today's
|
|
122
|
+
# synchronous inline cascade (instant revert).
|
|
123
|
+
ASYNC_ESCALATION_ENABLED = os.environ.get("ASYNC_ESCALATION_ENABLED", "") not in ("", "0", "false")
|
|
119
124
|
STUDENT_ENDPOINT = os.environ.get("STUDENT_ENDPOINT", "")
|
|
120
125
|
STUDENT_API_KEY = os.environ.get("STUDENT_API_KEY", "")
|
|
121
126
|
STUDENT_MODEL = os.environ.get("STUDENT_MODEL", "nuextract-2.0-4b-ft")
|
|
@@ -2622,7 +2627,7 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
|
|
|
2622
2627
|
SELECT DISTINCT client FROM distillation_queue
|
|
2623
2628
|
WHERE (
|
|
2624
2629
|
status = 'pending'
|
|
2625
|
-
OR (status = 'claimed' AND claim_expires_at < NOW())
|
|
2630
|
+
OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NULL)
|
|
2626
2631
|
) AND attempts < %s
|
|
2627
2632
|
) c
|
|
2628
2633
|
CROSS JOIN LATERAL (
|
|
@@ -2631,7 +2636,7 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
|
|
|
2631
2636
|
WHERE client IS NOT DISTINCT FROM c.client
|
|
2632
2637
|
AND (
|
|
2633
2638
|
status = 'pending'
|
|
2634
|
-
OR (status = 'claimed' AND claim_expires_at < NOW())
|
|
2639
|
+
OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NULL)
|
|
2635
2640
|
) AND attempts < %s
|
|
2636
2641
|
ORDER BY COALESCE(priority, 100), emit_ts, id
|
|
2637
2642
|
LIMIT %s
|
|
@@ -2642,11 +2647,11 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
|
|
|
2642
2647
|
))
|
|
2643
2648
|
AND (
|
|
2644
2649
|
status = 'pending'
|
|
2645
|
-
OR (status = 'claimed' AND claim_expires_at < NOW())
|
|
2650
|
+
OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NULL)
|
|
2646
2651
|
) AND attempts < %s
|
|
2647
2652
|
FOR UPDATE SKIP LOCKED
|
|
2648
2653
|
)
|
|
2649
|
-
RETURNING id, event_id, attempts
|
|
2654
|
+
RETURNING id, event_id, attempts, transient_attempts
|
|
2650
2655
|
""",
|
|
2651
2656
|
(WORKER_ID, CLAIM_TTL_SEC, MAX_ATTEMPTS, MAX_ATTEMPTS,
|
|
2652
2657
|
BATCH_SIZE, BATCH_SIZE, MAX_ATTEMPTS),
|
|
@@ -2654,6 +2659,47 @@ def claim_next_batch(conn: psycopg.Connection) -> list[dict[str, Any]]:
|
|
|
2654
2659
|
return cur.fetchall()
|
|
2655
2660
|
|
|
2656
2661
|
|
|
2662
|
+
def claim_escalated_batch(
|
|
2663
|
+
conn: psycopg.Connection,
|
|
2664
|
+
worker_id: str,
|
|
2665
|
+
batch_size: int,
|
|
2666
|
+
claim_ttl_sec: int,
|
|
2667
|
+
max_attempts: int,
|
|
2668
|
+
) -> list[dict[str, Any]]:
|
|
2669
|
+
"""Teacher-mop-up claim: lease `escalated_pending` rows (priority, emit_ts,
|
|
2670
|
+
id order — no student fair-share needed, this is the teacher tier), leased
|
|
2671
|
+
the same way as `claim_next_batch` (claimed/claimed_by/claimed_at/
|
|
2672
|
+
claim_expires_at, FOR UPDATE SKIP LOCKED) so it is crash-safe. Also
|
|
2673
|
+
re-claims expired-claimed rows whose prior state was escalated (a mop-up
|
|
2674
|
+
worker that dies mid-batch leaves them `claimed`; once the lease expires
|
|
2675
|
+
they must come back through THIS claim, not the student's `claim_next_batch`
|
|
2676
|
+
— scoped via `escalation_reason IS NOT NULL`, which student-claimed rows
|
|
2677
|
+
never carry). Served by idx_distillation_escalated."""
|
|
2678
|
+
with conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
|
|
2679
|
+
cur.execute(
|
|
2680
|
+
"""
|
|
2681
|
+
UPDATE distillation_queue SET
|
|
2682
|
+
status = 'claimed',
|
|
2683
|
+
claimed_by = %s,
|
|
2684
|
+
claimed_at = NOW(),
|
|
2685
|
+
claim_expires_at = NOW() + (%s || ' seconds')::interval
|
|
2686
|
+
WHERE id IN (
|
|
2687
|
+
SELECT id FROM distillation_queue
|
|
2688
|
+
WHERE (
|
|
2689
|
+
(status = 'escalated_pending')
|
|
2690
|
+
OR (status = 'claimed' AND claim_expires_at < NOW() AND escalation_reason IS NOT NULL)
|
|
2691
|
+
) AND attempts < %s
|
|
2692
|
+
ORDER BY COALESCE(priority, 100), emit_ts, id
|
|
2693
|
+
LIMIT %s
|
|
2694
|
+
FOR UPDATE SKIP LOCKED
|
|
2695
|
+
)
|
|
2696
|
+
RETURNING id, event_id, attempts, transient_attempts, escalation_reason
|
|
2697
|
+
""",
|
|
2698
|
+
(worker_id, claim_ttl_sec, max_attempts, batch_size),
|
|
2699
|
+
)
|
|
2700
|
+
return cur.fetchall()
|
|
2701
|
+
|
|
2702
|
+
|
|
2657
2703
|
def fetch_event(conn: psycopg.Connection, event_id: str) -> dict[str, Any] | None:
|
|
2658
2704
|
with conn.cursor(row_factory=psycopg.rows.dict_row) as cur:
|
|
2659
2705
|
cur.execute(
|
|
@@ -2688,11 +2734,23 @@ def release_claim(conn: psycopg.Connection, queue_id: int, error: str) -> None:
|
|
|
2688
2734
|
# Recoverable genuine-failure path (will retry) → count the attempt. This is
|
|
2689
2735
|
# where the retry budget is spent — NOT at claim time — so a deploy-induced
|
|
2690
2736
|
# reclaim never consumes it.
|
|
2737
|
+
#
|
|
2738
|
+
# status: CASE-guarded rather than a bare 'pending'. This function is
|
|
2739
|
+
# reused by the teacher-mop-up path (_run_teacher, on non-exhausted
|
|
2740
|
+
# failures) — a row claimed out of `escalated_pending` still carries
|
|
2741
|
+
# `escalation_reason`. Sending it to bare 'pending' would let the
|
|
2742
|
+
# STUDENT's claim_next_batch (unconditional status='pending' branch)
|
|
2743
|
+
# claim and write it, silently defeating the escalation. Routing it back
|
|
2744
|
+
# to 'escalated_pending' instead means only claim_escalated_batch can
|
|
2745
|
+
# reclaim it. In SYNC mode escalation_reason is never set (see
|
|
2746
|
+
# mark_escalated_pending / ASYNC_ESCALATION_ENABLED gate), so the CASE
|
|
2747
|
+
# always yields 'pending' there — byte-for-byte identical to before.
|
|
2691
2748
|
with conn.cursor() as cur:
|
|
2692
2749
|
cur.execute(
|
|
2693
2750
|
"""
|
|
2694
2751
|
UPDATE distillation_queue SET
|
|
2695
|
-
status =
|
|
2752
|
+
status = CASE WHEN escalation_reason IS NOT NULL
|
|
2753
|
+
THEN 'escalated_pending' ELSE 'pending' END,
|
|
2696
2754
|
claimed_by = NULL,
|
|
2697
2755
|
claimed_at = NULL,
|
|
2698
2756
|
claim_expires_at = NULL,
|
|
@@ -2704,6 +2762,44 @@ def release_claim(conn: psycopg.Connection, queue_id: int, error: str) -> None:
|
|
|
2704
2762
|
)
|
|
2705
2763
|
|
|
2706
2764
|
|
|
2765
|
+
def mark_escalated_pending(
|
|
2766
|
+
conn: psycopg.Connection, queue_id: int, reason: str
|
|
2767
|
+
) -> None:
|
|
2768
|
+
"""Student ran; a teacher-only pass is owed. Move the row to
|
|
2769
|
+
'escalated_pending' WITHOUT touching `attempts` (the student didn't fail —
|
|
2770
|
+
it deferred) and clear the student's claim so the teacher-mop-up consumer
|
|
2771
|
+
can pick it up. The student's output is NOT written (deferred-write → the
|
|
2772
|
+
teacher's write is the first and only write; no supersede)."""
|
|
2773
|
+
with conn.cursor() as cur:
|
|
2774
|
+
cur.execute(
|
|
2775
|
+
"""
|
|
2776
|
+
UPDATE distillation_queue SET
|
|
2777
|
+
status = 'escalated_pending',
|
|
2778
|
+
escalation_reason = %s,
|
|
2779
|
+
claimed_by = NULL,
|
|
2780
|
+
claimed_at = NULL,
|
|
2781
|
+
claim_expires_at = NULL
|
|
2782
|
+
WHERE id = %s
|
|
2783
|
+
""",
|
|
2784
|
+
(reason, queue_id),
|
|
2785
|
+
)
|
|
2786
|
+
|
|
2787
|
+
|
|
2788
|
+
def _route_escalations(
|
|
2789
|
+
conn: psycopg.Connection,
|
|
2790
|
+
escalate_items: list[dict[str, Any]],
|
|
2791
|
+
reason_by_qid: dict[int, str],
|
|
2792
|
+
) -> list[dict[str, Any]]:
|
|
2793
|
+
"""Async: mark each escalated item 'escalated_pending' and return an empty
|
|
2794
|
+
inline-teacher batch. Sync (flag off): return the items unchanged so the
|
|
2795
|
+
caller runs the teacher inline exactly as before."""
|
|
2796
|
+
if not ASYNC_ESCALATION_ENABLED:
|
|
2797
|
+
return escalate_items
|
|
2798
|
+
for item in escalate_items:
|
|
2799
|
+
mark_escalated_pending(conn, item["id"], reason_by_qid.get(item["id"], "escalated"))
|
|
2800
|
+
return []
|
|
2801
|
+
|
|
2802
|
+
|
|
2707
2803
|
def _attempts_exhausted(attempts: int) -> bool:
|
|
2708
2804
|
"""Whether THIS processing failure should be terminal (mark_failed) rather
|
|
2709
2805
|
than retried (release_claim). `attempts` is the row's value at claim time =
|
|
@@ -2713,6 +2809,77 @@ def _attempts_exhausted(attempts: int) -> bool:
|
|
|
2713
2809
|
return attempts + 1 >= MAX_ATTEMPTS
|
|
2714
2810
|
|
|
2715
2811
|
|
|
2812
|
+
# ── transient vs genuine failures ─────────────────────────────────────
|
|
2813
|
+
# Transient = infra hiccups that resolve on their own: the teacher timing out
|
|
2814
|
+
# under load (ReadTimeout) or the teacher-lb returning 502/503/504 while the
|
|
2815
|
+
# fleet is cold-starting from zero. These must NOT burn the genuine `attempts`
|
|
2816
|
+
# budget (3) — 3 quick ReadTimeouts used to march a row to `failed`, and failed
|
|
2817
|
+
# rows are never requeued, so the extraction was silently lost. Instead they get
|
|
2818
|
+
# their own larger budget (MAX_TRANSIENT_ATTEMPTS) on the `transient_attempts`
|
|
2819
|
+
# counter (migration 013), leaving `attempts` untouched so the row stays
|
|
2820
|
+
# claimable (`attempts < MAX_ATTEMPTS` in claim_next_batch). ReadTimeouts are
|
|
2821
|
+
# self-spaced by LLM_TIMEOUT_SEC and the fair-share round-robin spreads re-claims
|
|
2822
|
+
# across the backlog, so this does not hot-loop; MAX_TRANSIENT_ATTEMPTS bounds a
|
|
2823
|
+
# genuine poison chunk (one that times out even on a healthy multi-box fleet).
|
|
2824
|
+
MAX_TRANSIENT_ATTEMPTS = int(os.environ.get("MAX_TRANSIENT_ATTEMPTS", "12"))
|
|
2825
|
+
_TRANSIENT_MARKERS = (
|
|
2826
|
+
"ReadTimeout",
|
|
2827
|
+
"TimeoutException",
|
|
2828
|
+
"ConnectTimeout",
|
|
2829
|
+
"ConnectError",
|
|
2830
|
+
"PoolTimeout",
|
|
2831
|
+
"502",
|
|
2832
|
+
"503",
|
|
2833
|
+
"504",
|
|
2834
|
+
"Bad Gateway",
|
|
2835
|
+
"Gateway Time-out",
|
|
2836
|
+
"Service Unavailable",
|
|
2837
|
+
"Temporarily Unavailable",
|
|
2838
|
+
)
|
|
2839
|
+
|
|
2840
|
+
|
|
2841
|
+
def _is_transient(err: str) -> bool:
|
|
2842
|
+
"""True if `err` looks like a recoverable infra hiccup (teacher timeout /
|
|
2843
|
+
LB gateway error) rather than a genuine extraction/DB failure."""
|
|
2844
|
+
return any(m in err for m in _TRANSIENT_MARKERS)
|
|
2845
|
+
|
|
2846
|
+
|
|
2847
|
+
def requeue_transient(
|
|
2848
|
+
conn: psycopg.Connection, queue_id: int, transient_attempts: int, error: str
|
|
2849
|
+
) -> None:
|
|
2850
|
+
"""Recoverable INFRA failure. Release the claim back to pending WITHOUT
|
|
2851
|
+
touching `attempts` (preserves the genuine-failure budget), counting instead
|
|
2852
|
+
against `transient_attempts`. Once that budget is exhausted the row IS marked
|
|
2853
|
+
failed — a chunk that times out this many times is genuinely unprocessable,
|
|
2854
|
+
not a passing hiccup."""
|
|
2855
|
+
if transient_attempts + 1 >= MAX_TRANSIENT_ATTEMPTS:
|
|
2856
|
+
mark_failed(conn, queue_id, f"transient-exhausted: {error}")
|
|
2857
|
+
return
|
|
2858
|
+
# status: CASE-guarded rather than a bare 'pending' — same reasoning as
|
|
2859
|
+
# release_claim above. _run_teacher calls requeue_transient on ReadTimeout/
|
|
2860
|
+
# 502/503/504 for BOTH student-primary and teacher-mop-up items; a row
|
|
2861
|
+
# claimed out of `escalated_pending` still carries `escalation_reason` and
|
|
2862
|
+
# must return there (only claim_escalated_batch may reclaim it), not to
|
|
2863
|
+
# bare 'pending' where the student's claim_next_batch would grab and write
|
|
2864
|
+
# it, silently defeating the escalation. In SYNC mode escalation_reason is
|
|
2865
|
+
# never set, so the CASE always yields 'pending' — identical to before.
|
|
2866
|
+
with conn.cursor() as cur:
|
|
2867
|
+
cur.execute(
|
|
2868
|
+
"""
|
|
2869
|
+
UPDATE distillation_queue SET
|
|
2870
|
+
status = CASE WHEN escalation_reason IS NOT NULL
|
|
2871
|
+
THEN 'escalated_pending' ELSE 'pending' END,
|
|
2872
|
+
claimed_by = NULL,
|
|
2873
|
+
claimed_at = NULL,
|
|
2874
|
+
claim_expires_at = NULL,
|
|
2875
|
+
transient_attempts = transient_attempts + 1,
|
|
2876
|
+
last_error = %s
|
|
2877
|
+
WHERE id = %s
|
|
2878
|
+
""",
|
|
2879
|
+
(error[:1024], queue_id),
|
|
2880
|
+
)
|
|
2881
|
+
|
|
2882
|
+
|
|
2716
2883
|
# --------------------------------------------------------------------
|
|
2717
2884
|
# Main loop
|
|
2718
2885
|
# --------------------------------------------------------------------
|
|
@@ -2869,6 +3036,11 @@ async def _process_cascade(
|
|
|
2869
3036
|
f"cascade: {len(callable_items) - len(escalate_items)} student-handled, "
|
|
2870
3037
|
f"{len(escalate_items)} escalated to teacher"
|
|
2871
3038
|
)
|
|
3039
|
+
# RFC-async-escalation.md: flag-off ⇒ passthrough (identity-preserved), so
|
|
3040
|
+
# the inline teacher call below runs exactly as before. Flag-on ⇒ each
|
|
3041
|
+
# escalated item is marked 'escalated_pending' here and this returns [],
|
|
3042
|
+
# so the inline teacher below is never called for this claim.
|
|
3043
|
+
escalate_items = _route_escalations(conn, escalate_items, reason_by_qid)
|
|
2872
3044
|
if escalate_items:
|
|
2873
3045
|
await _run_teacher(http, conn, escalate_items, events_by_qid, False, reason_by_qid)
|
|
2874
3046
|
|
|
@@ -2924,7 +3096,11 @@ async def _run_teacher(
|
|
|
2924
3096
|
log.warning(
|
|
2925
3097
|
f"extraction failed queue_id={queue_id} attempts={attempts}: {err}"
|
|
2926
3098
|
)
|
|
2927
|
-
if
|
|
3099
|
+
if _is_transient(err):
|
|
3100
|
+
requeue_transient(
|
|
3101
|
+
conn, queue_id, item.get("transient_attempts", 0), err
|
|
3102
|
+
)
|
|
3103
|
+
elif _attempts_exhausted(attempts):
|
|
2928
3104
|
mark_failed(conn, queue_id, err)
|
|
2929
3105
|
else:
|
|
2930
3106
|
release_claim(conn, queue_id, err)
|
|
@@ -3022,8 +3198,13 @@ def _apply_extraction(
|
|
|
3022
3198
|
f"db upsert failed queue_id={queue_id} attempts={attempts}: {err}"
|
|
3023
3199
|
)
|
|
3024
3200
|
# The transaction rolled back; the claim bookkeeping runs on the
|
|
3025
|
-
# now-autocommit connection so it persists independently.
|
|
3026
|
-
|
|
3201
|
+
# now-autocommit connection so it persists independently. A transient DB
|
|
3202
|
+
# blip (connection/pool timeout) requeues without burning the genuine
|
|
3203
|
+
# budget; real errors (UniqueViolation, schema mismatch) don't match
|
|
3204
|
+
# _is_transient and take the normal attempts path.
|
|
3205
|
+
if _is_transient(err):
|
|
3206
|
+
requeue_transient(conn, queue_id, item.get("transient_attempts", 0), err)
|
|
3207
|
+
elif _attempts_exhausted(attempts):
|
|
3027
3208
|
mark_failed(conn, queue_id, err)
|
|
3028
3209
|
else:
|
|
3029
3210
|
release_claim(conn, queue_id, err)
|
|
@@ -3135,18 +3316,68 @@ async def amain():
|
|
|
3135
3316
|
await process_batch(http, conn, items, stub_mode)
|
|
3136
3317
|
return len(items)
|
|
3137
3318
|
|
|
3319
|
+
async def _run_one_mopup_batch() -> int:
|
|
3320
|
+
"""Teacher-mop-up: claim + process one batch of `escalated_pending`
|
|
3321
|
+
rows on a DEDICATED connection (mirrors `_run_one_batch`). Only
|
|
3322
|
+
scheduled when ASYNC_ESCALATION_ENABLED. Reuses `_run_teacher` —
|
|
3323
|
+
the SAME teacher call → `_apply_extraction(producer="teacher")` →
|
|
3324
|
+
trace → `mark_done` path the inline cascade escalation uses — so
|
|
3325
|
+
there is exactly one teacher-write code path, invoked either
|
|
3326
|
+
inline (sync escalation) or here (async drain). `reason_by_qid`
|
|
3327
|
+
is populated from each row's `escalation_reason` (set by
|
|
3328
|
+
`mark_escalated_pending` when the student deferred it) so the
|
|
3329
|
+
cascade ledger still tags these as escalated."""
|
|
3330
|
+
with psycopg.connect(PG_DSN, autocommit=True) as conn:
|
|
3331
|
+
esc = claim_escalated_batch(
|
|
3332
|
+
conn, WORKER_ID, BATCH_SIZE, CLAIM_TTL_SEC, MAX_ATTEMPTS
|
|
3333
|
+
)
|
|
3334
|
+
if not esc:
|
|
3335
|
+
return 0
|
|
3336
|
+
log.info(f"mop-up claimed {len(esc)} escalated item(s)")
|
|
3337
|
+
events_by_qid: dict[int, dict[str, Any] | None] = {
|
|
3338
|
+
item["id"]: fetch_event(conn, item["event_id"]) for item in esc
|
|
3339
|
+
}
|
|
3340
|
+
reason_by_qid: dict[int, str] = {
|
|
3341
|
+
item["id"]: item.get("escalation_reason") or "escalated"
|
|
3342
|
+
for item in esc
|
|
3343
|
+
}
|
|
3344
|
+
# Mirror process_batch's missing-event guard — never hand
|
|
3345
|
+
# _run_teacher an id whose event vanished.
|
|
3346
|
+
callable_items = []
|
|
3347
|
+
for item in esc:
|
|
3348
|
+
if events_by_qid[item["id"]] is None:
|
|
3349
|
+
log.warning(
|
|
3350
|
+
f"event {item['event_id']} missing — "
|
|
3351
|
+
f"marking queue {item['id']} done"
|
|
3352
|
+
)
|
|
3353
|
+
mark_done(conn, item["id"])
|
|
3354
|
+
else:
|
|
3355
|
+
callable_items.append(item)
|
|
3356
|
+
if callable_items:
|
|
3357
|
+
await _run_teacher(
|
|
3358
|
+
http, conn, callable_items, events_by_qid, stub_mode, reason_by_qid
|
|
3359
|
+
)
|
|
3360
|
+
return len(esc)
|
|
3361
|
+
|
|
3138
3362
|
# Keep up to MAX_INFLIGHT_BATCHES independent claim→apply pipelines
|
|
3139
3363
|
# running at once. Their student/teacher calls all draw from the shared
|
|
3140
3364
|
# pools above, so while one batch drains its slow tail the others keep
|
|
3141
3365
|
# the GPUs fed. Total GPU concurrency is capped by the sems, NOT by this
|
|
3142
|
-
# count.
|
|
3366
|
+
# count. The teacher-mop-up drain (escalated_pending) is scheduled
|
|
3367
|
+
# OUTSIDE this budget — one at a time, tracked separately — since it is
|
|
3368
|
+
# the teacher tier, not the student fair-share pool; it only runs when
|
|
3369
|
+
# ASYNC_ESCALATION_ENABLED.
|
|
3143
3370
|
inflight: set[asyncio.Task] = set()
|
|
3371
|
+
mopup_task: asyncio.Task | None = None
|
|
3144
3372
|
while True:
|
|
3145
3373
|
try:
|
|
3146
3374
|
while len(inflight) < MAX_INFLIGHT_BATCHES:
|
|
3147
3375
|
inflight.add(asyncio.create_task(_run_one_batch()))
|
|
3148
|
-
|
|
3149
|
-
|
|
3376
|
+
if ASYNC_ESCALATION_ENABLED and mopup_task is None:
|
|
3377
|
+
mopup_task = asyncio.create_task(_run_one_mopup_batch())
|
|
3378
|
+
wait_set = inflight | ({mopup_task} if mopup_task is not None else set())
|
|
3379
|
+
done, pending = await asyncio.wait(
|
|
3380
|
+
wait_set, return_when=asyncio.FIRST_COMPLETED
|
|
3150
3381
|
)
|
|
3151
3382
|
claimed_any = False
|
|
3152
3383
|
for t in done:
|
|
@@ -3154,7 +3385,11 @@ async def amain():
|
|
|
3154
3385
|
if t.result() > 0:
|
|
3155
3386
|
claimed_any = True
|
|
3156
3387
|
except Exception as exc:
|
|
3157
|
-
|
|
3388
|
+
label = "mop-up batch" if t is mopup_task else "batch"
|
|
3389
|
+
log.error(f"{label} task error: {exc}")
|
|
3390
|
+
inflight = pending & inflight
|
|
3391
|
+
if mopup_task is not None and mopup_task not in pending:
|
|
3392
|
+
mopup_task = None
|
|
3158
3393
|
# All finished batches were empty claims ⇒ queue drained; back
|
|
3159
3394
|
# off before re-claiming so idle workers don't hot-spin. Under a
|
|
3160
3395
|
# real backlog every claim returns items and this never sleeps.
|
package/packages/memory-engine-v2/org-model/migrations/013_distillation_transient_attempts.sql
ADDED
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
-- 013 — separate the TRANSIENT retry budget from the genuine-failure budget.
|
|
2
|
+
--
|
|
3
|
+
-- The distiller shares one `attempts` counter for every failure, capped at
|
|
4
|
+
-- MAX_ATTEMPTS (3) and enforced by the claim query's `attempts < MAX_ATTEMPTS`
|
|
5
|
+
-- filter. That conflates two very different failures:
|
|
6
|
+
-- * genuine — parse error, NULL-subject fact, schema mismatch: real, should
|
|
7
|
+
-- give up after a few tries.
|
|
8
|
+
-- * transient — teacher ReadTimeout under load, teacher-lb 502/503/504 during
|
|
9
|
+
-- a scale-from-zero window: infra hiccups that resolve on their
|
|
10
|
+
-- own once the fleet catches up / a box finishes booting.
|
|
11
|
+
--
|
|
12
|
+
-- Because they shared a counter, a transient hiccup burned the genuine budget:
|
|
13
|
+
-- 3 quick ReadTimeouts → `failed`, and since failed rows are never requeued the
|
|
14
|
+
-- extraction was silently LOST (measured 154 ReadTimeout losses in a single day
|
|
15
|
+
-- during a bulk-import whale, 2026-07-14).
|
|
16
|
+
--
|
|
17
|
+
-- This adds a dedicated `transient_attempts` counter. worker.py routes transient
|
|
18
|
+
-- errors here (incrementing THIS, leaving `attempts` untouched so the row stays
|
|
19
|
+
-- claimable — `attempts` remains < MAX_ATTEMPTS) up to MAX_TRANSIENT_ATTEMPTS,
|
|
20
|
+
-- after which the row is marked failed like any genuinely stuck work. Genuine
|
|
21
|
+
-- errors keep using `attempts` exactly as before.
|
|
22
|
+
--
|
|
23
|
+
-- Metadata-only on PostgreSQL (constant DEFAULT, no table rewrite) — safe to run
|
|
24
|
+
-- online against the ~250k-row prod table. Idempotent.
|
|
25
|
+
|
|
26
|
+
ALTER TABLE distillation_queue
|
|
27
|
+
ADD COLUMN IF NOT EXISTS transient_attempts integer NOT NULL DEFAULT 0;
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
-- 014 — escalated_pending state for async escalation (RFC-async-escalation.md).
|
|
2
|
+
--
|
|
3
|
+
-- The student-primary consumer marks escalated events 'escalated_pending' (no
|
|
4
|
+
-- write) instead of calling the teacher inline; the teacher-mop-up consumer
|
|
5
|
+
-- claims and writes them once. Additive + safe online; apply BEFORE the code.
|
|
6
|
+
|
|
7
|
+
ALTER TABLE distillation_queue
|
|
8
|
+
ADD COLUMN IF NOT EXISTS escalation_reason text;
|
|
9
|
+
|
|
10
|
+
-- Recreate the status CHECK to admit the new value. NOT VALID + VALIDATE avoids
|
|
11
|
+
-- the long ACCESS EXCLUSIVE validation lock on the large prod table.
|
|
12
|
+
ALTER TABLE distillation_queue DROP CONSTRAINT IF EXISTS valid_status;
|
|
13
|
+
ALTER TABLE distillation_queue ADD CONSTRAINT valid_status
|
|
14
|
+
CHECK (status = ANY (ARRAY['pending','claimed','done','failed','escalated_pending']))
|
|
15
|
+
NOT VALID;
|
|
16
|
+
ALTER TABLE distillation_queue VALIDATE CONSTRAINT valid_status;
|
|
17
|
+
|
|
18
|
+
-- Mop-up claim + autoscaler count. Mirrors idx_distillation_fair_claim shape.
|
|
19
|
+
-- NOTE: CREATE INDEX CONCURRENTLY cannot run inside a transaction block.
|
|
20
|
+
-- The schema manager applies this as a standalone statement.
|
|
21
|
+
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_distillation_escalated
|
|
22
|
+
ON distillation_queue (priority, emit_ts, id) WHERE status = 'escalated_pending';
|
|
@@ -49,6 +49,21 @@ from datetime import datetime, timezone
|
|
|
49
49
|
|
|
50
50
|
CLEAN_PROMPT_HASH = "6ccfe70f1286a131" # SDK 0.10.19 (#126 + #129); verify vs worker.SYSTEM_PROMPT_HASH
|
|
51
51
|
|
|
52
|
+
# ASYNC_ESCALATION_ENABLED env-string check identical to extractor-async/worker.py
|
|
53
|
+
# (RFC-async-escalation.md Task 6). When on, redistill's re-enqueue must go
|
|
54
|
+
# straight to the teacher (escalated_pending) — the student-primary consumer
|
|
55
|
+
# claims plain 'pending' rows and would otherwise hit these first, which is
|
|
56
|
+
# the opposite of what a stale/dirty-hash re-distill wants.
|
|
57
|
+
ASYNC_ESCALATION_ENABLED = os.environ.get("ASYNC_ESCALATION_ENABLED", "") not in ("", "0", "false")
|
|
58
|
+
|
|
59
|
+
ENQUEUE_SQL = (
|
|
60
|
+
"INSERT INTO distillation_queue (event_id, status, attempts, priority, escalation_reason) "
|
|
61
|
+
"VALUES (%s, 'escalated_pending', 0, 200, 'redistill')"
|
|
62
|
+
if ASYNC_ESCALATION_ENABLED
|
|
63
|
+
else "INSERT INTO distillation_queue (event_id, status, attempts, priority) "
|
|
64
|
+
"VALUES (%s, 'pending', 0, 200)"
|
|
65
|
+
)
|
|
66
|
+
|
|
52
67
|
|
|
53
68
|
def _connect(dsn: str):
|
|
54
69
|
import psycopg
|
|
@@ -97,7 +112,7 @@ def triage(cur, arena: str, clean_hash: str, dirty_hashes: list[str] | None, lim
|
|
|
97
112
|
WHERE t.event_id = e.id AND t.system_prompt_hash = %s)
|
|
98
113
|
AND NOT EXISTS ( -- not already queued (so loops/batches advance)
|
|
99
114
|
SELECT 1 FROM distillation_queue q
|
|
100
|
-
WHERE q.event_id = e.id AND q.status IN ('pending','claimed'))
|
|
115
|
+
WHERE q.event_id = e.id AND q.status IN ('pending','claimed','escalated_pending'))
|
|
101
116
|
AND (
|
|
102
117
|
EXISTS (SELECT 1 FROM distillation_traces t2 -- dirty teacher trace
|
|
103
118
|
WHERE t2.event_id = e.id {dirty_trace})
|
|
@@ -120,7 +135,7 @@ def triage(cur, arena: str, clean_hash: str, dirty_hashes: list[str] | None, lim
|
|
|
120
135
|
AND cardinality(f.provenance_event_ids) = 1) AS solo_fact_n,
|
|
121
136
|
EXISTS (SELECT 1 FROM distillation_queue q
|
|
122
137
|
WHERE q.event_id = s.event_id
|
|
123
|
-
AND q.status IN ('pending','claimed')) AS in_flight
|
|
138
|
+
AND q.status IN ('pending','claimed','escalated_pending')) AS in_flight
|
|
124
139
|
FROM stale s
|
|
125
140
|
ORDER BY fact_n DESC
|
|
126
141
|
""",
|
|
@@ -211,13 +226,14 @@ def main() -> int:
|
|
|
211
226
|
eid = c["event_id"]
|
|
212
227
|
if args.supersede_facts:
|
|
213
228
|
deleted += dump_and_delete_solo_facts(cur, eid, rb)
|
|
214
|
-
|
|
215
|
-
|
|
216
|
-
|
|
217
|
-
|
|
218
|
-
|
|
219
|
-
|
|
220
|
-
|
|
229
|
+
# priority=200 (low) in both modes: re-distills are
|
|
230
|
+
# maintenance — they must never compete with live client
|
|
231
|
+
# traffic. Explicit value wins over the enrichment
|
|
232
|
+
# trigger's mapping (migration 012). Under
|
|
233
|
+
# ASYNC_ESCALATION_ENABLED, ENQUEUE_SQL routes straight to
|
|
234
|
+
# 'escalated_pending' (teacher-only, bypasses the
|
|
235
|
+
# student) — see module docstring + Task 6.
|
|
236
|
+
cur.execute(ENQUEUE_SQL, (eid,))
|
|
221
237
|
enq += 1
|
|
222
238
|
led.write(json.dumps({
|
|
223
239
|
"event_id": eid, "arena": args.arena,
|
|
@@ -0,0 +1,81 @@
|
|
|
1
|
+
"""Guards redistill's enqueue against the latent "hits student first" gap
|
|
2
|
+
(RFC-async-escalation.md Task 6): under async escalation, re-enqueued stale
|
|
3
|
+
events must go straight to the teacher (`escalated_pending`), not `pending`
|
|
4
|
+
(which the student-primary consumer would claim first). Pure-logic tests —
|
|
5
|
+
no DB — importing `redistill` must not require psycopg (it's imported lazily
|
|
6
|
+
inside `_connect`)."""
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import importlib
|
|
10
|
+
import os
|
|
11
|
+
import sys
|
|
12
|
+
|
|
13
|
+
HERE = os.path.dirname(__file__)
|
|
14
|
+
|
|
15
|
+
|
|
16
|
+
def _load():
|
|
17
|
+
sys.path.insert(0, HERE)
|
|
18
|
+
sys.modules.pop("redistill", None)
|
|
19
|
+
return importlib.import_module("redistill")
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
def test_redistill_enqueues_teacher_only(monkeypatch) -> None:
|
|
23
|
+
monkeypatch.setenv("ASYNC_ESCALATION_ENABLED", "1")
|
|
24
|
+
redistill = _load()
|
|
25
|
+
assert "escalated_pending" in redistill.ENQUEUE_SQL
|
|
26
|
+
assert "'redistill'" in redistill.ENQUEUE_SQL
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
def test_redistill_enqueues_pending_when_flag_off(monkeypatch) -> None:
|
|
30
|
+
monkeypatch.delenv("ASYNC_ESCALATION_ENABLED", raising=False)
|
|
31
|
+
redistill = _load()
|
|
32
|
+
assert "'pending'" in redistill.ENQUEUE_SQL
|
|
33
|
+
assert "escalated_pending" not in redistill.ENQUEUE_SQL
|
|
34
|
+
assert "escalation_reason" not in redistill.ENQUEUE_SQL
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def test_redistill_flag_off_explicitly(monkeypatch) -> None:
|
|
38
|
+
for off_value in ("0", "false"):
|
|
39
|
+
monkeypatch.setenv("ASYNC_ESCALATION_ENABLED", off_value)
|
|
40
|
+
redistill = _load()
|
|
41
|
+
assert "'pending'" in redistill.ENQUEUE_SQL
|
|
42
|
+
assert "escalated_pending" not in redistill.ENQUEUE_SQL
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
class _FakeCursor:
|
|
46
|
+
"""Captures the SQL text passed to execute() without touching a DB."""
|
|
47
|
+
|
|
48
|
+
def __init__(self) -> None:
|
|
49
|
+
self.query: str | None = None
|
|
50
|
+
self.params: list | None = None
|
|
51
|
+
|
|
52
|
+
def execute(self, query, params=None) -> None:
|
|
53
|
+
self.query = query
|
|
54
|
+
self.params = params
|
|
55
|
+
|
|
56
|
+
def fetchall(self):
|
|
57
|
+
return []
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def test_triage_not_already_queued_filter_includes_escalated_pending(monkeypatch) -> None:
|
|
61
|
+
"""The CTE's `NOT EXISTS (... status IN (...))` filter must also treat
|
|
62
|
+
`escalated_pending` as in-flight, or a redistill re-run before the teacher
|
|
63
|
+
mop-up drains the queue would re-select and duplicate-enqueue events
|
|
64
|
+
already sitting in escalated_pending (violates exactly-once)."""
|
|
65
|
+
monkeypatch.delenv("ASYNC_ESCALATION_ENABLED", raising=False)
|
|
66
|
+
redistill = _load()
|
|
67
|
+
cur = _FakeCursor()
|
|
68
|
+
redistill.triage(cur, "arena%", redistill.CLEAN_PROMPT_HASH, None, 50)
|
|
69
|
+
assert "AND q.status IN ('pending','claimed','escalated_pending')" in cur.query
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
def test_triage_in_flight_column_includes_escalated_pending(monkeypatch) -> None:
|
|
73
|
+
"""The `in_flight` column subquery must widen the same way, so an event
|
|
74
|
+
already in escalated_pending is correctly reported as in-flight (not
|
|
75
|
+
re-enqueued) rather than silently duplicated."""
|
|
76
|
+
monkeypatch.delenv("ASYNC_ESCALATION_ENABLED", raising=False)
|
|
77
|
+
redistill = _load()
|
|
78
|
+
cur = _FakeCursor()
|
|
79
|
+
redistill.triage(cur, "arena%", redistill.CLEAN_PROMPT_HASH, None, 50)
|
|
80
|
+
in_flight_clause = "AND q.status IN ('pending','claimed','escalated_pending')) AS in_flight"
|
|
81
|
+
assert in_flight_clause in cur.query
|