diff --git a/tests/accuracy/sdk/accuracy-result-storage/disk-storage.ts b/tests/accuracy/sdk/accuracy-result-storage/disk-storage.ts index 06a5480d6..bc09da9e8 100644 --- a/tests/accuracy/sdk/accuracy-result-storage/disk-storage.ts +++ b/tests/accuracy/sdk/accuracy-result-storage/disk-storage.ts @@ -58,9 +58,7 @@ export class DiskBasedResultStorage implements AccuracyResultStorage { async updateRunStatus(commitSHA: string, runId: string, status: AccuracyRunStatuses): Promise { const resultFilePath = this.getAccuracyResultFilePath(commitSHA, runId); - let releaseLock: (() => Promise) | undefined; - try { - releaseLock = await lock(resultFilePath, { retries: 10 }); + await this.withFileLock(resultFilePath, async () => { const accuracyResult = await this.getAccuracyResult(commitSHA, runId, false); if (!accuracyResult) { throw new Error("Results not found!"); @@ -78,23 +76,16 @@ export class DiskBasedResultStorage implements AccuracyResultStorage { ), { encoding: "utf8" } ); - } catch (error) { - console.warn( - `Could not update run status to ${status} for commit - ${commitSHA}, runId - ${runId}.`, - error - ); - throw error; - } finally { - await releaseLock?.(); - } + }); // This bit is important to mark the current run as the latest run for a // commit so that we can use that during baseline comparison. if (status === AccuracyRunStatus.Done) { - await this.atomicUpdateLink( - this.getAccuracyResultFilePath(commitSHA, runId), - this.getLatestResultFilePath(commitSHA) - ); + const latestResultFilePath = this.getLatestResultFilePath(commitSHA); + await this.withFileLock(latestResultFilePath, async () => { + await fs.unlink(latestResultFilePath); + await fs.link(resultFilePath, latestResultFilePath); + }); } } @@ -134,50 +125,36 @@ export class DiskBasedResultStorage implements AccuracyResultStorage { return; } - let releaseLock: (() => Promise) | undefined; - try { - releaseLock = await lock(resultFilePath, { retries: 10 }); - const accuracyResult = await this.getAccuracyResult(commitSHA, runId, false); + await this.withFileLock(resultFilePath, async () => { + let accuracyResult = await this.getAccuracyResult(commitSHA, runId, false); if (!accuracyResult) { throw new Error("Expected at-least initial accuracy result to be present"); } const existingPromptIdx = accuracyResult.promptResults.findIndex((result) => result.prompt === prompt); const promptResult = accuracyResult.promptResults[existingPromptIdx]; - if (!promptResult) { - return await fs.writeFile( - resultFilePath, - JSON.stringify( + if (promptResult) { + accuracyResult.promptResults.splice(existingPromptIdx, 1, { + prompt: promptResult.prompt, + expectedToolCalls: promptResult.expectedToolCalls, + modelResponses: [...promptResult.modelResponses, modelResponse], + }); + } else { + accuracyResult = { + ...accuracyResult, + promptResults: [ + ...accuracyResult.promptResults, { - ...accuracyResult, - promptResults: [ - ...accuracyResult.promptResults, - { - prompt, - expectedToolCalls, - modelResponses: [modelResponse], - }, - ], + prompt, + expectedToolCalls, + modelResponses: [modelResponse], }, - null, - 2 - ) - ); + ], + }; } - accuracyResult.promptResults.splice(existingPromptIdx, 1, { - prompt: promptResult.prompt, - expectedToolCalls: promptResult.expectedToolCalls, - modelResponses: [...promptResult.modelResponses, modelResponse], - }); - - return await fs.writeFile(resultFilePath, JSON.stringify(accuracyResult, null, 2)); - } catch (error) { - console.warn(`Could not save model response for commit - ${commitSHA}, runId - ${runId}.`, error); - throw error; - } finally { - await releaseLock?.(); - } + await fs.writeFile(resultFilePath, JSON.stringify(accuracyResult, null, 2)); + }); } close(): Promise { @@ -206,20 +183,16 @@ export class DiskBasedResultStorage implements AccuracyResultStorage { } } - private async atomicUpdateLink(filePath: string, linkPath: string) { - for (let attempt = 0; attempt < 10; attempt++) { - try { - const tempLinkPath = `${linkPath}~${Date.now()}`; - await fs.link(filePath, tempLinkPath); - await fs.rename(tempLinkPath, linkPath); - return; - } catch (error) { - if (attempt < 10) { - await this.waitFor(100 + Math.random() * 200); - } else { - throw error; - } - } + private async withFileLock(filePath: string, callback: () => Promise): Promise { + let releaseLock: (() => Promise) | undefined; + try { + releaseLock = await lock(filePath, { retries: 10 }); + await callback(); + } catch (error) { + console.warn(`Could not acquire lock for file - ${filePath}.`, error); + throw error; + } finally { + await releaseLock?.(); } } @@ -230,8 +203,4 @@ export class DiskBasedResultStorage implements AccuracyResultStorage { private getLatestResultFilePath(commitSHA: string): string { return path.join(ACCURACY_RESULTS_DIR, commitSHA, `${LATEST_ACCURACY_RUN_NAME}.json`); } - - private waitFor(ms: number) { - return new Promise((resolve) => setTimeout(resolve, ms)); - } }