Something went wrong. Try again.
Japanese learning that pulls JA Bluesky posts, scores them per-token, and serves them thricely.
Something went wrong. Try again.
5.1 kB · 166 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167// Embedding drain loop. Claims rows from CRDB embedding_queue, reads// readings.text, calls embed-stack, writes readings.embedding, deletes// the queue row. Pattern: cidx-cloud/services/embed-runner.
import type { Pool } from "pg"import type { Config } from "./config"
const MAX_BACKOFF_MS = 2 * 60 * 60_000 // 2 hours
export function createDrainFn(config: Config, pool: Pool) { let draining = false let backoffUntil = 0 let consecutiveFailures = 0
return async function drain(): Promise<void> { if (draining) return if (Date.now() < backoffUntil) return draining = true try { await drainLoop(config, pool) consecutiveFailures = 0 backoffUntil = 0 } catch (e) { consecutiveFailures++ const delay = Math.min(1000 * 2 ** consecutiveFailures, MAX_BACKOFF_MS) backoffUntil = Date.now() + delay console.error( `embed drain backing off ${Math.round(delay / 1000)}s after ${consecutiveFailures} consecutive failures` ) } finally { draining = false } }}
async function drainLoop(config: Config, pool: Pool): Promise<void> { while (true) { const client = await pool.connect() try { // Claim a batch using FOR UPDATE SKIP LOCKED. const { rows: batch } = await client.query<{ reading_id: string input_tokens: number }>( `UPDATE embedding_queue SET claimed_by = $1, claimed_at = now() WHERE reading_id IN ( SELECT reading_id FROM embedding_queue WHERE claimed_at IS NULL OR claimed_at < now() - INTERVAL '5 minutes' ORDER BY enqueued_at LIMIT $2 FOR UPDATE SKIP LOCKED ) RETURNING reading_id, input_tokens`, [config.hostname, config.batchSize] )
if (batch.length === 0) break
// Fetch text for embedding. const ids = batch.map((r) => r.reading_id) const placeholders = ids.map((_, i) => `$${i + 1}`).join(",") const { rows: readings } = await client.query<{ id: string text: string }>( `SELECT id, text FROM readings WHERE id IN (${placeholders}) AND embedding IS NULL`, ids )
const textById = new Map(readings.map((r) => [r.id, r.text]))
// Separate: rows with text vs orphans (reading GC'd or already embedded). const toEmbed: { id: string; text: string }[] = [] const skipIds: string[] = []
for (const row of batch) { const text = textById.get(row.reading_id) if (text) { toEmbed.push({ id: row.reading_id, text }) } else { skipIds.push(row.reading_id) } }
// Delete orphan/already-embedded queue rows. if (skipIds.length > 0) { const ph = skipIds.map((_, i) => `$${i + 1}`).join(",") await client.query( `DELETE FROM embedding_queue WHERE reading_id IN (${ph})`, skipIds ) }
if (toEmbed.length === 0) continue
// Call embed-stack. let vectors: number[][] try { const res = await fetch(`${config.embedStackUrl}/embed`, { method: "POST", headers: { "Content-Type": "application/json" }, body: JSON.stringify({ texts: toEmbed.map((e) => e.text), input_type: "document", }), }) if (!res.ok) { throw new Error( `embed-stack ${res.status}: ${(await res.text()).slice(0, 500)}` ) } const data = (await res.json()) as { vectors: number[][] dim: number } vectors = data.vectors } catch (e) { console.error("embed-stack call failed, requeuing batch:", e) // Requeue by clearing claim fields. const ph = toEmbed.map((_, i) => `$${i + 1}`).join(",") await client.query( `UPDATE embedding_queue SET claimed_by = NULL, claimed_at = NULL WHERE reading_id IN (${ph})`, toEmbed.map((e) => e.id) ) throw e // propagate so backoff kicks in }
// Write embeddings and delete queue rows. let embedded = 0 for (let i = 0; i < toEmbed.length; i++) { const vec = vectors[i] if (!vec) continue
await client.query(`UPDATE readings SET embedding = $1 WHERE id = $2`, [ `[${vec.join(",")}]`, toEmbed[i]!.id, ]) await client.query( `DELETE FROM embedding_queue WHERE reading_id = $1`, [toEmbed[i]!.id] ) embedded++ }
console.log(`Embedded ${embedded} readings`) } finally { client.release() } }}
/** Requeue rows this runner had claimed before crashing. */export async function selfHeal(pool: Pool, hostname: string): Promise<number> { const { rowCount } = await pool.query( `UPDATE embedding_queue SET claimed_by = NULL, claimed_at = NULL WHERE claimed_by = $1`, [hostname] ) return rowCount ?? 0}