diff --git a/src/orchestrator/phases/ingest.test.ts b/src/orchestrator/phases/ingest.test.ts new file mode 100644 index 0000000..b3912bb --- /dev/null +++ b/src/orchestrator/phases/ingest.test.ts @@ -0,0 +1,196 @@ +import { afterEach, describe, expect, test } from "bun:test" +import { mkdtempSync, rmSync } from "fs" +import { tmpdir } from "os" +import { join } from "path" +import { runIngestPhase } from "./ingest" +import { CheckpointManager } from "../checkpoint" +import type { Benchmark } from "../../types/benchmark" +import type { Provider, IngestResult } from "../../types/provider" +import type { RunCheckpoint } from "../../types/checkpoint" +import type { UnifiedQuestion, UnifiedSession } from "../../types/unified" + +const QUESTION_ID = "q1" +const SESSION_IDS = ["s1", "s2", "s3", "s4"] + +const tempDirs: string[] = [] + +function makeManager(): CheckpointManager { + const dir = mkdtempSync(join(tmpdir(), "memorybench-ingest-")) + tempDirs.push(dir) + return new CheckpointManager(dir) +} + +/** + * Wait for the queued checkpoint writes to reach disk. save() returns before the + * write happens, so a test that reads the file has to drain the queue first. + */ +async function settleSaves(manager: CheckpointManager, runId: string): Promise { + const locks = (manager as unknown as { saveLock: Map> }).saveLock + for (let i = 0; i < 20; i++) { + const pending = locks.get(runId) + if (!pending) return + await pending.catch(() => {}) + } +} + +function makeQuestion(): UnifiedQuestion { + return { + questionId: QUESTION_ID, + question: "what did I say?", + questionType: "single-session-user", + groundTruth: "something", + haystackSessionIds: [...SESSION_IDS], + } +} + +function makeBenchmark(): Benchmark { + const sessions: UnifiedSession[] = SESSION_IDS.map((sessionId) => ({ + sessionId, + messages: [{ role: "user", content: `message in ${sessionId}` }], + })) + + return { + name: "test", + load: async () => {}, + getQuestions: () => [makeQuestion()], + getHaystackSessions: () => sessions, + getGroundTruth: () => "something", + getQuestionTypes: () => ({}), + } +} + +/** + * Ingests one session at a time, minting an id per session, and throws on the + * session ids listed in `failOn` — the shape of a provider that dies partway + * through a question (rate limit, oversized payload, dropped connection). + */ +function makeProvider(failOn: string[] = []): Provider { + return { + name: "test", + initialize: async () => {}, + ingest: async (sessions: UnifiedSession[]): Promise => { + const ids: string[] = [] + const tasks: string[] = [] + for (const session of sessions) { + if (failOn.includes(session.sessionId)) { + throw new Error(`provider refused ${session.sessionId}`) + } + ids.push(`doc-${session.sessionId}`) + tasks.push(`task-${session.sessionId}`) + } + return { documentIds: ids, taskIds: tasks } + }, + awaitIndexing: async () => {}, + search: async () => [], + clear: async () => {}, + } +} + +function makeCheckpoint(manager: CheckpointManager, runId: string): RunCheckpoint { + const checkpoint = manager.create(runId, "test", "test", "gpt-4o", "gpt-4o") + manager.initQuestion(checkpoint, QUESTION_ID, `${QUESTION_ID}-${runId}`, { + question: "what did I say?", + groundTruth: "something", + questionType: "single-session-user", + }) + return checkpoint +} + +/** Reload the checkpoint from disk, the way resuming a run does. */ +async function reload(manager: CheckpointManager, runId: string): Promise { + await settleSaves(manager, runId) + const loaded = manager.load(runId) + if (!loaded) throw new Error("checkpoint was not written") + return loaded +} + +afterEach(() => { + while (tempDirs.length) { + rmSync(tempDirs.pop()!, { recursive: true, force: true }) + } +}) + +describe("runIngestPhase", () => { + test("keeps the document ids of sessions ingested before a later session failed", async () => { + const manager = makeManager() + const runId = "run-partial" + const checkpoint = makeCheckpoint(manager, runId) + + await expect( + runIngestPhase(makeProvider(["s3"]), makeBenchmark(), checkpoint, manager) + ).rejects.toThrow(/provider refused s3/) + + const persisted = (await reload(manager, runId)).questions[QUESTION_ID].phases.ingest + + // The two ids must be on disk, not just in the failed attempt's memory: the + // sessions that produced them are already recorded as completed, so a resume + // skips them and never mints the ids again. + expect(persisted.completedSessions).toEqual(["s1", "s2"]) + expect(persisted.ingestResult?.documentIds).toEqual(["doc-s1", "doc-s2"]) + expect(persisted.ingestResult?.taskIds).toEqual(["task-s1", "task-s2"]) + }) + + test("resuming after a failure ends up with every session's ids exactly once", async () => { + const manager = makeManager() + const runId = "run-resume" + + await expect( + runIngestPhase(makeProvider(["s3"]), makeBenchmark(), makeCheckpoint(manager, runId), manager) + ).rejects.toThrow() + + const resumed = await reload(manager, runId) + await runIngestPhase(makeProvider(), makeBenchmark(), resumed, manager) + + const persisted = (await reload(manager, runId)).questions[QUESTION_ID].phases.ingest + + expect(persisted.status).toBe("completed") + expect(persisted.completedSessions).toEqual(SESSION_IDS) + // Indexing waits on exactly this list before search is allowed to run, so a + // short list lets search query sessions the provider has not indexed yet. + expect(persisted.ingestResult?.documentIds).toEqual(SESSION_IDS.map((s) => `doc-${s}`)) + expect(persisted.ingestResult?.taskIds).toEqual(SESSION_IDS.map((s) => `task-${s}`)) + }) + + test("does not double-count ids when the phase completes in a single attempt", async () => { + const manager = makeManager() + const runId = "run-clean" + const checkpoint = makeCheckpoint(manager, runId) + + await runIngestPhase(makeProvider(), makeBenchmark(), checkpoint, manager) + + const persisted = (await reload(manager, runId)).questions[QUESTION_ID].phases.ingest + + expect(persisted.ingestResult?.documentIds).toEqual(SESSION_IDS.map((s) => `doc-${s}`)) + }) + + test("re-running a completed question is a no-op rather than a second copy", async () => { + const manager = makeManager() + const runId = "run-rerun" + const checkpoint = makeCheckpoint(manager, runId) + + await runIngestPhase(makeProvider(), makeBenchmark(), checkpoint, manager) + await runIngestPhase(makeProvider(), makeBenchmark(), await reload(manager, runId), manager) + + const persisted = (await reload(manager, runId)).questions[QUESTION_ID].phases.ingest + + expect(persisted.ingestResult?.documentIds).toEqual(SESSION_IDS.map((s) => `doc-${s}`)) + }) + + test("omits taskIds entirely for providers that do not report them", async () => { + const manager = makeManager() + const runId = "run-no-tasks" + const checkpoint = makeCheckpoint(manager, runId) + + const provider = makeProvider() + provider.ingest = async (sessions: UnifiedSession[]) => ({ + documentIds: sessions.map((s) => `doc-${s.sessionId}`), + }) + + await runIngestPhase(provider, makeBenchmark(), checkpoint, manager) + + const persisted = (await reload(manager, runId)).questions[QUESTION_ID].phases.ingest + + expect(persisted.ingestResult?.documentIds).toEqual(SESSION_IDS.map((s) => `doc-${s}`)) + expect(persisted.ingestResult?.taskIds).toBeUndefined() + }) +}) diff --git a/src/orchestrator/phases/ingest.ts b/src/orchestrator/phases/ingest.ts index e38815d..fb64b0f 100644 --- a/src/orchestrator/phases/ingest.ts +++ b/src/orchestrator/phases/ingest.ts @@ -58,9 +58,19 @@ export async function runIngestPhase( }) try { - const completedSessions = - checkpoint.questions[question.questionId].phases.ingest.completedSessions - const combinedResult: IngestResult = { documentIds: [], taskIds: [] } + const ingestPhase = checkpoint.questions[question.questionId].phases.ingest + const completedSessions = ingestPhase.completedSessions + + // Start from what earlier attempts already persisted rather than merging it + // in at the end, so the accumulated result covers every session listed in + // completedSessions at every point in the loop, not just once the loop has + // run to completion. The arrays are copied so the working result is not an + // alias of the checkpoint's until it is written back below. + const previousResult = ingestPhase.ingestResult + const combinedResult: IngestResult = { + documentIds: [...(previousResult?.documentIds || [])], + taskIds: [...(previousResult?.taskIds || [])], + } for (const session of sessions) { if (completedSessions.includes(session.sessionId)) { @@ -74,9 +84,16 @@ export async function runIngestPhase( combinedResult.taskIds!.push(...result.taskIds) } + // Persist the ids together with the session that produced them. A session + // recorded in completedSessions is skipped on resume, so ids left only in + // this in-memory object would be lost for good if a later session threw — + // and indexing would then wait on a short list and let search run before + // those sessions were queryable. save() serialises at call time, so + // handing it the working object records the ids gathered so far. completedSessions.push(session.sessionId) checkpointManager.updatePhase(checkpoint, question.questionId, "ingest", { completedSessions, + ingestResult: combinedResult, }) } @@ -84,20 +101,6 @@ export async function runIngestPhase( delete combinedResult.taskIds } - const existingResult = checkpoint.questions[question.questionId].phases.ingest.ingestResult - if (existingResult) { - combinedResult.documentIds = [ - ...existingResult.documentIds, - ...combinedResult.documentIds, - ] - if (existingResult.taskIds || combinedResult.taskIds) { - combinedResult.taskIds = [ - ...(existingResult.taskIds || []), - ...(combinedResult.taskIds || []), - ] - } - } - const durationMs = Date.now() - startTime checkpointManager.updatePhase(checkpoint, question.questionId, "ingest", { status: "completed",