@aztec/prover-client 0.0.0-test.1 → 0.0.1-commit.b655e406

package/src/proving_broker/proving_agent.ts

@@ -1,10 +1,11 @@
+import { AbortError } from '@aztec/foundation/error';
 import { createLogger } from '@aztec/foundation/log';
 import { RunningPromise } from '@aztec/foundation/running-promise';
 import { truncate } from '@aztec/foundation/string';
-import { Timer } from '@aztec/foundation/timer';
 import { ProvingError } from '@aztec/stdlib/errors';
 import type {
-  ProvingJob,
+  GetProvingJobResponse,
+  ProverAgentStatus,
   ProvingJobConsumer,
   ProvingJobId,
   ProvingJobInputs,
@@ -31,7 +32,6 @@ export class ProvingAgent implements Traceable {
   private currentJobController?: ProvingJobController;
   private runningPromise: RunningPromise;
   private instrumentation: ProvingAgentInstrumentation;
-  private idleTimer: Timer | undefined;
 
   public readonly tracer: Tracer;
 
@@ -64,7 +64,6 @@ export class ProvingAgent implements Traceable {
   }
 
   public start(): void {
-    this.idleTimer = new Timer();
     this.runningPromise.start();
   }
 
@@ -73,41 +72,78 @@ export class ProvingAgent implements Traceable {
     await this.runningPromise.stop();
   }
 
+  public getStatus(): ProverAgentStatus {
+    if (this.currentJobController) {
+      return {
+        status: 'proving',
+        jobId: this.currentJobController.getJobId(),
+        proofType: this.currentJobController.getProofType(),
+        startedAtISO: new Date(this.currentJobController.getStartedAt()).toISOString(),
+      };
+    }
+
+    return this.runningPromise.isRunning() ? { status: 'running' } : { status: 'stopped' };
+  }
+
   @trackSpan('ProvingAgent.safeWork')
   private async work() {
-    // every tick we need to
-    // (1) either do a heartbeat, telling the broker that we're working
-    // (2) get a new job
-    // If during (1) the broker returns a new job that means we can cancel the current job and start the new one
-    let maybeJob: { job: ProvingJob; time: number } | undefined;
-    if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
-      maybeJob = await this.broker.reportProvingJobProgress(
-        this.currentJobController.getJobId(),
-        this.currentJobController.getStartedAt(),
-        { allowList: this.proofAllowList },
-      );
+    // every tick we need to take one of the following actions:
+    // 1. send a hearbeat to the broker that we're working on some job
+    // 2. if the job is complete, send its result to the broker
+    // 3. get a job from the broker
+    // Any one of these actions could give us a new job to work on. If that happens we abort the current job.
+    //
+    // This loop gets triggered in one of two ways:
+    // - either on a timer (see pollIntervalMs)
+    // - or when a proof completes
+    let maybeJob: GetProvingJobResponse | undefined;
+
+    if (this.currentJobController) {
+      const status = this.currentJobController.getStatus();
+      const jobId = this.currentJobController.getJobId();
+      const proofType = this.currentJobController.getProofType();
+      const startedAt = this.currentJobController.getStartedAt();
+      const result = this.currentJobController.getResult();
+
+      if (status === ProvingJobControllerStatus.RUNNING) {
+        maybeJob = await this.broker.reportProvingJobProgress(jobId, startedAt, { allowList: this.proofAllowList });
+      } else if (status === ProvingJobControllerStatus.DONE) {
+        if (result) {
+          maybeJob = await this.reportResult(jobId, proofType, result);
+        } else {
+          this.log.warn(
+            `Job controller for job ${this.currentJobController.getJobId()} is done but doesn't have a result`,
+            { jobId },
+          );
+          maybeJob = await this.reportResult(
+            jobId,
+            proofType,
+            new ProvingError('No result found after proving', undefined, /* retry */ true),
+          );
+        }
+
+        this.currentJobController = undefined;
+      } else {
+        // IDLE status should not be seen because a job is started as soon as it is created
+        this.log.warn(`Idle job controller for job: ${this.currentJobController.getJobId()}. Skipping main loop work`, {
+          jobId: this.currentJobController.getJobId(),
+        });
+        return;
+      }
     } else {
       maybeJob = await this.broker.getProvingJob({ allowList: this.proofAllowList });
     }
 
-    if (!maybeJob) {
-      return;
-    }
-
-    if (this.idleTimer) {
-      this.instrumentation.recordIdleTime(this.idleTimer);
+    if (maybeJob) {
+      await this.startJob(maybeJob);
     }
-    this.idleTimer = undefined;
-
-    const { job, time } = maybeJob;
-    await this.startJob(job, time);
   }
 
-  private async startJob(job: ProvingJob, startedAt: number): Promise<void> {
+  private async startJob({ job, time: startedAt }: GetProvingJobResponse): Promise<void> {
     let abortedProofJobId: string | undefined;
     let abortedProofName: string | undefined;
 
-    if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
+    if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.RUNNING) {
       abortedProofJobId = this.currentJobController.getJobId();
       abortedProofName = this.currentJobController.getProofTypeName();
       this.currentJobController?.abort();
@@ -116,13 +152,13 @@ export class ProvingAgent implements Traceable {
     let inputs: ProvingJobInputs;
     try {
       inputs = await this.proofStore.getProofInput(job.inputsUri);
-    } catch (err) {
+    } catch {
       const maybeJob = await this.broker.reportProvingJobError(job.id, 'Failed to load proof inputs', true, {
         allowList: this.proofAllowList,
       });
 
       if (maybeJob) {
-        return this.startJob(maybeJob.job, maybeJob.time);
+        return this.startJob(maybeJob);
       }
 
       return;
@@ -134,7 +170,11 @@ export class ProvingAgent implements Traceable {
       job.epochNumber,
       startedAt,
       this.circuitProver,
-      this.handleJobResult,
+      () => {
+        // trigger a run of the main work loop when proving completes
+        // no need to await this here. The controller will stay alive (in DONE state) until the result is send to the broker
+        void this.runningPromise.trigger();
+      },
     );
 
     if (abortedProofJobId) {
@@ -154,28 +194,30 @@ export class ProvingAgent implements Traceable {
     this.currentJobController.start();
   }
 
-  handleJobResult = async <T extends ProvingRequestType>(
+  private async reportResult<T extends ProvingRequestType>(
     jobId: ProvingJobId,
     type: T,
-    err: Error | undefined,
-    result: ProvingJobResultsMap[T] | undefined,
-  ) => {
-    let maybeJob: { job: ProvingJob; time: number } | undefined;
-    if (err) {
-      const retry = err.name === ProvingError.NAME ? (err as ProvingError).retry : false;
-      this.log.error(`Job id=${jobId} type=${ProvingRequestType[type]} failed err=${err.message} retry=${retry}`, err);
-      maybeJob = await this.broker.reportProvingJobError(jobId, err.message, retry, { allowList: this.proofAllowList });
-    } else if (result) {
+    result: ProvingJobResultsMap[T] | Error,
+  ): Promise<GetProvingJobResponse | undefined> {
+    let maybeJob: GetProvingJobResponse | undefined;
+    if (result instanceof AbortError) {
+      // no-op
+      this.log.warn(`Job id=${jobId} was aborted. Not reporting result back to broker`, result);
+    } else if (result instanceof Error) {
+      const retry = result.name === ProvingError.NAME ? (result as ProvingError).retry : false;
+      this.log.error(
+        `Job id=${jobId} type=${ProvingRequestType[type]} failed err=${result.message} retry=${retry}`,
+        result,
+      );
+      maybeJob = await this.broker.reportProvingJobError(jobId, result.message, retry, {
+        allowList: this.proofAllowList,
+      });
+    } else {
       const outputUri = await this.proofStore.saveProofOutput(jobId, type, result);
       this.log.info(`Job id=${jobId} type=${ProvingRequestType[type]} completed outputUri=${truncate(outputUri)}`);
       maybeJob = await this.broker.reportProvingJobSuccess(jobId, outputUri, { allowList: this.proofAllowList });
     }
 
-    if (maybeJob) {
-      const { job, time } = maybeJob;
-      await this.startJob(job, time);
-    } else {
-      this.idleTimer = new Timer();
-    }
-  };
+    return maybeJob;
+  }
 }

package/src/proving_broker/proving_broker.ts

@@ -2,16 +2,17 @@ import { createLogger } from '@aztec/foundation/log';
 import { type PromiseWithResolvers, RunningPromise, promiseWithResolvers } from '@aztec/foundation/promise';
 import { PriorityMemoryQueue } from '@aztec/foundation/queue';
 import { Timer } from '@aztec/foundation/timer';
-import type {
-  GetProvingJobResponse,
-  ProofUri,
-  ProvingJob,
-  ProvingJobConsumer,
-  ProvingJobFilter,
-  ProvingJobId,
-  ProvingJobProducer,
-  ProvingJobSettledResult,
-  ProvingJobStatus,
+import {
+  type GetProvingJobResponse,
+  type ProofUri,
+  type ProvingJob,
+  type ProvingJobConsumer,
+  type ProvingJobFilter,
+  type ProvingJobId,
+  type ProvingJobProducer,
+  type ProvingJobSettledResult,
+  type ProvingJobStatus,
+  tryStop,
 } from '@aztec/stdlib/interfaces/server';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
 import {
@@ -43,20 +44,33 @@ type EnqueuedProvingJob = Pick<ProvingJob, 'id' | 'epochNumber'>;
 export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Traceable {
   private queues: ProvingQueues = {
     [ProvingRequestType.PUBLIC_VM]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.TUBE_PROOF]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PUBLIC_CHONK_VERIFIER]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
-    [ProvingRequestType.PRIVATE_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.PUBLIC_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PRIVATE_TX_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PUBLIC_TX_BASE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.TX_MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
     [ProvingRequestType.ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
     [ProvingRequestType.BLOCK_MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.BLOCK_ROOT_FIRST_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.BLOCK_ROOT_SINGLE_TX_FIRST_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(
+      provingJobComparator,
+    ),
+    [ProvingRequestType.BLOCK_ROOT_EMPTY_TX_FIRST_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(
+      provingJobComparator,
+    ),
     [ProvingRequestType.BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.SINGLE_TX_BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.BLOCK_ROOT_SINGLE_TX_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+
+    [ProvingRequestType.CHECKPOINT_ROOT_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.CHECKPOINT_ROOT_SINGLE_BLOCK_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(
+      provingJobComparator,
+    ),
+    [ProvingRequestType.CHECKPOINT_MERGE_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.CHECKPOINT_PADDING_ROLLUP]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
 
-    [ProvingRequestType.BASE_PARITY]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
-    [ProvingRequestType.ROOT_PARITY]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PARITY_BASE]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
+    [ProvingRequestType.PARITY_ROOT]: new PriorityMemoryQueue<EnqueuedProvingJob>(provingJobComparator),
   };
 
   // holds a copy of the database in memory in order to quickly fulfill requests
@@ -93,7 +107,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
   /**
    * The broker keeps track of the highest epoch its seen.
    * This information is used for garbage collection: once it reaches the next epoch, it can start pruning the database of old state.
-   * It is important that this value is initialised to zero. This ensures that we don't delete any old jobs until the current
+   * It is important that this value is initialized to zero. This ensures that we don't delete any old jobs until the current
    * process instance receives a job request informing it of the actual current highest epoch
    * Example:
    * proving epoch 11 - the broker will wipe all jobs for epochs 9 and lower
@@ -184,7 +198,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
       this.logger.warn('ProvingBroker not started');
       return Promise.resolve();
     }
-    await this.cleanupPromise.stop();
+    await tryStop(this.cleanupPromise);
   }
 
   public enqueueProvingJob(job: ProvingJob): Promise<ProvingJobStatus> {
@@ -313,7 +327,6 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
     return Promise.resolve(notifications.concat(completedJobs));
   }
 
-  // eslint-disable-next-line require-await
   #getProvingJob(filter: ProvingJobFilter = { allowList: [] }): { job: ProvingJob; time: number } | undefined {
     const allowedProofs: ProvingRequestType[] =
       Array.isArray(filter.allowList) && filter.allowList.length > 0
@@ -571,7 +584,7 @@ export class ProvingBroker implements ProvingJobProducer, ProvingJobConsumer, Tr
 
     if (jobsToClean.length > 0) {
       this.cleanUpProvingJobState(jobsToClean);
-      this.logger.info(`Cleaned up jobs=${jobsToClean.length}`);
+      this.logger.verbose(`Cleaned up proving jobs=${jobsToClean.length}`);
     }
   }
 
@@ -671,17 +684,23 @@ function proofTypeComparator(a: ProvingRequestType, b: ProvingRequestType): -1 |
  * The aim is that this will speed up block proving as the closer we get to a block's root proof the more likely it
  * is to get picked up by agents
  */
-const PROOF_TYPES_IN_PRIORITY_ORDER: ProvingRequestType[] = [
+export const PROOF_TYPES_IN_PRIORITY_ORDER: ProvingRequestType[] = [
+  ProvingRequestType.ROOT_ROLLUP,
+  ProvingRequestType.BLOCK_ROOT_FIRST_ROLLUP,
+  ProvingRequestType.BLOCK_ROOT_SINGLE_TX_FIRST_ROLLUP,
+  ProvingRequestType.BLOCK_ROOT_EMPTY_TX_FIRST_ROLLUP,
   ProvingRequestType.BLOCK_ROOT_ROLLUP,
-  ProvingRequestType.SINGLE_TX_BLOCK_ROOT_ROLLUP,
+  ProvingRequestType.BLOCK_ROOT_SINGLE_TX_ROLLUP,
   ProvingRequestType.BLOCK_MERGE_ROLLUP,
-  ProvingRequestType.ROOT_ROLLUP,
-  ProvingRequestType.MERGE_ROLLUP,
-  ProvingRequestType.PUBLIC_BASE_ROLLUP,
-  ProvingRequestType.PRIVATE_BASE_ROLLUP,
+  ProvingRequestType.CHECKPOINT_ROOT_ROLLUP,
+  ProvingRequestType.CHECKPOINT_ROOT_SINGLE_BLOCK_ROLLUP,
+  ProvingRequestType.CHECKPOINT_MERGE_ROLLUP,
+  ProvingRequestType.CHECKPOINT_PADDING_ROLLUP,
+  ProvingRequestType.TX_MERGE_ROLLUP,
+  ProvingRequestType.PUBLIC_TX_BASE_ROLLUP,
+  ProvingRequestType.PRIVATE_TX_BASE_ROLLUP,
   ProvingRequestType.PUBLIC_VM,
-  ProvingRequestType.TUBE_PROOF,
-  ProvingRequestType.ROOT_PARITY,
-  ProvingRequestType.BASE_PARITY,
-  ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP,
+  ProvingRequestType.PUBLIC_CHONK_VERIFIER,
+  ProvingRequestType.PARITY_ROOT,
+  ProvingRequestType.PARITY_BASE,
 ];

package/src/proving_broker/proving_broker_database/memory.ts

@@ -52,7 +52,7 @@ export class InMemoryBrokerDatabase implements ProvingBrokerDatabase {
   }
 
   async *allProvingJobs(): AsyncIterableIterator<[ProvingJob, ProvingJobSettledResult | undefined]> {
-    for await (const item of this.jobs.values()) {
+    for (const item of this.jobs.values()) {
       yield [item, this.results.get(item.id)] as const;
     }
   }

package/src/proving_broker/proving_broker_database/persisted.ts

@@ -46,9 +46,9 @@ class SingleEpochDatabase {
 
   async *allProvingJobs(): AsyncIterableIterator<[ProvingJob, ProvingJobSettledResult | undefined]> {
     for await (const jobStr of this.jobs.valuesAsync()) {
-      const job = await jsonParseWithSchema(jobStr, ProvingJob);
+      const job = jsonParseWithSchema(jobStr, ProvingJob);
       const resultStr = await this.jobResults.getAsync(job.id);
-      const result = resultStr ? await jsonParseWithSchema(resultStr, ProvingJobSettledResult) : undefined;
+      const result = resultStr ? jsonParseWithSchema(resultStr, ProvingJobSettledResult) : undefined;
       yield [job, result];
     }
   }
@@ -111,7 +111,8 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
   private async estimateSize() {
     const sizes = await Promise.all(Array.from(this.epochs.values()).map(x => x.estimateSize()));
     return {
-      mappingSize: this.config.dataStoreMapSizeKB,
+      mappingSize: this.config.dataStoreMapSizeKb,
+      physicalFileSize: sizes.reduce((prev, curr) => prev + curr.physicalFileSize, 0),
       numItems: sizes.reduce((prev, curr) => prev + curr.numItems, 0),
       actualSize: sizes.reduce((prev, curr) => prev + curr.actualSize, 0),
     };
@@ -136,13 +137,13 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
         continue;
       }
       logger.info(
-        `Loading broker database for epoch ${epochNumber} from ${fullDirectory} with map size ${config.dataStoreMapSizeKB}KB`,
+        `Loading broker database for epoch ${epochNumber} from ${fullDirectory} with map size ${config.dataStoreMapSizeKb}KB`,
       );
       const db = await openVersionedStoreAt(
         fullDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
         config.l1Contracts.rollupAddress,
-        config.dataStoreMapSizeKB,
+        config.dataStoreMapSizeKb,
       );
       const epochDb = new SingleEpochDatabase(db);
       epochs.set(epochNumber, epochDb);
@@ -170,7 +171,7 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       if (!db) {
         continue;
       }
-      this.logger.info(`Deleting broker database for epoch ${old}`);
+      this.logger.verbose(`Deleting broker database for epoch ${old}`);
       await db.delete();
       this.epochs.delete(old);
     }
@@ -201,13 +202,13 @@ export class KVBrokerDatabase implements ProvingBrokerDatabase {
       const newEpochDirectory = join(this.config.dataDirectory!, epochNumber.toString());
       await mkdir(newEpochDirectory, { recursive: true });
       this.logger.info(
-        `Creating broker database for epoch ${epochNumber} at ${newEpochDirectory} with map size ${this.config.dataStoreMapSizeKB}`,
+        `Creating broker database for epoch ${epochNumber} at ${newEpochDirectory} with map size ${this.config.dataStoreMapSizeKb}`,
       );
       const db = await openVersionedStoreAt(
         newEpochDirectory,
         SingleEpochDatabase.SCHEMA_VERSION,
         this.config.l1Contracts.rollupAddress,
-        this.config.dataStoreMapSizeKB,
+        this.config.dataStoreMapSizeKb,
       );
       epochDb = new SingleEpochDatabase(db);
       this.epochs.set(epochNumber, epochDb);