From 68632015d495c7414fb3270b03d61790f2e3d5f8 Mon Sep 17 00:00:00 2001 From: Alex Gherghisan Date: Thu, 28 Nov 2024 11:52:46 +0000 Subject: [PATCH] refactor: use common DATA_DIRECTORY env var --- .../src/interfaces/prover-broker.ts | 14 +++++--- .../src/interfaces/prover-client.ts | 22 ++++++++----- yarn-project/foundation/src/config/env_var.ts | 2 -- .../src/prover-client/prover-client.ts | 5 --- .../src/proving_broker/factory.ts | 14 ++++++-- yarn-project/prover-node/src/factory.ts | 4 +-- .../src/prover-cache/cache_manager.ts | 33 +++++++++++-------- 7 files changed, 55 insertions(+), 39 deletions(-) diff --git a/yarn-project/circuit-types/src/interfaces/prover-broker.ts b/yarn-project/circuit-types/src/interfaces/prover-broker.ts index fb2fdadef67..4b55e9124db 100644 --- a/yarn-project/circuit-types/src/interfaces/prover-broker.ts +++ b/yarn-project/circuit-types/src/interfaces/prover-broker.ts @@ -18,7 +18,8 @@ export const ProverBrokerConfig = z.object({ /** If starting a prover broker locally, the interval the broker checks for timed out jobs */ proverBrokerPollIntervalMs: z.number(), /** If starting a prover broker locally, the directory to store broker data */ - proverBrokerDataDirectory: z.string().optional(), + dataDirectory: z.string().optional(), + dataStoreMapSizeKB: z.number(), }); export type ProverBrokerConfig = z.infer; @@ -39,9 +40,14 @@ export const proverBrokerConfigMappings: ConfigMappingsType description: 'If starting a prover broker locally, the max number of retries per proving job', ...numberConfigHelper(3), }, - proverBrokerDataDirectory: { - env: 'PROVER_BROKER_DATA_DIRECTORY', - description: 'If starting a prover broker locally, the directory to store broker data', + dataDirectory: { + env: 'DATA_DIRECTORY', + description: 'Optional dir to store data. If omitted will store in memory.', + }, + dataStoreMapSizeKB: { + env: 'DATA_STORE_MAP_SIZE_KB', + description: 'DB mapping size to be applied to all key/value stores', + ...numberConfigHelper(128 * 1_024 * 1_024), // Defaulted to 128 GB }, }; diff --git a/yarn-project/circuit-types/src/interfaces/prover-client.ts b/yarn-project/circuit-types/src/interfaces/prover-client.ts index 29f8cc4fb53..08a10410257 100644 --- a/yarn-project/circuit-types/src/interfaces/prover-client.ts +++ b/yarn-project/circuit-types/src/interfaces/prover-client.ts @@ -24,9 +24,8 @@ export type ProverConfig = ActualProverConfig & { nodeUrl?: string; /** Identifier of the prover */ proverId: Fr; - /** Where to store temporary data */ - cacheDir?: string; - + dataDirectory?: string; + dataStoreMapSizeKB: number; proverAgentCount: number; }; @@ -35,7 +34,8 @@ export const ProverConfigSchema = z.object({ realProofs: z.boolean(), proverId: schemas.Fr, proverTestDelayMs: z.number(), - cacheDir: z.string().optional(), + dataDirectory: z.string().optional(), + dataStoreMapSizeKB: z.number(), proverAgentCount: z.number(), }) satisfies ZodFor; @@ -60,16 +60,20 @@ export const proverConfigMappings: ConfigMappingsType = { description: 'Artificial delay to introduce to all operations to the test prover.', ...numberConfigHelper(0), }, - cacheDir: { - env: 'PROVER_CACHE_DIR', - description: 'Where to store cache data generated while proving', - defaultValue: '/tmp/aztec-prover', - }, proverAgentCount: { env: 'PROVER_AGENT_COUNT', description: 'The number of prover agents to start', ...numberConfigHelper(1), }, + dataDirectory: { + env: 'DATA_DIRECTORY', + description: 'Optional dir to store data. If omitted will store in memory.', + }, + dataStoreMapSizeKB: { + env: 'DATA_STORE_MAP_SIZE_KB', + description: 'DB mapping size to be applied to all key/value stores', + ...numberConfigHelper(128 * 1_024 * 1_024), // Defaulted to 128 GB + }, }; function parseProverId(str: string) { diff --git a/yarn-project/foundation/src/config/env_var.ts b/yarn-project/foundation/src/config/env_var.ts index cb43781825f..d5135deae6e 100644 --- a/yarn-project/foundation/src/config/env_var.ts +++ b/yarn-project/foundation/src/config/env_var.ts @@ -111,7 +111,6 @@ export type EnvVar = | 'PROVER_BROKER_JOB_TIMEOUT_MS' | 'PROVER_BROKER_POLL_INTERVAL_MS' | 'PROVER_BROKER_JOB_MAX_RETRIES' - | 'PROVER_BROKER_DATA_DIRECTORY' | 'PROVER_COORDINATION_NODE_URL' | 'PROVER_DISABLED' | 'PROVER_ID' @@ -125,7 +124,6 @@ export type EnvVar = | 'PROVER_REAL_PROOFS' | 'PROVER_REQUIRED_CONFIRMATIONS' | 'PROVER_TEST_DELAY_MS' - | 'PROVER_CACHE_DIR' | 'PXE_BLOCK_POLLING_INTERVAL_MS' | 'PXE_L2_STARTING_BLOCK' | 'PXE_PROVER_ENABLED' diff --git a/yarn-project/prover-client/src/prover-client/prover-client.ts b/yarn-project/prover-client/src/prover-client/prover-client.ts index b5d81a34905..daf715ce4e4 100644 --- a/yarn-project/prover-client/src/prover-client/prover-client.ts +++ b/yarn-project/prover-client/src/prover-client/prover-client.ts @@ -16,8 +16,6 @@ import { createLogger } from '@aztec/foundation/log'; import { NativeACVMSimulator } from '@aztec/simulator'; import { type TelemetryClient } from '@aztec/telemetry-client'; -import { join } from 'path'; - import { type ProverClientConfig } from '../config.js'; import { ProvingOrchestrator } from '../orchestrator/orchestrator.js'; import { CachingBrokerFacade } from '../proving_broker/caching_broker_facade.js'; @@ -30,8 +28,6 @@ export class ProverClient implements EpochProverManager { private running = false; private agents: ProvingAgent[] = []; - private cacheDir?: string; - private constructor( private config: ProverClientConfig, private worldState: ForkMerkleTreeOperations, @@ -42,7 +38,6 @@ export class ProverClient implements EpochProverManager { ) { // TODO(palla/prover-node): Cache the paddingTx here, and not in each proving orchestrator, // so it can be reused across multiple ones and not recomputed every time. - this.cacheDir = this.config.cacheDir ? join(this.config.cacheDir, `tx_prover_${this.config.proverId}`) : undefined; } public createEpochProver(cache: ProverCache = new InMemoryProverCache()): EpochProver { diff --git a/yarn-project/prover-client/src/proving_broker/factory.ts b/yarn-project/prover-client/src/proving_broker/factory.ts index 67295fb6011..67069d86dd6 100644 --- a/yarn-project/prover-client/src/proving_broker/factory.ts +++ b/yarn-project/prover-client/src/proving_broker/factory.ts @@ -2,7 +2,10 @@ import { type ProverBrokerConfig } from '@aztec/circuit-types'; import { AztecLmdbStore } from '@aztec/kv-store/lmdb'; import { type TelemetryClient } from '@aztec/telemetry-client'; +import { join } from 'path'; + import { ProvingBroker } from './proving_broker.js'; +import { ProvingBrokerDatabase } from './proving_broker_database.js'; import { InMemoryBrokerDatabase } from './proving_broker_database/memory.js'; import { KVBrokerDatabase } from './proving_broker_database/persisted.js'; @@ -10,9 +13,14 @@ export async function createAndStartProvingBroker( config: ProverBrokerConfig, client: TelemetryClient, ): Promise { - const database = config.proverBrokerDataDirectory - ? new KVBrokerDatabase(AztecLmdbStore.open(config.proverBrokerDataDirectory), client) - : new InMemoryBrokerDatabase(); + let database: ProvingBrokerDatabase; + if (config.dataDirectory) { + const dataDir = join(config.dataDirectory, 'prover_broker'); + const store = AztecLmdbStore.open(dataDir, config.dataStoreMapSizeKB); + database = new KVBrokerDatabase(store, client); + } else { + database = new InMemoryBrokerDatabase(); + } const broker = new ProvingBroker(database, client, { jobTimeoutMs: config.proverBrokerJobTimeoutMs, diff --git a/yarn-project/prover-node/src/factory.ts b/yarn-project/prover-node/src/factory.ts index 57429290be7..2caf2adf91d 100644 --- a/yarn-project/prover-node/src/factory.ts +++ b/yarn-project/prover-node/src/factory.ts @@ -78,8 +78,8 @@ export async function createProverNode( const walletClient = publisher.getClient(); const bondManager = await createBondManager(rollupContract, walletClient, config); - const cacheDir = config.cacheDir ? join(config.cacheDir, `prover_${config.proverId}`) : undefined; - const cacheManager = new ProverCacheManager(cacheDir); + const cacheDir = config.dataDirectory ? join(config.dataDirectory, `prover_${config.proverId}`) : undefined; + const cacheManager = new ProverCacheManager(cacheDir, config.dataStoreMapSizeKB); return new ProverNode( prover, diff --git a/yarn-project/prover-node/src/prover-cache/cache_manager.ts b/yarn-project/prover-node/src/prover-cache/cache_manager.ts index 497300d1e42..8ba134c53a2 100644 --- a/yarn-project/prover-node/src/prover-cache/cache_manager.ts +++ b/yarn-project/prover-node/src/prover-cache/cache_manager.ts @@ -14,26 +14,29 @@ const EPOCH_DIR_SEPARATOR = '_'; const EPOCH_HASH_FILENAME = 'epoch_hash.txt'; export class ProverCacheManager { - constructor(private cacheDir?: string, private log = createLogger('prover-node:cache-manager')) {} + constructor( + private dataRootDir?: string, + private cacheMapSize?: number, + private log = createLogger('aztec:prover-node:cache-manager'), + ) {} public async openCache(epochNumber: bigint, epochHash: Buffer): Promise { - if (!this.cacheDir) { + if (!this.dataRootDir) { return new InMemoryProverCache(); } - const epochDir = EPOCH_DIR_PREFIX + EPOCH_DIR_SEPARATOR + epochNumber; - const dataDir = join(this.cacheDir, epochDir); + const epochDataDir = join(this.dataRootDir, EPOCH_DIR_PREFIX + EPOCH_DIR_SEPARATOR + epochNumber); - const storedEpochHash = await readFile(join(dataDir, EPOCH_HASH_FILENAME), 'hex').catch(() => Buffer.alloc(0)); + const storedEpochHash = await readFile(join(epochDataDir, EPOCH_HASH_FILENAME), 'hex').catch(() => Buffer.alloc(0)); if (storedEpochHash.toString() !== epochHash.toString()) { - await rm(dataDir, { recursive: true, force: true }); + await rm(epochDataDir, { recursive: true, force: true }); } - await mkdir(dataDir, { recursive: true }); - await writeFile(join(dataDir, EPOCH_HASH_FILENAME), epochHash.toString('hex')); + await mkdir(epochDataDir, { recursive: true }); + await writeFile(join(epochDataDir, EPOCH_HASH_FILENAME), epochHash.toString('hex')); - const store = AztecLmdbStore.open(dataDir); - this.log.debug(`Created new database for epoch ${epochNumber} at ${dataDir}`); + const store = AztecLmdbStore.open(epochDataDir, this.cacheMapSize); + this.log.debug(`Created new database for epoch ${epochNumber} at ${epochDataDir}`); const cleanup = () => store.close(); return new KVProverCache(store, cleanup); } @@ -43,11 +46,11 @@ export class ProverCacheManager { * @param upToAndIncludingEpoch - The epoch number up to which to remove caches */ public async removeStaleCaches(upToAndIncludingEpoch: bigint): Promise { - if (!this.cacheDir) { + if (!this.dataRootDir) { return; } - const entries: Dirent[] = await readdir(this.cacheDir, { withFileTypes: true }).catch(() => []); + const entries: Dirent[] = await readdir(this.dataRootDir, { withFileTypes: true }).catch(() => []); for (const item of entries) { if (!item.isDirectory()) { @@ -61,8 +64,10 @@ export class ProverCacheManager { const epochNumberInt = BigInt(epochNumber); if (epochNumberInt <= upToAndIncludingEpoch) { - this.log.info(`Removing old epoch database for epoch ${epochNumberInt} at ${join(this.cacheDir, item.name)}`); - await rm(join(this.cacheDir, item.name), { recursive: true }); + this.log.info( + `Removing old epoch database for epoch ${epochNumberInt} at ${join(this.dataRootDir, item.name)}`, + ); + await rm(join(this.dataRootDir, item.name), { recursive: true }); } } }