@aztec/prover-client 0.0.1-commit.b33fc05d0 → 0.0.1-commit.b3d3157a

package/src/prover-client/prover-client.ts

@@ -1,5 +1,7 @@
 import { type ACVMConfig, type BBConfig, BBNativeRollupProver, TestCircuitProver } from '@aztec/bb-prover';
+import type { EpochNumber } from '@aztec/foundation/branded-types';
 import { times } from '@aztec/foundation/collection';
+import type { Fr } from '@aztec/foundation/curves/bn254';
 import type { EthAddress } from '@aztec/foundation/eth-address';
 import { type Logger, createLogger } from '@aztec/foundation/log';
 import { NativeACVMSimulator } from '@aztec/simulator/server';
@@ -15,19 +17,64 @@ import {
   type ServerCircuitProver,
   tryStop,
 } from '@aztec/stdlib/interfaces/server';
+import type { CheckpointConstantData } from '@aztec/stdlib/rollup';
+import type { BlockHeader } from '@aztec/stdlib/tx';
 import { type TelemetryClient, getTelemetryClient } from '@aztec/telemetry-client';
 
 import type { ProverClientConfig } from '../config.js';
+import { CheckpointSubTreeOrchestrator } from '../orchestrator/checkpoint-sub-tree-orchestrator.js';
+import { EpochProvingContext } from '../orchestrator/epoch-proving-context.js';
 import { ProvingOrchestrator } from '../orchestrator/orchestrator.js';
+import { TopTreeOrchestrator } from '../orchestrator/top-tree-orchestrator.js';
 import { BrokerCircuitProverFacade } from '../proving_broker/broker_prover_facade.js';
 import { InlineProofStore, type ProofStore, createProofStore } from '../proving_broker/proof_store/index.js';
 import { ProvingAgent } from '../proving_broker/proving_agent.js';
 import { ServerEpochProver } from './server-epoch-prover.js';
 
+/**
+ * The factory surface that `EpochProvingJob` (in `prover-node`) depends on. Implemented
+ * by `ProverClient`. Defined here rather than in stdlib because the return types
+ * (`CheckpointSubTreeOrchestrator`, `TopTreeOrchestrator`) are concrete classes from
+ * this package.
+ *
+ * A single `BrokerCircuitProverFacade` is owned by `ProverClient` and shared across
+ * every orchestrator (every sub-tree and every top-tree across every concurrent epoch
+ * job). The broker delivers each completed-job notification exactly once (drained on
+ * the first `getCompletedJobs` poll), so multiple facades polling the same broker
+ * race and lose notifications
+ *
+ * The facade's job map cleans up entries on resolve/reject, and the prover-node
+ * keeps `ProverClient` alive for its whole lifetime
+ */
+export interface EpochProverFactory {
+  getProverId(): EthAddress;
+  /**
+   * Constructs a per-epoch shared context for the caching of e.g. chonk verifier results
+   */
+  createEpochProvingContext(epochNumber: EpochNumber): EpochProvingContext;
+  /**
+   * Constructs and starts a `CheckpointSubTreeOrchestrator` for a single checkpoint.
+   */
+  createCheckpointSubTreeOrchestrator(
+    epochContext: EpochProvingContext,
+    checkpointConstants: CheckpointConstantData,
+    l1ToL2Messages: Fr[],
+    totalNumBlocks: number,
+    headerOfLastBlockInPreviousCheckpoint: BlockHeader,
+  ): Promise<CheckpointSubTreeOrchestrator>;
+  createTopTreeOrchestrator(): TopTreeOrchestrator;
+}
+
 /** Manages proving of epochs by orchestrating the proving of individual blocks relying on a pool of prover agents. */
