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

package/src/proving_broker/proving_broker_instrumentation.ts

@@ -8,7 +8,7 @@ import {
   type ObservableResult,
   type TelemetryClient,
   type UpDownCounter,
-  ValueType,
+  createUpDownCounterWithDefault,
 } from '@aztec/telemetry-client';
 
 export type MonitorCallback = (proofType: ProvingRequestType) => number;
@@ -28,49 +28,28 @@ export class ProvingBrokerInstrumentation {
   constructor(client: TelemetryClient, name = 'ProvingBroker') {
     const meter = client.getMeter(name);
 
-    this.queueSize = meter.createObservableGauge(Metrics.PROVING_QUEUE_SIZE, {
-      valueType: ValueType.INT,
-    });
+    this.queueSize = meter.createObservableGauge(Metrics.PROVING_QUEUE_SIZE);
 
-    this.activeJobs = meter.createObservableGauge(Metrics.PROVING_QUEUE_ACTIVE_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.activeJobs = meter.createObservableGauge(Metrics.PROVING_QUEUE_ACTIVE_JOBS);
 
-    this.resolvedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_RESOLVED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    const provingJobTypes = Object.values(ProvingRequestType).filter(v => typeof v === 'string');
+    const provingJobAttrs = { [Attributes.PROVING_JOB_TYPE]: provingJobTypes };
 
-    this.rejectedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_REJECTED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.resolvedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_RESOLVED_JOBS, provingJobAttrs);
 
-    this.retriedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_RETRIED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.rejectedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_REJECTED_JOBS, provingJobAttrs);
 
-    this.timedOutJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_TIMED_OUT_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.retriedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_RETRIED_JOBS, provingJobAttrs);
 
-    this.cachedJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_CACHED_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.timedOutJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_TIMED_OUT_JOBS, provingJobAttrs);
 
-    this.totalJobs = meter.createUpDownCounter(Metrics.PROVING_QUEUE_TOTAL_JOBS, {
-      valueType: ValueType.INT,
-    });
+    this.cachedJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_CACHED_JOBS, provingJobAttrs);
 
-    this.jobWait = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_WAIT, {
-      description: 'Records how long a job sits in the queue',
-      unit: 'ms',
-      valueType: ValueType.INT,
-    });
+    this.totalJobs = createUpDownCounterWithDefault(meter, Metrics.PROVING_QUEUE_TOTAL_JOBS, provingJobAttrs);
 
