From dfb93eab79011ed5056db4baffe14068f6e54a7a Mon Sep 17 00:00:00 2001 From: Patrick Erichsen Date: Tue, 23 Jun 2026 15:19:20 -0700 Subject: [PATCH] fix: make security dataset export more adaptive Merge PR #2814 after exact-head checks passed. --- .../workflows/security-dataset-snapshot.yml | 33 ++++++++++++-- scripts/security-dataset/export-snapshot.ts | 44 ++++++++++++++++--- .../exportSnapshotCli.test.ts | 23 ++++++++++ 3 files changed, 91 insertions(+), 9 deletions(-) diff --git a/.github/workflows/security-dataset-snapshot.yml b/.github/workflows/security-dataset-snapshot.yml index c486e31c..9e984af0 100644 --- a/.github/workflows/security-dataset-snapshot.yml +++ b/.github/workflows/security-dataset-snapshot.yml @@ -19,6 +19,26 @@ on: description: "Hugging Face branch/revision to upload to" required: true default: "main" + page-size: + description: "Live export page size per Convex request" + required: true + default: "25" + min-page-size: + description: "Smallest page size to retry after Convex timeouts" + required: true + default: "5" + batch-pages: + description: "Live export pages per Convex request" + required: true + default: "1" + concurrency: + description: "Concurrent live export shards" + required: true + default: "2" + shards: + description: "Created-at shards per source kind" + required: true + default: "24" schedule: - cron: "17 9 * * *" @@ -44,6 +64,11 @@ jobs: HF_REVISION: ${{ inputs['hf-revision'] || 'main' }} HF_UPLOAD: ${{ github.event_name == 'schedule' || inputs.upload == 'true' }} SNAPSHOT_LIMIT: ${{ inputs.limit || '' }} + SNAPSHOT_PAGE_SIZE: ${{ inputs['page-size'] || '25' }} + SNAPSHOT_MIN_PAGE_SIZE: ${{ inputs['min-page-size'] || '5' }} + SNAPSHOT_BATCH_PAGES: ${{ inputs['batch-pages'] || '1' }} + SNAPSHOT_CONCURRENCY: ${{ inputs.concurrency || '2' }} + SNAPSHOT_SHARDS: ${{ inputs.shards || '24' }} SANITIZED_OUT_DIR: /tmp/clawhub-security-dataset/sanitized WORK_DIR: /tmp/clawhub-security-dataset steps: @@ -82,9 +107,11 @@ jobs: --hf-dataset --hf-repo "$HF_DATASET_REPO" --hf-revision "$HF_REVISION" - --page-size 25 - --batch-pages 2 - --concurrency 3 + --page-size "$SNAPSHOT_PAGE_SIZE" + --min-page-size "$SNAPSHOT_MIN_PAGE_SIZE" + --batch-pages "$SNAPSHOT_BATCH_PAGES" + --concurrency "$SNAPSHOT_CONCURRENCY" + --shards "$SNAPSHOT_SHARDS" ) if [[ -n "$SNAPSHOT_LIMIT" ]]; then args+=(--limit "$SNAPSHOT_LIMIT") diff --git a/scripts/security-dataset/export-snapshot.ts b/scripts/security-dataset/export-snapshot.ts index c1ed1001..4e307837 100644 --- a/scripts/security-dataset/export-snapshot.ts +++ b/scripts/security-dataset/export-snapshot.ts @@ -68,6 +68,7 @@ type Options = { mode: "public"; limit: number | null; pageSize: number; + minPageSize: number; batchPages: number; concurrency: number; shards: number; @@ -239,9 +240,11 @@ async function exportShard(input: { }) { const { options, shard, state, writers } = input; let cursor: string | null = null; + let pageSize = options.pageSize; let batchPages = options.batchPages; while (!isLimitReached(options, state)) { - const result = await runConvexPage(options, shard, cursor, options.pageSize, batchPages); + const result = await runConvexPage(options, shard, cursor, pageSize, batchPages); + pageSize = result.pageSize; batchPages = result.batchPages; const page = result.page; const inputs = reserveExportInputs(page.page, state, options.limit); @@ -262,9 +265,10 @@ async function runConvexPage( cursor: string | null, numItems: number, batchPages: number, -): Promise<{ page: ConvexPage; batchPages: number }> { +): Promise<{ page: ConvexPage; pageSize: number; batchPages: number }> { const functionName = "securityDatasetNode:listArtifactExportBatchCompressedInternal"; const workerFunctionName = "securityDatasetNode:listArtifactExportBatchCompressed"; + let pageSize = numItems; let pageCount = batchPages; while (true) { @@ -273,7 +277,7 @@ async function runConvexPage( mode: options.mode, createdAtGte: shard.createdAtGte, createdAtLt: shard.createdAtLt, - paginationOpts: { cursor, numItems }, + paginationOpts: { cursor, numItems: pageSize }, pageCount, }; @@ -293,13 +297,21 @@ async function runConvexPage( args, isCompressedConvexPage, ); - return { page: decodeCompressedConvexPage(compressed), batchPages: pageCount }; + return { + page: decodeCompressedConvexPage(compressed), + pageSize, + batchPages: pageCount, + }; } catch (error) { lastError = error; - if (isLikelyOversizedConvexBatch(error) && pageCount > 1) break; + if ( + isLikelyOversizedConvexBatch(error) && + canReduceConvexBatch(options, pageSize, pageCount) + ) + break; if (attempt === DEFAULT_MAX_CONVEX_ATTEMPTS) break; console.error( - `[snapshot] retrying ${functionName} batch-pages=${pageCount} after attempt ${attempt}: ${errorMessage(error)}`, + `[snapshot] retrying ${functionName} page-size=${pageSize} batch-pages=${pageCount} after attempt ${attempt}: ${errorMessage(error)}`, ); await delay(attempt * 500); } @@ -314,6 +326,16 @@ async function runConvexPage( continue; } + if (isLikelyOversizedConvexBatch(lastError) && pageSize > options.minPageSize) { + const nextPageSize = Math.max(options.minPageSize, Math.floor(pageSize / 2)); + console.error( + `[snapshot] ${shard.label} reducing page-size ${pageSize}->${nextPageSize}: ${errorMessage(lastError)}`, + ); + pageSize = nextPageSize; + pageCount = 1; + continue; + } + writeCommandErrorOutput(lastError); throw lastError; } @@ -768,6 +790,10 @@ function isLikelyOversizedConvexBatch(error: unknown) { return isLikelyTruncatedConvexOutput(error) || isLikelyConvexOperationTimeout(error); } +function canReduceConvexBatch(options: Options, pageSize: number, pageCount: number) { + return pageCount > 1 || pageSize > options.minPageSize; +} + function writeCommandErrorOutput(error: unknown) { if (!isRecord(error)) return; if (typeof error.stderr === "string" && error.stderr.length > 0) { @@ -794,6 +820,7 @@ function parseArgs(args: string[]): Options { mode: "public", limit: null, pageSize: DEFAULT_PAGE_SIZE, + minPageSize: Math.min(10, DEFAULT_PAGE_SIZE), batchPages: DEFAULT_BATCH_PAGES, concurrency: DEFAULT_CONCURRENCY, shards: DEFAULT_SHARDS, @@ -825,6 +852,8 @@ function parseArgs(args: string[]): Options { options.limit = readPositiveInt(readValue(args, ++index, arg), arg); } else if (arg === "--page-size") { options.pageSize = readPositiveInt(readValue(args, ++index, arg), arg); + } else if (arg === "--min-page-size") { + options.minPageSize = readPositiveInt(readValue(args, ++index, arg), arg); } else if (arg === "--batch-pages") { options.batchPages = readPositiveInt(readValue(args, ++index, arg), arg); } else if (arg === "--concurrency") { @@ -864,6 +893,9 @@ function parseArgs(args: string[]): Options { if (options.workerToken && (options.prod || options.deployment || options.push)) { throw new Error("Use --worker-token with --convex-url instead of --prod/--deployment/--push."); } + if (options.minPageSize > options.pageSize) { + throw new Error("--min-page-size must be less than or equal to --page-size."); + } assertCreatedTimeWindow(options.timeWindow); return options; } diff --git a/scripts/security-dataset/exportSnapshotCli.test.ts b/scripts/security-dataset/exportSnapshotCli.test.ts index ace7de22..36a447dd 100644 --- a/scripts/security-dataset/exportSnapshotCli.test.ts +++ b/scripts/security-dataset/exportSnapshotCli.test.ts @@ -88,6 +88,29 @@ describe("security dataset snapshot CLI", () => { await rm(directory, { recursive: true, force: true }); } }); + + it("rejects a minimum page size larger than the starting page size", async () => { + await expect( + execFileAsync( + "bun", + [ + "scripts/security-dataset/export-snapshot.ts", + "--page-size", + "5", + "--min-page-size", + "6", + "--dry-run", + ], + { + cwd: process.cwd(), + encoding: "utf8", + maxBuffer: 16 * 1024 * 1024, + }, + ), + ).rejects.toMatchObject({ + stderr: expect.stringContaining("--min-page-size must be less than or equal to --page-size."), + }); + }); }); function buildTinyConvexSnapshotZip() {