From 2d901ac7facb1fb3e0603eb957e614dfa375ec9a Mon Sep 17 00:00:00 2001 From: Chris Campbell Date: Thu, 27 Aug 2026 12:20:37 -0700 Subject: [PATCH] fix: correct scheduling of ref data requests when concurrency > 1 --- .../check-core/src/suite/suite-runner.spec.ts | 141 +++++++++++++++++- packages/check-core/src/suite/suite-runner.ts | 77 +++++++--- 2 files changed, 189 insertions(+), 29 deletions(-) diff --git a/packages/check-core/src/suite/suite-runner.spec.ts b/packages/check-core/src/suite/suite-runner.spec.ts index 5ad4d17a..37a07445 100644 --- a/packages/check-core/src/suite/suite-runner.spec.ts +++ b/packages/check-core/src/suite/suite-runner.spec.ts @@ -4,6 +4,7 @@ import { describe, expect, it } from 'vitest' import type { TaskExecutor, TaskExecutorKey } from '../_shared/task-queue' import { createExecutor, TaskQueue } from '../_shared/task-queue' +import type { DatasetKey, DatasetMap } from '../_shared/types' import type { Bundle, BundleModel, ModelSpec } from '../bundle/bundle-types' @@ -21,40 +22,69 @@ import type { RunSuiteCallbacks } from './suite-runner' import { runSuiteWithTaskQueue } from './suite-runner' interface MockConfigOptions { + /** If true, configure the check suite with no tests. */ emptyTests?: boolean + /** If true, configure the check suite with a test that fails to parse. */ invalidTests?: boolean + /** If true, use a check test that includes a predicate that references another dataset. */ + testsWithRefData?: boolean + /** The number of model instances to initialize (defaults to 1). */ + concurrency?: number + /** The delay (in msecs) applied to every model run. */ delayInGetDatasets?: number + /** The delay (in msecs) applied to a model run that includes the given dataset key. */ + delaysByDatasetKey?: Map + /** If true, throw an error from each model run in the "current" bundle. */ throwInCurrentGetDatasets?: boolean + /** If true, configure the suite with comparison tests in addition to check tests. */ includeComparisons?: boolean + /** Called at the start of each model run, before any configured delay is applied. */ onGetDatasets?: () => void + /** Called when a model run starts and when it completes. */ + onModelRun?: (event: string) => void } function mockBundleModel(modelSpec: ModelSpec, mockOptions: MockConfigOptions): BundleModel { return { modelSpec, - getDatasetsForScenario: async () => { + getDatasetsForScenario: async (_scenarioSpec, datasetKeys) => { if (mockOptions.onGetDatasets) { mockOptions.onGetDatasets() } - if (mockOptions.delayInGetDatasets) { - await new Promise(resolve => setTimeout(resolve, mockOptions.delayInGetDatasets)) + mockOptions.onModelRun?.(`start ${datasetKeys.join(',')}`) + + // Use the longest delay that is configured for the requested dataset keys + let delay = mockOptions.delayInGetDatasets || 0 + for (const datasetKey of datasetKeys) { + const delayForKey = mockOptions.delaysByDatasetKey?.get(datasetKey) + if (delayForKey !== undefined && delayForKey > delay) { + delay = delayForKey + } + } + if (delay > 0) { + await new Promise(resolve => setTimeout(resolve, delay)) } + if (mockOptions.throwInCurrentGetDatasets === true) { throw new Error('Fake error') } + + const datasetMap: DatasetMap = new Map(datasetKeys.map(datasetKey => [datasetKey, new Map([[2000, 0]])])) + mockOptions.onModelRun?.(`end ${datasetKeys.join(',')}`) return { - datasetMap: new Map([['Model_v1', new Map([[2000, 0]])]]) + datasetMap } } } } function mockBundle(mockOptions: MockConfigOptions): Bundle { + const outputVars = mockOptions.testsWithRefData === true ? [outputVar('V1'), outputVar('V2')] : [outputVar('V1')] const modelSpec: ModelSpec = { modelSizeInBytes: 0, dataSizeInBytes: 0, inputVars: new Map([inputVar('1', 'I1'), inputVar('2', 'I2')]), - outputVars: new Map([outputVar('V1')]), + outputVars: new Map(outputVars), implVars: new Map(), inputGroups: new Map(), datasetGroups: new Map() @@ -74,6 +104,26 @@ async function mockConfig(mockOptions: MockConfigOptions): Promise { tests = [] } else if (mockOptions.invalidTests === true) { tests = ['INVALID'] + } else if (mockOptions.testsWithRefData === true) { + // This test includes a predicate that compares one dataset against another + // dataset, which means the reference data for V2 must be fetched (and held + // in memory) before the check on V1 is performed + const test = ` +- describe: group1 + tests: + - it: test1 + scenarios: + - with: I1 + at: min + datasets: + - name: V1 + predicates: + - eq: + dataset: + name: V2 + scenario: inherit +` + tests = [test] } else { const test = ` - describe: group1 @@ -147,7 +197,8 @@ async function mockConfig(mockOptions: MockConfigOptions): Promise { check: { tests }, - comparison: comparisonOptions + comparison: comparisonOptions, + concurrency: mockOptions.concurrency } return createConfig(configOptions) } @@ -196,6 +247,84 @@ describe('runSuite', () => { expect(getDatasetsCallCount).toBe(2) }) + it('should fetch all reference data before running checks that depend on it', async () => { + // Use more than one model instance so that multiple data requests can be processed + // concurrently, and make the reference data request (for V2) take longer than the + // check data request (for V1). If the two were allowed to run at the same time, the + // check would be performed before the reference data was available. + const events: string[] = [] + const config = await mockConfig({ + testsWithRefData: true, + concurrency: 2, + delaysByDatasetKey: new Map([['Model_v2', 20]]), + onModelRun: event => { + events.push(event) + } + }) + const taskQueue = mockTaskQueue(config) + + const progressPcts: number[] = [] + const report: SuiteReport = await new Promise((resolve, reject) => { + const callbacks: RunSuiteCallbacks = { + onProgress: pct => { + progressPcts.push(pct) + }, + onComplete: suiteReport => { + resolve(suiteReport) + }, + onError: () => { + reject(new Error('onError should not be called')) + } + } + runSuiteWithTaskQueue(config, taskQueue, callbacks) + }) + + // Verify that the reference data run completed before the check data run started + expect(events).toEqual(['start Model_v2', 'end Model_v2', 'start Model_v1', 'end Model_v1']) + + // Verify that the check passed (it will be reported as an error if the reference + // data was not resolved by the time the check was performed) + const test1 = report.checkReport.groups[0].tests[0] + expect(test1.status).toBe('passed') + + // Verify that progress is reported across both sets of data requests + expect(progressPcts).toEqual([0, 0.5, 1]) + }) + + it('should build reports when there is reference data but all checks are skipped', async () => { + const events: string[] = [] + const config = await mockConfig({ + testsWithRefData: true, + concurrency: 2, + onModelRun: event => { + events.push(event) + } + }) + const taskQueue = mockTaskQueue(config) + + const report: SuiteReport = await new Promise((resolve, reject) => { + const callbacks: RunSuiteCallbacks = { + onComplete: suiteReport => { + resolve(suiteReport) + }, + onError: () => { + reject(new Error('onError should not be called')) + } + } + runSuiteWithTaskQueue(config, taskQueue, callbacks, { + skipChecks: [{ groupName: 'group1', testName: 'test1' }] + }) + }) + + // Verify that only the reference data run was performed (there are no check data + // requests because the only check was skipped) + expect(events).toEqual(['start Model_v2', 'end Model_v2']) + + // Verify that the report was still built + expect(report.checkReport.groups.length).toBe(1) + expect(report.checkReport.groups[0].tests[0].status).toBe('skipped') + }) + it('should notify progress and completion callbacks even when there are no tests', async () => { const config = await mockConfig({ emptyTests: true }) const taskQueue = mockTaskQueue(config) diff --git a/packages/check-core/src/suite/suite-runner.ts b/packages/check-core/src/suite/suite-runner.ts index a30d0aca..26c60640 100644 --- a/packages/check-core/src/suite/suite-runner.ts +++ b/packages/check-core/src/suite/suite-runner.ts @@ -141,8 +141,9 @@ class SuiteRunner { // the reference data is available in memory when checks are performed. const refDataPlan = refDataPlanner.buildPlan() const dataPlan = dataPlanner.buildPlan() - const dataRequests = [...refDataPlan.requests, ...dataPlan.requests] - const taskCount = dataRequests.length + const refDataRequests = refDataPlan.requests + const dataRequests = dataPlan.requests + const taskCount = refDataRequests.length + dataRequests.length if (taskCount === 0) { // There are no data requests. This can occur when there are no checks or comparisons defined, // or when all checks and comparisons are skipped. We still build the reports so that the @@ -183,32 +184,62 @@ class SuiteRunner { }) } - // Schedule a task for each data request + // Schedule a task for each of the given data requests. The `onGroupComplete` + // function will be called after all of the given requests have been processed. let tasksCompleted = 0 let dataTaskId = 1 - for (const dataRequest of dataRequests) { - const task: Task = { - key: `suite-runner-${dataTaskId++}`, - kind: 'suite-runner', - process: async bundleModels => { - // Remove the task key from the set of pending task keys - this.pendingTaskKeys.delete(task.key) - - // Process the request - await this.processRequest(dataRequest, bundleModels) - - // Notify the progress callback after each task is processed - tasksCompleted++ - this.callbacks.onProgress?.(tasksCompleted / taskCount) - - // Notify the completion callback when all tasks have been processed - if (tasksCompleted === taskCount) { - buildReport() + const scheduleRequests = (requests: DataRequest[], onGroupComplete: () => void) => { + let groupTasksCompleted = 0 + for (const dataRequest of requests) { + const task: Task = { + key: `suite-runner-${dataTaskId++}`, + kind: 'suite-runner', + process: async bundleModels => { + // Remove the task key from the set of pending task keys + this.pendingTaskKeys.delete(task.key) + + // Process the request + await this.processRequest(dataRequest, bundleModels) + + // Notify the progress callback after each task is processed + tasksCompleted++ + this.callbacks.onProgress?.(tasksCompleted / taskCount) + + // Notify the callback when all tasks in this group have been processed + groupTasksCompleted++ + if (groupTasksCompleted === requests.length) { + onGroupComplete() + } } } + this.taskQueue.addTask(task) + this.pendingTaskKeys.add(task.key) } - this.taskQueue.addTask(task) - this.pendingTaskKeys.add(task.key) + } + + // Schedule the tasks that fetch the data for the checks and comparisons. Note + // that this is only called after all reference data has been fetched. + const scheduleDataRequests = () => { + if (this.stopped) { + return + } + if (dataRequests.length > 0) { + scheduleRequests(dataRequests, buildReport) + } else { + // There is reference data but no checks or comparisons to run (as is the + // case when all checks are skipped), so build the report now + buildReport() + } + } + + // Schedule the reference data tasks first. These must all be processed before + // any check tasks are scheduled; otherwise, when more than one task can be + // processed at a time, a check could run before the data that it references + // has been fetched. + if (refDataRequests.length > 0) { + scheduleRequests(refDataRequests, scheduleDataRequests) + } else { + scheduleDataRequests() } }