-    this.jobDuration = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_DURATION, {
-      description: 'Records how long a job takes to complete',
-      unit: 'ms',
-      valueType: ValueType.INT,
-    });
+    this.jobWait = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_WAIT);
+
+    this.jobDuration = meter.createHistogram(Metrics.PROVING_QUEUE_JOB_DURATION);
   }
 
   monitorQueueDepth(fn: MonitorCallback) {

package/src/proving_broker/proving_job_controller.ts

@@ -1,4 +1,7 @@
-import { createLogger } from '@aztec/foundation/log';
+import { EpochNumber } from '@aztec/foundation/branded-types';
+import { randomBytes } from '@aztec/foundation/crypto/random';
+import { AbortError } from '@aztec/foundation/error';
+import { type Logger, type LoggerBindings, createLogger } from '@aztec/foundation/log';
 import type {
   ProvingJobId,
   ProvingJobInputs,
@@ -9,38 +12,35 @@ import { ProvingRequestType } from '@aztec/stdlib/proofs';
 
 export enum ProvingJobControllerStatus {
   IDLE = 'idle',
-  PROVING = 'proving',
+  RUNNING = 'running',
   DONE = 'done',
-  ABORTED = 'aborted',
-}
-
-interface ProvingJobCompletionCallback<T extends ProvingRequestType = ProvingRequestType> {
-  (
-    jobId: ProvingJobId,
-    type: T,
-    error: Error | undefined,
-    result: ProvingJobResultsMap[T] | undefined,
-  ): void | Promise<void>;
 }
 
 export class ProvingJobController {
   private status: ProvingJobControllerStatus = ProvingJobControllerStatus.IDLE;
   private promise?: Promise<void>;
   private abortController = new AbortController();
+  private result?: ProvingJobResultsMap[ProvingRequestType] | Error;
+  private log: Logger;
 
   constructor(
     private jobId: ProvingJobId,
     private inputs: ProvingJobInputs,
-    private epochNumber: number,
+    private epochNumber: EpochNumber,
     private startedAt: number,
     private circuitProver: ServerCircuitProver,
-    private onComplete: ProvingJobCompletionCallback,
-    private log = createLogger('prover-client:proving-agent:job-controller'),
-  ) {}
+    private onComplete: () => void,
+    bindings?: LoggerBindings,
+  ) {
+    this.log = createLogger('prover-client:proving-agent:job-controller', {
+      instanceId: randomBytes(4).toString('hex'),
+      ...bindings,
+    });
+  }
 
   public start(): void {
     if (this.status !== ProvingJobControllerStatus.IDLE) {
-      this.log.verbose(
+      this.log.warn(
         `Job controller for jobId=${this.jobId} not starting because it is not idle currentStatus=${this.status}`,
         {
           currentStatus: this.status,
@@ -50,63 +50,23 @@ export class ProvingJobController {
       return;
     }
 
-    this.status = ProvingJobControllerStatus.PROVING;
-    this.log.verbose(`Job controller started jobId=${this.jobId}`, {
+    this.promise = this.run();
+
+    this.log.info(`Job controller started jobId=${this.jobId}`, {
       jobId: this.jobId,
     });
-
-    this.promise = this.generateProof()
-      .then(
-        result => {
-          if (this.status === ProvingJobControllerStatus.ABORTED) {
-            this.log.warn(`Job controller for jobId=${this.jobId} completed successfully but job was aborted`, {
-              currentStatus: this.status,
-              jobId: this.jobId,
-            });
-            return;
-          }
-          this.status = ProvingJobControllerStatus.DONE;
-          this.log.verbose(`Job controller for jobId=${this.jobId} completed successfully`, {
-            jobId: this.jobId,
-          });
-          return this.onComplete(this.jobId, this.inputs.type, undefined, result);
-        },
-        error => {
-          if (this.status === ProvingJobControllerStatus.ABORTED) {
-            this.log.warn(`Job controller for jobId=${this.jobId} finished with an error but job was aborted`, {
-              currentStatus: this.status,
-              jobId: this.jobId,
-            });
-            return;
-          }
-
-          if (error.name === 'AbortError') {
-            // Ignore abort errors
-            return;
-          }
-
-          this.log.verbose(`Job controller for jobId=${this.jobId} finished with an error`, {
-            jobId: this.jobId,
-            err: error,
-          });
-
-          this.status = ProvingJobControllerStatus.DONE;
-          return this.onComplete(this.jobId, this.inputs.type, error, undefined);
-        },
-      )
-      .catch(err => {
-        this.log.error(`Job constroller failed to send result for jobId=${this.jobId}: ${err}`, err, {
-          jobId: this.jobId,
-        });
-      });
   }
 
   public getStatus(): ProvingJobControllerStatus {
     return this.status;
   }
 
+  public getResult(): ProvingJobResultsMap[ProvingRequestType] | Error | undefined {
+    return this.result;
+  }
+
   public abort(): void {
-    if (this.status !== ProvingJobControllerStatus.PROVING) {
+    if (this.status !== ProvingJobControllerStatus.RUNNING) {
       this.log.warn(`Tried to abort job controller for jobId=${this.jobId} but it is not running`, {
         currentStatus: this.status,
         jobId: this.jobId,
@@ -114,9 +74,8 @@ export class ProvingJobController {
       return;
     }
 
-    this.status = ProvingJobControllerStatus.ABORTED;
     this.abortController.abort();
-    this.log.verbose(`Aborted job controller for jobId=${this.jobId}`, {
+    this.log.warn(`Aborted job controller for jobId=${this.jobId}`, {
       jobId: this.jobId,
     });
   }
@@ -125,6 +84,10 @@ export class ProvingJobController {
     return this.jobId;
   }
 
+  public getProofType(): ProvingRequestType {
+    return this.inputs.type;
+  }
+
   public getStartedAt(): number {
     return this.startedAt;
   }
@@ -133,6 +96,36 @@ export class ProvingJobController {
     return ProvingRequestType[this.inputs.type];
   }
 
+  private run = async () => {
+    this.status = ProvingJobControllerStatus.RUNNING;
+    let result: ProvingJobResultsMap[ProvingRequestType] | Error;
+    try {
+      result = await this.generateProof();
+    } catch (err) {
+      if (err && err instanceof Error) {
+        result = err;
+      } else {
+        result = new Error('Unknown proving error: ' + String(err), { cause: err });
+      }
+    }
+
+    if (this.abortController.signal.aborted) {
+      this.log.warn(`Job controller for jobId=${this.jobId} completed but job was aborted`, {
+        currentStatus: this.status,
+        jobId: this.jobId,
+      });
+      result = new AbortError('Proof was aborted');
+    }
+
+    this.result = result;
+    this.status = ProvingJobControllerStatus.DONE;
+    try {
+      this.onComplete();
+    } catch (err) {
+      this.log.warn(`On complete handler error: ${err}`, { jobId: this.jobId });
+    }
+  };
+
   private async generateProof(): Promise<ProvingJobResultsMap[ProvingRequestType]> {
     const { type, inputs } = this.inputs;
     const signal = this.abortController.signal;
@@ -141,50 +134,74 @@ export class ProvingJobController {
         return await this.circuitProver.getAvmProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.PRIVATE_BASE_ROLLUP: {
-        return await this.circuitProver.getPrivateBaseRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.PUBLIC_CHONK_VERIFIER: {
+        return await this.circuitProver.getPublicChonkVerifierProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.PRIVATE_TX_BASE_ROLLUP: {
+        return await this.circuitProver.getPrivateTxBaseRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.PUBLIC_TX_BASE_ROLLUP: {
+        return await this.circuitProver.getPublicTxBaseRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.TX_MERGE_ROLLUP: {
+        return await this.circuitProver.getTxMergeRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.PUBLIC_BASE_ROLLUP: {
-        return await this.circuitProver.getPublicBaseRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.BLOCK_ROOT_FIRST_ROLLUP: {
+        return await this.circuitProver.getBlockRootFirstRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.MERGE_ROLLUP: {
-        return await this.circuitProver.getMergeRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.BLOCK_ROOT_SINGLE_TX_FIRST_ROLLUP: {
+        return await this.circuitProver.getBlockRootSingleTxFirstRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.EMPTY_BLOCK_ROOT_ROLLUP: {
-        return await this.circuitProver.getEmptyBlockRootRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.BLOCK_ROOT_EMPTY_TX_FIRST_ROLLUP: {
+        return await this.circuitProver.getBlockRootEmptyTxFirstRollupProof(inputs, signal, this.epochNumber);
       }
 
       case ProvingRequestType.BLOCK_ROOT_ROLLUP: {
         return await this.circuitProver.getBlockRootRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.SINGLE_TX_BLOCK_ROOT_ROLLUP: {
-        return await this.circuitProver.getSingleTxBlockRootRollupProof(inputs, signal, this.epochNumber);
+      case ProvingRequestType.BLOCK_ROOT_SINGLE_TX_ROLLUP: {
+        return await this.circuitProver.getBlockRootSingleTxRollupProof(inputs, signal, this.epochNumber);
       }
 
       case ProvingRequestType.BLOCK_MERGE_ROLLUP: {
         return await this.circuitProver.getBlockMergeRollupProof(inputs, signal, this.epochNumber);
       }
 
+      case ProvingRequestType.CHECKPOINT_ROOT_ROLLUP: {
+        return await this.circuitProver.getCheckpointRootRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.CHECKPOINT_ROOT_SINGLE_BLOCK_ROLLUP: {
+        return await this.circuitProver.getCheckpointRootSingleBlockRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.CHECKPOINT_PADDING_ROLLUP: {
+        return await this.circuitProver.getCheckpointPaddingRollupProof(inputs, signal, this.epochNumber);
+      }
+
+      case ProvingRequestType.CHECKPOINT_MERGE_ROLLUP: {
+        return await this.circuitProver.getCheckpointMergeRollupProof(inputs, signal, this.epochNumber);
+      }
+
       case ProvingRequestType.ROOT_ROLLUP: {
         return await this.circuitProver.getRootRollupProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.BASE_PARITY: {
+      case ProvingRequestType.PARITY_BASE: {
         return await this.circuitProver.getBaseParityProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.ROOT_PARITY: {
+      case ProvingRequestType.PARITY_ROOT: {
         return await this.circuitProver.getRootParityProof(inputs, signal, this.epochNumber);
       }
 
-      case ProvingRequestType.TUBE_PROOF: {
-        return await this.circuitProver.getTubeProof(inputs, signal, this.epochNumber);
-      }
-
       default: {
         const _exhaustive: never = type;
         return Promise.reject(new Error(`Invalid proof request type: ${type}`));

package/src/proving_broker/rpc.ts

@@ -1,5 +1,4 @@
 import { createSafeJsonRpcClient } from '@aztec/foundation/json-rpc/client';
-import type { SafeJsonRpcServer } from '@aztec/foundation/json-rpc/server';
 import {
   type GetProvingJobResponse,
   ProofUri,
@@ -13,7 +12,7 @@ import {
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
 import { type ApiSchemaFor, optional } from '@aztec/stdlib/schemas';
 import { type ComponentsVersions, getVersioningResponseHandler } from '@aztec/stdlib/versioning';
-import { createTracedJsonRpcServer, makeTracedFetch } from '@aztec/telemetry-client';
+import { makeTracedFetch } from '@aztec/telemetry-client';
 
 import { z } from 'zod';
 
@@ -54,10 +53,6 @@ export const ProvingJobBrokerSchema: ApiSchemaFor<ProvingJobBroker> = {
   ...ProvingJobProducerSchema,
 };
 
-export function createProvingBrokerServer(broker: ProvingJobBroker): SafeJsonRpcServer {
-  return createTracedJsonRpcServer(broker, ProvingJobBrokerSchema);
-}
-
 export function createProvingJobBrokerClient(
   url: string,
   versions: Partial<ComponentsVersions>,

package/src/test/mock_proof_store.ts

@@ -0,0 +1,14 @@
+import type { ProvingJobId } from '@aztec/stdlib/interfaces/server';
+import { ProvingRequestType } from '@aztec/stdlib/proofs';
+
+// Mock ProofStore for faster benchmarks with realistic Cloud Storage URIs
+export class MockProofStore {
+  private mockCounter = 0;
+  private readonly bucketName = 'aztec-proving-benchmarks';
+  private readonly basePath = 'proving-jobs';
+
+  saveProofInput(jobId: ProvingJobId, type: ProvingRequestType): Promise<string> {
+    const uri = `gs://${this.bucketName}/${this.basePath}/inputs/${type}/${jobId}`;
+    return Promise.resolve(uri as any);
+  }
+}