@aztec/prover-client 0.67.0 → 0.67.1-devnet

package/src/proving_broker/{caching_broker_facade.ts → broker_prover_facade.ts}

@@ -1,6 +1,5 @@
 import {
   type ProofAndVerificationKey,
-  type ProverCache,
   type ProvingJobId,
   type ProvingJobInputsMap,
   type ProvingJobProducer,
@@ -35,24 +34,20 @@ import {
 import { sha256 } from '@aztec/foundation/crypto';
 import { createLogger } from '@aztec/foundation/log';
 import { retryUntil } from '@aztec/foundation/retry';
+import { truncate } from '@aztec/foundation/string';
 
 import { InlineProofStore, type ProofStore } from './proof_store.js';
-import { InMemoryProverCache } from './prover_cache/memory.js';
 
 // 20 minutes, roughly the length of an Aztec epoch. If a proof isn't ready in this amount of time then we've failed to prove the whole epoch
 const MAX_WAIT_MS = 1_200_000;
 
-/**
- * A facade around a job broker that generates stable job ids and caches results
- */
-export class CachingBrokerFacade implements ServerCircuitProver {
+export class BrokerCircuitProverFacade implements ServerCircuitProver {
   constructor(
     private broker: ProvingJobProducer,
-    private cache: ProverCache = new InMemoryProverCache(),
     private proofStore: ProofStore = new InlineProofStore(),
     private waitTimeoutMs = MAX_WAIT_MS,
     private pollIntervalMs = 1000,
-    private log = createLogger('prover-client:caching-prover-broker'),
+    private log = createLogger('prover-client:broker-circuit-prover-facade'),
   ) {}
 
   private async enqueueAndWaitForJob<T extends ProvingRequestType>(
@@ -62,49 +57,23 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber = 0,
     signal?: AbortSignal,
   ): Promise<ProvingJobResultsMap[T]> {
-    // first try the cache
-    let jobEnqueued = false;
-    try {
-      const cachedResult = await this.cache.getProvingJobStatus(id);
-      if (cachedResult.status !== 'not-found') {
-        this.log.debug(`Found cached result for job=${id}: status=${cachedResult.status}`);
-      }
-
-      if (cachedResult.status === 'fulfilled') {
-        const output = await this.proofStore.getProofOutput(cachedResult.value);
-        if (output.type === type) {
-          return output.result as ProvingJobResultsMap[T];
-        } else {
-          this.log.warn(`Cached result type mismatch for job=${id}. Expected=${type} but got=${output.type}`);
-        }
-      } else if (cachedResult.status === 'rejected') {
-        // prefer returning a rejected promises so that we don't trigger the catch block below
-        return Promise.reject(new Error(cachedResult.reason));
-      } else if (cachedResult.status === 'in-progress' || cachedResult.status === 'in-queue') {
-        jobEnqueued = true;
-      } else {
-        jobEnqueued = false;
-      }
-    } catch (err) {
-      this.log.warn(`Failed to get cached proving job id=${id}: ${err}. Re-running job`);
-    }
+    const inputsUri = await this.proofStore.saveProofInput(id, type, inputs);
+    await this.broker.enqueueProvingJob({
+      id,
+      type,
+      inputsUri,
+      epochNumber,
+    });
 
-    if (!jobEnqueued) {
-      try {
-        const inputsUri = await this.proofStore.saveProofInput(id, type, inputs);
-        await this.broker.enqueueProvingJob({
-          id,
-          type,
-          inputsUri,
-          epochNumber,
-        });
-        await this.cache.setProvingJobStatus(id, { status: 'in-queue' });
-      } catch (err) {
-        this.log.error(`Failed to enqueue proving job id=${id}: ${err}`);
-        await this.cache.setProvingJobStatus(id, { status: 'not-found' });
-        throw err;
-      }
-    }
+    this.log.verbose(
+      `Sent proving job to broker id=${id} type=${ProvingRequestType[type]} epochNumber=${epochNumber}`,
+      {
+        provingJobId: id,
+        provingJobType: ProvingRequestType[type],
+        epochNumber,
+        inputsUri: truncate(inputsUri),
+      },
+    );
 
     // notify broker of cancelled job
     const abortFn = async () => {
@@ -131,26 +100,18 @@ export class CachingBrokerFacade implements ServerCircuitProver {
         this.pollIntervalMs / 1000,
       );
 
-      try {
-        await this.cache.setProvingJobStatus(id, result);
-      } catch (err) {
-        this.log.warn(`Failed to cache proving job id=${id} resultStatus=${result.status}: ${err}`);
-      }
-
       if (result.status === 'fulfilled') {
         const output = await this.proofStore.getProofOutput(result.value);
         if (output.type === type) {
           return output.result as ProvingJobResultsMap[T];
         } else {
-          return Promise.reject(new Error(`Unexpected proof type: ${output.type}. Expected: ${type}`));
+          throw new Error(`Unexpected proof type: ${output.type}. Expected: ${type}`);
         }
       } else {
-        return Promise.reject(new Error(result.reason));
+        throw new Error(result.reason);
       }
     } finally {
       signal?.removeEventListener('abort', abortFn);
-      // we've saved the result in our cache. We can tell the broker to clear its state
-      await this.broker.cleanUpProvingJobState(id);
     }
   }
 
@@ -160,7 +121,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<ProofAndVerificationKey<typeof AVM_PROOF_LENGTH_IN_FIELDS>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.PUBLIC_VM, inputs),
+      this.generateId(ProvingRequestType.PUBLIC_VM, inputs, epochNumber),
       ProvingRequestType.PUBLIC_VM,
       inputs,
       epochNumber,
@@ -174,7 +135,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<ParityPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.BASE_PARITY, inputs),
+      this.generateId(ProvingRequestType.BASE_PARITY, inputs, epochNumber),
       ProvingRequestType.BASE_PARITY,
       inputs,
       epochNumber,
@@ -188,7 +149,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.BLOCK_MERGE_ROLLUP, input),
+      this.generateId(ProvingRequestType.BLOCK_MERGE_ROLLUP, input, epochNumber),
       ProvingRequestType.BLOCK_MERGE_ROLLUP,
       input,
       epochNumber,
@@ -202,7 +163,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.BLOCK_ROOT_ROLLUP, input),
+      this.generateId(ProvingRequestType.BLOCK_ROOT_ROLLUP, input, epochNumber),
       ProvingRequestType.BLOCK_ROOT_ROLLUP,
       input,
       epochNumber,
