diff --git a/open-sse/services/batchProcessor.ts b/open-sse/services/batchProcessor.ts index 5d98fd3c8e..315ba94a26 100644 --- a/open-sse/services/batchProcessor.ts +++ b/open-sse/services/batchProcessor.ts @@ -1,12 +1,12 @@ import { v4 as uuidv4 } from "uuid"; import { getPendingBatches, + getTerminalBatches, updateBatch, getFileContent, createFile, getApiKeyById, getBatch, - listBatches, listFiles, deleteFile, updateFileStatus, @@ -90,7 +90,7 @@ export async function processPendingBatches() { async function cleanupExpiredBatches() { try { const now = Math.floor(Date.now() / 1000); - const batches = listBatches(undefined, 200); // Check last 200 batches + const batches = getTerminalBatches(); const parseWindow = (window: string): number => { if (!window) return 86400; @@ -104,10 +104,9 @@ async function cleanupExpiredBatches() { return 86400; }; + // Delete files for terminal batches that have exceeded their completion window for (const batch of batches) { const windowSeconds = parseWindow(batch.completionWindow); - - // Cleanup completed/failed/cancelled batches after their completion window const completionTime = batch.completedAt || batch.failedAt || batch.cancelledAt || batch.expiredAt; if (completionTime && now - completionTime > windowSeconds) { @@ -115,9 +114,12 @@ async function cleanupExpiredBatches() { if (batch.outputFileId) deleteFile(batch.outputFileId); if (batch.errorFileId) deleteFile(batch.errorFileId); } + } - // Expire pending batches that have exceeded their completion window + // Expire validating batches that have exceeded their completion window + for (const batch of getPendingBatches()) { if (batch.status === "validating") { + const windowSeconds = parseWindow(batch.completionWindow); if (now - batch.createdAt > windowSeconds) { updateBatch(batch.id, { status: "expired", expiredAt: now }); } @@ -125,7 +127,8 @@ async function cleanupExpiredBatches() { } // Cleanup orphan files (batch-purpose files stuck in validating after 48h) - const allFiles = listFiles(); + // Use asc order so oldest files are processed first; use a high limit to avoid missing old orphans. + const allFiles = listFiles({ order: "asc", limit: 100 }); for (const file of allFiles) { if ( file.purpose === "batch" && @@ -162,8 +165,10 @@ async function startBatch(batch: any) { // Fire-and-forget: process items in the background so the poll loop isn't blocked. // isProcessing prevents a second poll tick from overlapping. - // noinspection ES6MissingAwait - processBatchItems(batch, lines); + processBatchItems(batch, lines).catch((err) => { + console.error(`[BATCH] Critical error in processBatchItems for ${batch.id}:`, err); + failBatch(batch.id, String(err)); + }); } catch (err) { console.error(`[BATCH] Error starting batch ${batch.id}:`, err); failBatch(batch.id, err instanceof Error ? err.message : String(err)); @@ -272,21 +277,37 @@ async function processBatchItems(batch: any, lines: string[]) { failedCount++; } - // Periodic progress update - updateBatch(batch.id, { - requestCountsCompleted: completedCount, - requestCountsFailed: failedCount, - model: usedModel, - usage: { - input_tokens: totalInputTokens, - output_tokens: totalOutputTokens, - total_tokens: totalInputTokens + totalOutputTokens, - input_tokens_details: { cached_tokens: 0 }, - output_tokens_details: { reasoning_tokens: totalReasoningTokens }, - }, - }); + // Throttle progress updates to every 50 items to reduce DB contention + if ((completedCount + failedCount) % 50 === 0) { + updateBatch(batch.id, { + requestCountsCompleted: completedCount, + requestCountsFailed: failedCount, + model: usedModel, + usage: { + input_tokens: totalInputTokens, + output_tokens: totalOutputTokens, + total_tokens: totalInputTokens + totalOutputTokens, + input_tokens_details: { cached_tokens: 0 }, + output_tokens_details: { reasoning_tokens: totalReasoningTokens }, + }, + }); + } } + // Final progress update to capture accurate counts before finalization + updateBatch(batch.id, { + requestCountsCompleted: completedCount, + requestCountsFailed: failedCount, + model: usedModel, + usage: { + input_tokens: totalInputTokens, + output_tokens: totalOutputTokens, + total_tokens: totalInputTokens + totalOutputTokens, + input_tokens_details: { cached_tokens: 0 }, + output_tokens_details: { reasoning_tokens: totalReasoningTokens }, + }, + }); + // Finalize await finalizeBatch(batch.id, results, errors); } diff --git a/src/app/api/v1/files/route.ts b/src/app/api/v1/files/route.ts index 4b5aa73355..8fe4f36565 100644 --- a/src/app/api/v1/files/route.ts +++ b/src/app/api/v1/files/route.ts @@ -32,10 +32,18 @@ export async function POST(request: Request) { const bytes = file.size; const filename = file.name; - const arrayBuffer = await file.arrayBuffer(); - const content = Buffer.from(arrayBuffer); const mimeType = file.type; + // Stream the upload into memory in chunks to avoid allocating a large contiguous ArrayBuffer + const reader = file.stream().getReader(); + const chunks: Uint8Array[] = []; + while (true) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(value); + } + const content = Buffer.concat(chunks.map((c) => Buffer.from(c))); + let expiresAt: number | undefined; if (expiresAfterAnchor === "created_at" && expiresAfterSeconds) { const seconds = Number.parseInt(expiresAfterSeconds); diff --git a/src/lib/db/batches.ts b/src/lib/db/batches.ts index 6df156b4c5..039ae0136e 100644 --- a/src/lib/db/batches.ts +++ b/src/lib/db/batches.ts @@ -140,3 +140,11 @@ export function getPendingBatches(): BatchRecord[] { ).all(); return rows.map(row => parseBatchRow(row)); } + +export function getTerminalBatches(): BatchRecord[] { + const db = getDbInstance(); + const rows = db.prepare( + "SELECT * FROM batches WHERE status IN ('completed', 'failed', 'cancelled', 'expired') ORDER BY created_at ASC" + ).all(); + return rows.map(row => parseBatchRow(row)); +} diff --git a/src/lib/localDb.ts b/src/lib/localDb.ts index 08f809a737..d64857f4e3 100755 --- a/src/lib/localDb.ts +++ b/src/lib/localDb.ts @@ -220,6 +220,7 @@ export { updateBatch, listBatches, getPendingBatches, + getTerminalBatches, } from "./db/batches"; export type { FileRecord } from "./db/files"; diff --git a/tests/unit/batch_api.test.ts b/tests/unit/batch_api.test.ts index 7b9920b643..f40fee7b94 100644 --- a/tests/unit/batch_api.test.ts +++ b/tests/unit/batch_api.test.ts @@ -1,6 +1,6 @@ import { test } from "node:test"; import assert from "node:assert"; -import { createFile, createBatch, getBatch, getFileContent, updateBatch, createProviderConnection, createApiKey, getFile, listFiles, formatFileResponse, deleteFile } from "../../src/lib/localDb"; +import { createFile, createBatch, getBatch, getFileContent, updateBatch, createProviderConnection, createApiKey, getFile, listFiles, formatFileResponse, deleteFile, getTerminalBatches } from "../../src/lib/localDb"; import { getDbInstance } from "@/lib/db/core.ts"; import { initBatchProcessor, stopBatchProcessor } from "../../open-sse/services/batchProcessor"; @@ -619,3 +619,55 @@ test("Retrieve file content spec compliance", async () => { const unauthorized = (file.apiKeyId !== null && file.apiKeyId !== otherApiKey.id); assert.ok(unauthorized); }); + +test("getTerminalBatches returns only terminal statuses ordered oldest first", async () => { + const apiKey = await createApiKey("Terminal Batches Test Key", "test-machine"); + + const file = createFile({ + bytes: 10, + filename: "terminal_mock.jsonl", + purpose: "batch", + content: Buffer.from("{}"), + apiKeyId: apiKey.id + }); + + // Create batches in different terminal and non-terminal states + const completedBatch = createBatch({ endpoint: "/v1/chat/completions", completionWindow: "24h", inputFileId: file.id, apiKeyId: apiKey.id }); + updateBatch(completedBatch.id, { status: "completed", completedAt: Math.floor(Date.now() / 1000) }); + + const failedBatch = createBatch({ endpoint: "/v1/chat/completions", completionWindow: "24h", inputFileId: file.id, apiKeyId: apiKey.id }); + updateBatch(failedBatch.id, { status: "failed", failedAt: Math.floor(Date.now() / 1000) }); + + const cancelledBatch = createBatch({ endpoint: "/v1/chat/completions", completionWindow: "24h", inputFileId: file.id, apiKeyId: apiKey.id }); + updateBatch(cancelledBatch.id, { status: "cancelled", cancelledAt: Math.floor(Date.now() / 1000) }); + + const expiredBatch = createBatch({ endpoint: "/v1/chat/completions", completionWindow: "24h", inputFileId: file.id, apiKeyId: apiKey.id }); + updateBatch(expiredBatch.id, { status: "expired", expiredAt: Math.floor(Date.now() / 1000) }); + + // This one should NOT appear in terminal batches + const pendingBatch = createBatch({ endpoint: "/v1/chat/completions", completionWindow: "24h", inputFileId: file.id, apiKeyId: apiKey.id }); + + const terminalIds = new Set([completedBatch.id, failedBatch.id, cancelledBatch.id, expiredBatch.id]); + const terminal = getTerminalBatches(); + + // All returned batches must be terminal + for (const b of terminal) { + assert.ok( + ["completed", "failed", "cancelled", "expired"].includes(b.status), + `Unexpected status: ${b.status}` + ); + } + + // Our four terminal batches must all be present + for (const id of terminalIds) { + assert.ok(terminal.some(b => b.id === id), `Missing terminal batch ${id}`); + } + + // The pending batch must not appear + assert.ok(!terminal.some(b => b.id === pendingBatch.id), "Pending batch should not be in terminal list"); + + // Results must be ordered oldest first (created_at ASC) + for (let i = 1; i < terminal.length; i++) { + assert.ok(terminal[i].createdAt >= terminal[i - 1].createdAt, "Results should be ordered oldest first"); + } +});