refactor: prover db config (#11126)

This PR refactors the types of the prover broker and agent config to
reuse more of the existing helpers.

Built on top of #10981 

Fix #10267

spartan/aztec-network/templates/prover-broker.yaml

@@ -103,9 +103,9 @@ spec:
               value: "{{ .Values.proverBroker.jobTimeoutMs }}"
             - name: PROVER_BROKER_JOB_MAX_RETRIES
               value: "{{ .Values.proverBroker.jobMaxRetries }}"
-            - name: PROVER_BROKER_DATA_DIRECTORY
+            - name: DATA_DIRECTORY
               value: "{{ .Values.proverBroker.dataDir }}"
-            - name: PROVER_BROKER_DATA_MAP_SIZE_KB
+            - name: DATA_STORE_MAP_SIZE_KB
               value: "{{ .Values.storage.dataStoreMapSize }}"
             - name: OTEL_RESOURCE_ATTRIBUTES
               value: service.name={{ .Release.Name }},service.namespace={{ .Release.Namespace }},service.version={{ .Chart.AppVersion }},environment={{ .Values.environment | default "production" }}

spartan/aztec-network/templates/prover-node.yaml

@@ -158,7 +158,7 @@ spec:
               value: "{{ .Values.proverNode.proverBroker.jobTimeoutMs }}"
             - name: PROVER_BROKER_JOB_MAX_RETRIES
               value: "{{ .Values.proverNode.proverBroker.jobMaxRetries }}"
-            - name: PROVER_BROKER_DATA_DIRECTORY
+            - name: DATA_DIRECTORY
               value: "{{ .Values.proverNode.proverBroker.dataDirectory }}"
             - name: PROVER_PUBLISHER_PRIVATE_KEY
               value: "{{ .Values.proverNode.proverPublisherPrivateKey }}"

yarn-project/aztec/src/cli/aztec_start_options.ts

@@ -2,12 +2,6 @@ import { type ArchiverConfig, archiverConfigMappings } from '@aztec/archiver/con
 import { faucetConfigMapping } from '@aztec/aztec-faucet/config';
 import { sequencerClientConfigMappings } from '@aztec/aztec-node/config';
 import { botConfigMappings } from '@aztec/bot/config';
-import {
-  type ProverAgentConfig,
-  type ProverBrokerConfig,
-  proverAgentConfigMappings,
-  proverBrokerConfigMappings,
-} from '@aztec/circuit-types/config';
 import {
   type ConfigMapping,
   type EnvVar,
@@ -17,6 +11,12 @@ import {
 } from '@aztec/foundation/config';
 import { bootnodeConfigMappings, p2pConfigMappings } from '@aztec/p2p/config';
 import { proofVerifierConfigMappings } from '@aztec/proof-verifier/config';
+import {
+  type ProverAgentConfig,
+  type ProverBrokerConfig,
+  proverAgentConfigMappings,
+  proverBrokerConfigMappings,
+} from '@aztec/prover-client/broker';
 import { proverNodeConfigMappings } from '@aztec/prover-node/config';
 import { allPxeConfigMappings } from '@aztec/pxe/config';
 import { telemetryClientConfigMappings } from '@aztec/telemetry-client/start';

yarn-project/aztec/src/cli/cmds/start_prover_agent.ts

@@ -1,9 +1,14 @@
-import { type ProverAgentConfig, proverAgentConfigMappings } from '@aztec/circuit-types';
 import { times } from '@aztec/foundation/collection';
 import { type NamespacedApiHandlers } from '@aztec/foundation/json-rpc/server';
 import { type LogFn } from '@aztec/foundation/log';
 import { buildServerCircuitProver } from '@aztec/prover-client';
-import { InlineProofStore, ProvingAgent, createProvingJobBrokerClient } from '@aztec/prover-client/broker';
+import {
+  InlineProofStore,
+  type ProverAgentConfig,
+  ProvingAgent,
+  createProvingJobBrokerClient,
+  proverAgentConfigMappings,
+} from '@aztec/prover-client/broker';
 import { getProverNodeAgentConfigFromEnv } from '@aztec/prover-node';
 import { createAndStartTelemetryClient, telemetryClientConfigMappings } from '@aztec/telemetry-client/start';
 

yarn-project/aztec/src/cli/cmds/start_prover_broker.ts

@@ -1,7 +1,12 @@
-import { type ProverBrokerConfig, type ProvingJobBroker, proverBrokerConfigMappings } from '@aztec/circuit-types';
+import { type ProvingJobBroker } from '@aztec/circuit-types';
 import { type NamespacedApiHandlers } from '@aztec/foundation/json-rpc/server';
 import { type LogFn } from '@aztec/foundation/log';
-import { ProvingJobBrokerSchema, createAndStartProvingBroker } from '@aztec/prover-client/broker';
+import {
+  type ProverBrokerConfig,
+  ProvingJobBrokerSchema,
+  createAndStartProvingBroker,
+  proverBrokerConfigMappings,
+} from '@aztec/prover-client/broker';
 import { getProverNodeBrokerConfigFromEnv } from '@aztec/prover-node';
 import {
   createAndStartTelemetryClient,

yarn-project/circuit-types/src/config.ts

@@ -1,3 +1 @@
-export { ProverAgentConfig, proverAgentConfigMappings } from './interfaces/prover-agent.js';
-export { ProverBrokerConfig, proverBrokerConfigMappings } from './interfaces/prover-broker.js';
 export { SequencerConfig, AllowedElement, SequencerConfigSchema } from './interfaces/configs.js';

yarn-project/circuit-types/src/interfaces/prover-agent.ts

@@ -1,63 +1,7 @@
-import { type ConfigMappingsType, booleanConfigHelper, numberConfigHelper } from '@aztec/foundation/config';
 import { type ApiSchemaFor } from '@aztec/foundation/schemas';
 
 import { z } from 'zod';
 
-import { ProvingRequestType } from './proving-job.js';
-
-export const ProverAgentConfig = z.object({
-  /** The number of prover agents to start */
-  proverAgentCount: z.number(),
-  /** The types of proofs the prover agent can generate */
-  proverAgentProofTypes: z.array(z.nativeEnum(ProvingRequestType)),
-  /** How often the prover agents poll for jobs */
-  proverAgentPollIntervalMs: z.number(),
-  /** The URL where this agent takes jobs from */
-  proverBrokerUrl: z.string().optional(),
-  /** Whether to construct real proofs */
-  realProofs: z.boolean(),
-  /** Artificial delay to introduce to all operations to the test prover. */
-  proverTestDelayMs: z.number(),
-});
-
-export type ProverAgentConfig = z.infer<typeof ProverAgentConfig>;
-
-export const proverAgentConfigMappings: ConfigMappingsType<ProverAgentConfig> = {
-  proverAgentCount: {
-    env: 'PROVER_AGENT_COUNT',
-    description: 'Whether this prover has a local prover agent',
-    ...numberConfigHelper(1),
-  },
-  proverAgentPollIntervalMs: {
-    env: 'PROVER_AGENT_POLL_INTERVAL_MS',
-    description: 'The interval agents poll for jobs at',
-    ...numberConfigHelper(100),
-  },
-  proverAgentProofTypes: {
-    env: 'PROVER_AGENT_PROOF_TYPES',
-    description: 'The types of proofs the prover agent can generate',
-    parseEnv: (val: string) =>
-      val
-        .split(',')
-        .map(v => ProvingRequestType[v as any])
-        .filter(v => typeof v === 'number'),
-  },
-  proverBrokerUrl: {
-    env: 'PROVER_BROKER_HOST',
-    description: 'The URL where this agent takes jobs from',
-  },
-  realProofs: {
-    env: 'PROVER_REAL_PROOFS',
-    description: 'Whether to construct real proofs',
-    ...booleanConfigHelper(false),
-  },
-  proverTestDelayMs: {
-    env: 'PROVER_TEST_DELAY_MS',
-    description: 'Artificial delay to introduce to all operations to the test prover.',
-    ...numberConfigHelper(0),
-  },
-};
-
 export interface ProverAgentApi {
   setMaxConcurrency(maxConcurrency: number): Promise<void>;
 

yarn-project/circuit-types/src/interfaces/prover-broker.ts

@@ -5,51 +5,6 @@ import {
   type ProvingJobStatus,
   type ProvingRequestType,
 } from '@aztec/circuit-types';
-import { type ConfigMappingsType, numberConfigHelper } from '@aztec/foundation/config';
-
-import { z } from 'zod';
-
-export const ProverBrokerConfig = z.object({
-  /** If starting a prover broker locally, the max number of retries per proving job */
-  proverBrokerJobMaxRetries: z.number(),
-  /** If starting a prover broker locally, the time after which a job times out and gets assigned to a different agent */
-  proverBrokerJobTimeoutMs: z.number(),
-  /** 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(),
-  /** The size of the data store map */
-  proverBrokerDataMapSizeKB: z.number(),
-});
-
-export type ProverBrokerConfig = z.infer<typeof ProverBrokerConfig>;
-
-export const proverBrokerConfigMappings: ConfigMappingsType<ProverBrokerConfig> = {
-  proverBrokerJobTimeoutMs: {
-    env: 'PROVER_BROKER_JOB_TIMEOUT_MS',
-    description: 'Jobs are retried if not kept alive for this long',
-    ...numberConfigHelper(30_000),
-  },
-  proverBrokerPollIntervalMs: {
-    env: 'PROVER_BROKER_POLL_INTERVAL_MS',
-    description: 'The interval to check job health status',
-    ...numberConfigHelper(1_000),
-  },
-  proverBrokerJobMaxRetries: {
-    env: 'PROVER_BROKER_JOB_MAX_RETRIES',
-    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',
-  },
-  proverBrokerDataMapSizeKB: {
-    env: 'PROVER_BROKER_DATA_MAP_SIZE_KB',
-    description: 'The size of the data store map',
-    ...numberConfigHelper(128 * 1_024 * 1_024), // Defaulted to 128 GB
-  },
-};
 
 /**
  * An interface for the proving orchestrator. The producer uses this to enqueue jobs for agents

yarn-project/foundation/src/config/env_var.ts

@@ -113,8 +113,6 @@ export type EnvVar =
   | 'PROVER_BROKER_JOB_TIMEOUT_MS'
   | 'PROVER_BROKER_POLL_INTERVAL_MS'
   | 'PROVER_BROKER_JOB_MAX_RETRIES'
-  | 'PROVER_BROKER_DATA_DIRECTORY'
-  | 'PROVER_BROKER_DATA_MAP_SIZE_KB'
   | 'PROVER_COORDINATION_NODE_URL'
   | 'PROVER_DISABLED'
   | 'PROVER_ID'

yarn-project/prover-client/src/config.ts

@@ -1,13 +1,13 @@
 import { type ACVMConfig, type BBConfig } from '@aztec/bb-prover';
+import { type ProverConfig, proverConfigMappings } from '@aztec/circuit-types';
+import { type ConfigMappingsType, booleanConfigHelper, getConfigFromMappings } from '@aztec/foundation/config';
+
 import {
   type ProverAgentConfig,
   type ProverBrokerConfig,
-  type ProverConfig,
   proverAgentConfigMappings,
   proverBrokerConfigMappings,
-  proverConfigMappings,
-} from '@aztec/circuit-types';
-import { type ConfigMappingsType, booleanConfigHelper, getConfigFromMappings } from '@aztec/foundation/config';
+} from './proving_broker/config.js';
 
 /**
  * The prover configuration.

yarn-project/prover-client/src/proving_broker/config.ts

@@ -0,0 +1,93 @@
+import { ProvingRequestType } from '@aztec/circuit-types';
+import { type ConfigMappingsType, booleanConfigHelper, numberConfigHelper } from '@aztec/foundation/config';
+import { type DataStoreConfig, dataConfigMappings } from '@aztec/kv-store/config';
+
+import { z } from 'zod';
+
+export const ProverBrokerConfig = z.object({
+  /** If starting a prover broker locally, the max number of retries per proving job */
+  proverBrokerJobMaxRetries: z.number(),
+  /** If starting a prover broker locally, the time after which a job times out and gets assigned to a different agent */
+  proverBrokerJobTimeoutMs: z.number(),
+  /** 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 */
+  dataDirectory: z.string().optional(),
+  /** The size of the data store map */
+  dataStoreMapSizeKB: z.number(),
+});
+
+export type ProverBrokerConfig = z.infer<typeof ProverBrokerConfig> &
+  Pick<DataStoreConfig, 'dataStoreMapSizeKB' | 'dataDirectory'>;
+
+export const proverBrokerConfigMappings: ConfigMappingsType<ProverBrokerConfig> = {
+  proverBrokerJobTimeoutMs: {
+    env: 'PROVER_BROKER_JOB_TIMEOUT_MS',
+    description: 'Jobs are retried if not kept alive for this long',
+    ...numberConfigHelper(30_000),
+  },
+  proverBrokerPollIntervalMs: {
+    env: 'PROVER_BROKER_POLL_INTERVAL_MS',
+    description: 'The interval to check job health status',
+    ...numberConfigHelper(1_000),
+  },
+  proverBrokerJobMaxRetries: {
+    env: 'PROVER_BROKER_JOB_MAX_RETRIES',
+    description: 'If starting a prover broker locally, the max number of retries per proving job',
+    ...numberConfigHelper(3),
+  },
+  ...dataConfigMappings,
+};
+
+export const ProverAgentConfig = z.object({
+  /** The number of prover agents to start */
+  proverAgentCount: z.number(),
+  /** The types of proofs the prover agent can generate */
+  proverAgentProofTypes: z.array(z.nativeEnum(ProvingRequestType)),
+  /** How often the prover agents poll for jobs */
+  proverAgentPollIntervalMs: z.number(),
+  /** The URL where this agent takes jobs from */
+  proverBrokerUrl: z.string().optional(),
+  /** Whether to construct real proofs */
+  realProofs: z.boolean(),
+  /** Artificial delay to introduce to all operations to the test prover. */
+  proverTestDelayMs: z.number(),
+});
+
+export type ProverAgentConfig = z.infer<typeof ProverAgentConfig>;
+
+export const proverAgentConfigMappings: ConfigMappingsType<ProverAgentConfig> = {
+  proverAgentCount: {
+    env: 'PROVER_AGENT_COUNT',
+    description: 'Whether this prover has a local prover agent',
+    ...numberConfigHelper(1),
+  },
+  proverAgentPollIntervalMs: {
+    env: 'PROVER_AGENT_POLL_INTERVAL_MS',
+    description: 'The interval agents poll for jobs at',
+    ...numberConfigHelper(100),
+  },
+  proverAgentProofTypes: {
+    env: 'PROVER_AGENT_PROOF_TYPES',
+    description: 'The types of proofs the prover agent can generate',
+    parseEnv: (val: string) =>
+      val
+        .split(',')
+        .map(v => ProvingRequestType[v as any])
+        .filter(v => typeof v === 'number'),
+  },
+  proverBrokerUrl: {
+    env: 'PROVER_BROKER_HOST',
+    description: 'The URL where this agent takes jobs from',
+  },
+  realProofs: {
+    env: 'PROVER_REAL_PROOFS',
+    description: 'Whether to construct real proofs',
+    ...booleanConfigHelper(false),
+  },
+  proverTestDelayMs: {
+    env: 'PROVER_TEST_DELAY_MS',
+    description: 'Artificial delay to introduce to all operations to the test prover.',
+    ...numberConfigHelper(0),
+  },
+};

yarn-project/prover-client/src/proving_broker/factory.ts

@@ -1,6 +1,6 @@
-import { type ProverBrokerConfig } from '@aztec/circuit-types';
 import { type TelemetryClient } from '@aztec/telemetry-client';
 
+import { type ProverBrokerConfig } from './config.js';
 import { ProvingBroker } from './proving_broker.js';
 import { InMemoryBrokerDatabase } from './proving_broker_database/memory.js';
 import { KVBrokerDatabase } from './proving_broker_database/persisted.js';
@@ -9,9 +9,7 @@ export async function createAndStartProvingBroker(
   config: ProverBrokerConfig,
   client: TelemetryClient,
 ): Promise<ProvingBroker> {
-  const database = config.proverBrokerDataDirectory
-    ? await KVBrokerDatabase.new(config, client)
-    : new InMemoryBrokerDatabase();
+  const database = config.dataDirectory ? await KVBrokerDatabase.new(config, client) : new InMemoryBrokerDatabase();
 
   const broker = new ProvingBroker(database, client, {
     jobTimeoutMs: config.proverBrokerJobTimeoutMs,

yarn-project/prover-client/src/proving_broker/index.ts

@@ -6,3 +6,4 @@ export * from './proving_broker_database/memory.js';
 export * from './proving_broker_database/persisted.js';
 export * from './proof_store.js';
 export * from './factory.js';
+export * from './config.js';

yarn-project/prover-client/src/proving_broker/proving_broker.test.ts

@@ -1,6 +1,5 @@
 import {
   type ProofUri,
-  type ProverBrokerConfig,
   type ProvingJob,
   type ProvingJobId,
   type ProvingJobStatus,
@@ -14,6 +13,7 @@ import { mkdtemp } from 'fs/promises';
 import { tmpdir } from 'os';
 import { join } from 'path';
 
+import { type ProverBrokerConfig } from './config.js';
 import { makeInputsUri, makeOutputsUri, makeRandomProvingJobId } from './fixtures.js';
 import { ProvingBroker } from './proving_broker.js';
 import { type ProvingBrokerDatabase } from './proving_broker_database.js';
@@ -25,8 +25,8 @@ describe.each([
   async () => {
     const directory = await mkdtemp(join(tmpdir(), 'proving-broker-test'));
     const config: ProverBrokerConfig = {
-      proverBrokerDataMapSizeKB: 1024 * 1024 * 1024, // 1GB
-      proverBrokerDataDirectory: directory,
+      dataStoreMapSizeKB: 1024 * 1024 * 1024, // 1GB
+      dataDirectory: directory,
       proverBrokerJobMaxRetries: 1,
       proverBrokerJobTimeoutMs: 1000,
       proverBrokerPollIntervalMs: 1000,