@@ -216,7 +177,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BlockRootOrBlockMergePublicInputs>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP, input),
+      this.generateId(ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP, input, epochNumber),
       ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP,
       input,
       epochNumber,
@@ -230,7 +191,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<KernelCircuitPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.PRIVATE_KERNEL_EMPTY, inputs),
+      this.generateId(ProvingRequestType.PRIVATE_KERNEL_EMPTY, inputs, epochNumber),
       ProvingRequestType.PRIVATE_KERNEL_EMPTY,
       inputs,
       epochNumber,
@@ -244,7 +205,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.MERGE_ROLLUP, input),
+      this.generateId(ProvingRequestType.MERGE_ROLLUP, input, epochNumber),
       ProvingRequestType.MERGE_ROLLUP,
       input,
       epochNumber,
@@ -257,7 +218,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.PRIVATE_BASE_ROLLUP, baseRollupInput),
+      this.generateId(ProvingRequestType.PRIVATE_BASE_ROLLUP, baseRollupInput, epochNumber),
       ProvingRequestType.PRIVATE_BASE_ROLLUP,
       baseRollupInput,
       epochNumber,
@@ -271,7 +232,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<BaseOrMergeRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.PUBLIC_BASE_ROLLUP, inputs),
+      this.generateId(ProvingRequestType.PUBLIC_BASE_ROLLUP, inputs, epochNumber),
       ProvingRequestType.PUBLIC_BASE_ROLLUP,
       inputs,
       epochNumber,