-export class ProverClient implements EpochProverManager {
+export class ProverClient implements EpochProverManager, EpochProverFactory {
   private running = false;
   private agents: ProvingAgent[] = [];
+  /**
+   * The single broker facade shared by every orchestrator created from this client.
+   * Constructed lazily on `start()` and torn down on `stop()` — see the comment on
+   * `EpochProverFactory` for why a single shared facade is required.
+   */
+  private facade: BrokerCircuitProverFacade | undefined;
 
   private constructor(
     private config: ProverClientConfig,
@@ -40,6 +87,38 @@ export class ProverClient implements EpochProverManager {
     private log: Logger = createLogger('prover-client:tx-prover'),
   ) {}
 
+  /**
+   * Lazy-init the shared facade. The broker delivers each completed-job notification
+   * exactly once (drained on the first `getCompletedJobs` poll), so we cannot start
+   * a shared facade alongside the per-call facades that `createEpochProver` builds —
+   * they would race for notifications and one side would silently drop them. Starting
+   * the shared facade only on first use of one of the new factory methods keeps the
+   * legacy `createEpochProver` path race-free.
+   */
+  private getFacade(): BrokerCircuitProverFacade {
+    if (!this.running) {
+      throw new Error('ProverClient is not running; call start() before constructing orchestrators.');
+    }
+    if (!this.facade) {
+      this.facade = new BrokerCircuitProverFacade(
+        this.orchestratorClient,
+        this.proofStore,
+        this.failedProofStore,
+        undefined,
+        this.log.getBindings(),
+      );
+      this.facade.start();
+    }
+    return this.facade;
+  }
+
+  /**
+   * Legacy single-class epoch prover. Each call constructs its own
+   * `BrokerCircuitProverFacade`; the new factory methods (`createCheckpointSubTreeOrchestrator`,
+   * `createTopTreeOrchestrator`, `createEpochProvingContext`) share a single facade
+   * owned by `ProverClient`. Both APIs coexist while the prover-node migrates onto
+   * the new pair.
+   */
   public createEpochProver(): EpochProver {
     const bindings = this.log.getBindings();
     const facade = new BrokerCircuitProverFacade(
@@ -61,6 +140,43 @@ export class ProverClient implements EpochProverManager {
     return new ServerEpochProver(facade, orchestrator);
   }
 
+  public createEpochProvingContext(epochNumber: EpochNumber): EpochProvingContext {
+    return new EpochProvingContext(this.getFacade(), epochNumber, this.log.getBindings());
+  }
+
+  public createCheckpointSubTreeOrchestrator(
+    epochContext: EpochProvingContext,
+    checkpointConstants: CheckpointConstantData,
+    l1ToL2Messages: Fr[],
+    totalNumBlocks: number,
+    headerOfLastBlockInPreviousCheckpoint: BlockHeader,
+  ): Promise<CheckpointSubTreeOrchestrator> {
+    return CheckpointSubTreeOrchestrator.start(
+      this.worldState,
+      this.getFacade(),
+      this.config.proverId,
+      epochContext,
+      this.config.cancelJobsOnStop,
+      this.config.enqueueConcurrency,
+      checkpointConstants,
+      l1ToL2Messages,
+      totalNumBlocks,
+      headerOfLastBlockInPreviousCheckpoint,
+      this.telemetry,
+      this.log.getBindings(),
+    );
+  }
+
+  public createTopTreeOrchestrator(): TopTreeOrchestrator {
+    return new TopTreeOrchestrator(
+      this.getFacade(),
+      this.config.proverId,
+      this.config.enqueueConcurrency,
+      this.telemetry,
+      this.log.getBindings(),
+    );
+  }
+
   public getProverId(): EthAddress {
     return this.config.proverId;
   }
@@ -100,6 +216,14 @@ export class ProverClient implements EpochProverManager {
     }
     this.running = false;
     await this.stopAgents();
+    if (this.facade) {
+      try {
+        await this.facade.stop();
+      } catch (err) {
+        this.log.error('Error stopping shared broker facade', err);
+      }
+      this.facade = undefined;
+    }
     await tryStop(this.orchestratorClient);
   }
 

package/src/proving_broker/broker_prover_facade.ts

@@ -4,7 +4,9 @@ import type {
   NESTED_RECURSIVE_ROLLUP_HONK_PROOF_LENGTH,
   RECURSIVE_PROOF_LENGTH,
 } from '@aztec/constants';
+import { asyncPool } from '@aztec/foundation/async-pool';
 import { EpochNumber } from '@aztec/foundation/branded-types';
+import { chunk } from '@aztec/foundation/collection';
 import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import { type PromiseWithResolvers, RunningPromise, promiseWithResolvers } from '@aztec/foundation/promise';
 import { truncate } from '@aztec/foundation/string';
@@ -226,17 +228,11 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
       // We collect all returned notifications and return them
       const allCompleted = new Set<ProvingJobId>();
       try {
-        let numRequests = 0;
-        while (ids.length > 0) {
-          const slice = ids.splice(0, SNAPSHOT_SYNC_CHECK_MAX_REQUEST_SIZE);
-          const completed = await this.broker.getCompletedJobs(slice);
+        const batches = ids.length > 0 ? chunk(ids, SNAPSHOT_SYNC_CHECK_MAX_REQUEST_SIZE) : [[]];
+        await asyncPool(1, batches, async batch => {
+          const completed = await this.broker.getCompletedJobs(batch);
           completed.forEach(id => allCompleted.add(id));
-          ++numRequests;
-        }
-        if (numRequests === 0) {
-          const final = await this.broker.getCompletedJobs([]);
-          final.forEach(id => allCompleted.add(id));
-        }
+        });
       } catch (err) {
         this.log.error(`Error thrown when requesting completed job notifications from the broker`, err);
       }
@@ -352,12 +348,8 @@ export class BrokerCircuitProverFacade implements ServerCircuitProver {
       .map(id => this.jobs.get(id)!)
       .filter(x => x !== undefined);
     const totalJobsToRetrieve = toBeRetrieved.length;
-    let totalJobsRetrieved = 0;
-    while (toBeRetrieved.length > 0) {
-      const slice = toBeRetrieved.splice(0, MAX_CONCURRENT_JOB_SETTLED_REQUESTS);
-      const results = await Promise.all(slice.map(job => processJob(job!)));
-      totalJobsRetrieved += results.filter(x => x).length;
-    }
+    const results = await asyncPool(MAX_CONCURRENT_JOB_SETTLED_REQUESTS, toBeRetrieved, job => processJob(job));
+    const totalJobsRetrieved = results.filter(x => x).length;
     if (totalJobsToRetrieve > 0) {
       this.log.verbose(
         `Successfully retrieved ${totalJobsRetrieved} of ${totalJobsToRetrieve} jobs that should be ready, total ready jobs is now: ${this.jobsToRetrieve.size}`,

package/src/proving_broker/config.ts

@@ -4,6 +4,7 @@ import {
   booleanConfigHelper,
   getDefaultConfig,
   numberConfigHelper,
+  optionalNumberConfigHelper,
 } from '@aztec/foundation/config';
 import { pickConfigMappings } from '@aztec/foundation/config';
 import { type ChainConfig, chainConfigMappings } from '@aztec/stdlib/config';
@@ -73,7 +74,7 @@ export const proverBrokerConfigMappings: ConfigMappingsType<ProverBrokerConfig>
   },
   proverBrokerStoreMapSizeKb: {
     env: 'PROVER_BROKER_STORE_MAP_SIZE_KB',
-    parseEnv: (val: string | undefined) => (val ? +val : undefined),
+    ...optionalNumberConfigHelper(),
     description: "The size of the prover broker's database. Will override the dataStoreMapSizeKb if set.",
   },
   proverBrokerDebugReplayEnabled: {

package/src/proving_broker/index.ts

@@ -7,3 +7,4 @@ export * from './proving_broker_database/persisted.js';
 export * from './proof_store/index.js';
 export * from './factory.js';
 export * from './config.js';
+export { BrokerCircuitProverFacade } from './broker_prover_facade.js';

package/src/proving_broker/proving_broker.ts

@@ -319,14 +319,20 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Pr
   }
 
   private cleanUpProvingJobState(ids: ProvingJobId[]) {
+    const idsToClean = new Set(ids);
     for (const id of ids) {
       this.jobsCache.delete(id);
+      const deferred = this.promises.get(id);
+      if (deferred) {
+        deferred.resolve({ status: 'rejected', reason: 'Proving job cleaned up' });
+      }
       this.promises.delete(id);
       this.resultsCache.delete(id);
       this.inProgress.delete(id);
       this.retries.delete(id);
       this.enqueuedAt.delete(id);
     }
+    this.completedJobNotifications = this.completedJobNotifications.filter(id => !idsToClean.has(id));
   }
 
   #getProvingJobStatus(id: ProvingJobId): ProvingJobStatus {
@@ -594,21 +600,21 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Pr
   }
 
   private async cleanupPass() {
-    this.cleanupStaleJobs();
     this.reEnqueueExpiredJobs();
     const oldestEpochToKeep = this.oldestEpochToKeep();
     if (oldestEpochToKeep > 0) {
+      this.cleanupJobsOlderThanEpoch(EpochNumber(oldestEpochToKeep));
       await this.database.deleteAllProvingJobsOlderThanEpoch(EpochNumber(oldestEpochToKeep));
       this.logger.trace(`Deleted all epochs older than ${oldestEpochToKeep}`);
     }
   }
 
-  private cleanupStaleJobs() {
+  private cleanupJobsOlderThanEpoch(epochNumber: EpochNumber) {
     const jobIds = Array.from(this.jobsCache.keys());
     const jobsToClean: ProvingJobId[] = [];
     for (const id of jobIds) {
       const job = this.jobsCache.get(id)!;
-      if (this.isJobStale(job)) {
+      if (job.epochNumber < epochNumber) {
         jobsToClean.push(id);
       }
     }
@@ -632,10 +638,26 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Pr
       const now = this.msTimeSource();
       const msSinceLastUpdate = now - metadata.lastUpdatedAt;
       if (msSinceLastUpdate >= this.jobTimeoutMs) {
-        this.logger.warn(`Proving job id=${id} timed out. Adding it back to the queue.`, { provingJobId: id });
         this.inProgress.delete(id);
-        this.enqueueJobInternal(item);
         this.instrumentation.incTimedOutJobs(item.type);
+
+        const retries = this.retries.get(id) ?? 0;
+        if (retries + 1 < this.maxRetries && !this.isJobStale(item)) {
+          this.logger.warn(`Proving job id=${id} timed out. Re-enqueueing (retry ${retries + 1}/${this.maxRetries}).`, {
+            provingJobId: id,
+          });
+          this.retries.set(id, retries + 1);
+          this.enqueueJobInternal(item);
+        } else {
+          this.logger.error(`Proving job id=${id} timed out after ${retries + 1} attempts. Marking as failed.`, {
+            provingJobId: id,
+          });
+          const result: ProvingJobSettledResult = { status: 'rejected', reason: 'Timed out' };
+          this.resultsCache.set(id, result);
+          this.promises.get(id)?.resolve(result);
+          this.completedJobNotifications.push(id);
+          this.instrumentation.incRejectedJobs(item.type);
+        }
       }
     }
   }

package/src/proving_broker/proving_broker_database/persisted.ts

@@ -154,7 +154,7 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       const db = await openVersionedStoreAt(
         fullDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
-        config.l1Contracts.rollupAddress,
+        config.rollupAddress,
         config.dataStoreMapSizeKb,
       );
       const epochDb = new SingleEpochDatabase(db);
@@ -222,7 +222,7 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       const db = await openVersionedStoreAt(
         newEpochDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
-        this.config.l1Contracts.rollupAddress,
+        this.config.rollupAddress,
         this.config.dataStoreMapSizeKb,
       );
       epochDb = new SingleEpochDatabase(db);

package/src/proving_broker/rpc.ts

@@ -18,6 +18,15 @@ import { makeTracedFetch } from '@aztec/telemetry-client';
 
 import { z } from 'zod';
 
+/** Indefinite backoff for broker communication: 1, 1, 1, 2, 4, 4, 4, ... seconds. */
+export function* proverBrokerBackoff() {
+  const v = [1, 1, 1, 2, 4];
+  let i = 0;
+  while (true) {
+    yield v[Math.min(i++, v.length - 1)];
+  }
+}
+
 const ProvingJobFilterSchema = z.object({
   allowList: z.array(z.nativeEnum(ProvingRequestType)),
 });
@@ -28,26 +37,29 @@ const GetProvingJobResponse = z.object({
 });
 
 export const ProvingJobProducerSchema: ApiSchemaFor<ProvingJobProducer> = {
-  enqueueProvingJob: z.function().args(ProvingJob).returns(ProvingJobStatus),
-  getProvingJobStatus: z.function().args(ProvingJobId).returns(ProvingJobStatus),
-  cancelProvingJob: z.function().args(ProvingJobId).returns(z.void()),
-  getCompletedJobs: z.function().args(z.array(ProvingJobId)).returns(z.array(ProvingJobId)),
+  enqueueProvingJob: z.function({ input: z.tuple([ProvingJob]), output: ProvingJobStatus }),
+  getProvingJobStatus: z.function({ input: z.tuple([ProvingJobId]), output: ProvingJobStatus }),
+  cancelProvingJob: z.function({ input: z.tuple([ProvingJobId]), output: z.void() }),
+  getCompletedJobs: z.function({ input: z.tuple([z.array(ProvingJobId)]), output: z.array(ProvingJobId) }),
 };
 
 export const ProvingJobConsumerSchema: ApiSchemaFor<ProvingJobConsumer> = {
-  getProvingJob: z.function().args(optional(ProvingJobFilterSchema)).returns(GetProvingJobResponse.optional()),
-  reportProvingJobError: z
-    .function()
-    .args(ProvingJobId, z.string(), optional(z.boolean()), optional(ProvingJobFilterSchema))
-    .returns(GetProvingJobResponse.optional()),
-  reportProvingJobProgress: z
-    .function()
-    .args(ProvingJobId, z.number(), optional(ProvingJobFilterSchema))
-    .returns(GetProvingJobResponse.optional()),
-  reportProvingJobSuccess: z
-    .function()
-    .args(ProvingJobId, ProofUri, optional(ProvingJobFilterSchema))
-    .returns(GetProvingJobResponse.optional()),
+  getProvingJob: z.function({
+    input: z.tuple([optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
+  reportProvingJobError: z.function({
+    input: z.tuple([ProvingJobId, z.string(), optional(z.boolean()), optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
+  reportProvingJobProgress: z.function({
+    input: z.tuple([ProvingJobId, z.number(), optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
+  reportProvingJobSuccess: z.function({
+    input: z.tuple([ProvingJobId, ProofUri, optional(ProvingJobFilterSchema)]),
+    output: GetProvingJobResponse.optional(),
+  }),
 };
 
 export const ProvingJobBrokerSchema: ApiSchemaFor<ProvingJobBroker> = {
@@ -56,10 +68,10 @@ export const ProvingJobBrokerSchema: ApiSchemaFor<ProvingJobBroker> = {
 };
 
 export const ProvingJobBrokerDebugSchema: ApiSchemaFor<ProvingJobBrokerDebug> = {
-  replayProvingJob: z
-    .function()
-    .args(ProvingJobId, z.nativeEnum(ProvingRequestType), EpochNumberSchema, ProofUri)
-    .returns(ProvingJobStatus),
+  replayProvingJob: z.function({
+    input: z.tuple([ProvingJobId, z.nativeEnum(ProvingRequestType), EpochNumberSchema, ProofUri]),
+    output: ProvingJobStatus,
+  }),
 };
 
 export const ProvingJobBrokerSchemaWithDebug: ApiSchemaFor<ProvingJobBroker & ProvingJobBrokerDebug> = {
@@ -70,7 +82,7 @@ export const ProvingJobBrokerSchemaWithDebug: ApiSchemaFor<ProvingJobBroker & Pr
 export function createProvingJobBrokerClient(
   url: string,
   versions: Partial<ComponentsVersions>,
-  fetch = makeTracedFetch([1, 2, 3], false),
+  fetch = makeTracedFetch(proverBrokerBackoff, false),
 ): ProvingJobBroker {
   return createSafeJsonRpcClient(url, ProvingJobBrokerSchema, {
     namespaceMethods: 'proverBroker',
@@ -82,7 +94,7 @@ export function createProvingJobBrokerClient(
 export function createProvingJobProducerClient(
   url: string,
   versions: Partial<ComponentsVersions>,
-  fetch = makeTracedFetch([1, 2, 3], false),
+  fetch = makeTracedFetch(proverBrokerBackoff, false),
 ): ProvingJobProducer {
   return createSafeJsonRpcClient(url, ProvingJobProducerSchema, {
     namespaceMethods: 'provingJobProducer',
@@ -94,7 +106,7 @@ export function createProvingJobProducerClient(
 export function createProvingJobConsumerClient(
   url: string,
   versions: Partial<ComponentsVersions>,
-  fetch = makeTracedFetch([1, 2, 3], false),
+  fetch = makeTracedFetch(proverBrokerBackoff, false),
 ): ProvingJobConsumer {
   return createSafeJsonRpcClient(url, ProvingJobConsumerSchema, {
     namespaceMethods: 'provingJobConsumer',