@@ -285,7 +246,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<ParityPublicInputs, typeof NESTED_RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.ROOT_PARITY, inputs),
+      this.generateId(ProvingRequestType.ROOT_PARITY, inputs, epochNumber),
       ProvingRequestType.ROOT_PARITY,
       inputs,
       epochNumber,
@@ -299,7 +260,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<PublicInputsAndRecursiveProof<RootRollupPublicInputs, typeof RECURSIVE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.ROOT_ROLLUP, input),
+      this.generateId(ProvingRequestType.ROOT_ROLLUP, input, epochNumber),
       ProvingRequestType.ROOT_ROLLUP,
       input,
       epochNumber,
@@ -313,7 +274,7 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     epochNumber?: number,
   ): Promise<ProofAndVerificationKey<typeof TUBE_PROOF_LENGTH>> {
     return this.enqueueAndWaitForJob(
-      this.generateId(ProvingRequestType.TUBE_PROOF, tubeInput),
+      this.generateId(ProvingRequestType.TUBE_PROOF, tubeInput, epochNumber),
       ProvingRequestType.TUBE_PROOF,
       tubeInput,
       epochNumber,
@@ -321,8 +282,8 @@ export class CachingBrokerFacade implements ServerCircuitProver {
     );
   }
 
-  private generateId(type: ProvingRequestType, inputs: { toBuffer(): Buffer }) {
+  private generateId(type: ProvingRequestType, inputs: { toBuffer(): Buffer }, epochNumber = 0) {
     const inputsHash = sha256(inputs.toBuffer());
-    return `${ProvingRequestType[type]}:${inputsHash.toString('hex')}`;
+    return `${epochNumber}:${ProvingRequestType[type]}:${inputsHash.toString('hex')}`;
   }
 }

package/src/proving_broker/proving_agent.ts

@@ -10,8 +10,9 @@ import {
 } from '@aztec/circuit-types';
 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 { type TelemetryClient } from '@aztec/telemetry-client';
+import { type TelemetryClient, type Traceable, type Tracer, trackSpan } from '@aztec/telemetry-client';
 
 import { type ProofStore } from './proof_store.js';
 import { ProvingAgentInstrumentation } from './proving_agent_instrumentation.js';
@@ -20,12 +21,14 @@ import { ProvingJobController, ProvingJobControllerStatus } from './proving_job_
 /**
  * A helper class that encapsulates a circuit prover and connects it to a job source.
  */
-export class ProvingAgent {
+export class ProvingAgent implements Traceable {
   private currentJobController?: ProvingJobController;
   private runningPromise: RunningPromise;
   private instrumentation: ProvingAgentInstrumentation;
   private idleTimer: Timer | undefined;
 
+  public readonly tracer: Tracer;
+
   constructor(
     /** The source of proving jobs */
     private broker: ProvingJobConsumer,
@@ -41,8 +44,9 @@ export class ProvingAgent {
     private pollIntervalMs = 1000,
     private log = createLogger('prover-client:proving-agent'),
   ) {
+    this.tracer = client.getTracer('ProvingAgent');
     this.instrumentation = new ProvingAgentInstrumentation(client);
-    this.runningPromise = new RunningPromise(this.safeWork, this.pollIntervalMs);
+    this.runningPromise = new RunningPromise(this.work.bind(this), this.log, this.pollIntervalMs);
   }
 
   public setCircuitProver(circuitProver: ServerCircuitProver): void {
@@ -63,76 +67,74 @@ export class ProvingAgent {
     await this.runningPromise.stop();
   }
 
-  private safeWork = async () => {
-    try {
-      // 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 },
-        );
-      } else {
-        maybeJob = await this.broker.getProvingJob({ allowList: this.proofAllowList });
-      }
-
-      if (!maybeJob) {
-        return;
-      }
-
-      let abortedProofJobId: string | undefined;
-      let abortedProofName: string | undefined;
-      if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
-        abortedProofJobId = this.currentJobController.getJobId();
-        abortedProofName = this.currentJobController.getProofTypeName();
-        this.currentJobController?.abort();
-      }
-
-      const { job, time } = maybeJob;
-      let inputs: ProvingJobInputs;
-      try {
-        inputs = await this.proofStore.getProofInput(job.inputsUri);
-      } catch (err) {
-        await this.broker.reportProvingJobError(job.id, 'Failed to load proof inputs', true);
-        return;
-      }
-
-      this.currentJobController = new ProvingJobController(
-        job.id,
-        inputs,
-        time,
-        this.circuitProver,
-        this.handleJobResult,
+  @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 },
       );
+    } else {
+      maybeJob = await this.broker.getProvingJob({ allowList: this.proofAllowList });
+    }
+
+    if (!maybeJob) {
+      return;
+    }
+
+    let abortedProofJobId: string | undefined;
+    let abortedProofName: string | undefined;
+    if (this.currentJobController?.getStatus() === ProvingJobControllerStatus.PROVING) {
+      abortedProofJobId = this.currentJobController.getJobId();
+      abortedProofName = this.currentJobController.getProofTypeName();
+      this.currentJobController?.abort();
+    }
 
-      if (abortedProofJobId) {
-        this.log.info(
-          `Aborting job id=${abortedProofJobId} type=${abortedProofName} to start new job id=${this.currentJobController.getJobId()} type=${this.currentJobController.getProofTypeName()} inputsUri=${truncateString(
-            job.inputsUri,
-          )}`,
-        );
-      } else {
-        this.log.info(
-          `Starting job id=${this.currentJobController.getJobId()} type=${this.currentJobController.getProofTypeName()} inputsUri=${truncateString(
-            job.inputsUri,
-          )}`,
-        );
-      }
-
-      if (this.idleTimer) {
-        this.instrumentation.recordIdleTime(this.idleTimer);
-      }
-      this.idleTimer = undefined;
-
-      this.currentJobController.start();
+    const { job, time } = maybeJob;
+    let inputs: ProvingJobInputs;
+    try {
+      inputs = await this.proofStore.getProofInput(job.inputsUri);
     } catch (err) {
-      this.log.error(`Error in ProvingAgent: ${String(err)}`);
+      await this.broker.reportProvingJobError(job.id, 'Failed to load proof inputs', true);
+      return;
     }
-  };
+
+    this.currentJobController = new ProvingJobController(
+      job.id,
+      inputs,
+      job.epochNumber,
+      time,
+      this.circuitProver,
+      this.handleJobResult,
+    );
+
+    if (abortedProofJobId) {
+      this.log.info(
+        `Aborting job id=${abortedProofJobId} type=${abortedProofName} to start new job id=${this.currentJobController.getJobId()} type=${this.currentJobController.getProofTypeName()} inputsUri=${truncate(
+          job.inputsUri,
+        )}`,
+      );
+    } else {
+      this.log.info(
+        `Starting job id=${this.currentJobController.getJobId()} type=${this.currentJobController.getProofTypeName()} inputsUri=${truncate(
+          job.inputsUri,
+        )}`,
+      );
+    }
+
+    if (this.idleTimer) {
+      this.instrumentation.recordIdleTime(this.idleTimer);
+    }
+    this.idleTimer = undefined;
+
+    this.currentJobController.start();
+  }
 
   handleJobResult = async <T extends ProvingRequestType>(
     jobId: ProvingJobId,
@@ -147,14 +149,8 @@ export class ProvingAgent {
       return this.broker.reportProvingJobError(jobId, err.message, retry);
     } else if (result) {
       const outputUri = await this.proofStore.saveProofOutput(jobId, type, result);
-      this.log.info(
-        `Job id=${jobId} type=${ProvingRequestType[type]} completed outputUri=${truncateString(outputUri)}`,
-      );
+      this.log.info(`Job id=${jobId} type=${ProvingRequestType[type]} completed outputUri=${truncate(outputUri)}`);
       return this.broker.reportProvingJobSuccess(jobId, outputUri);
     }
   };
 }
-
-function truncateString(str: string, length: number = 64): string {
-  return str.length > length ? str.slice(0, length) + '...' : str;
-}