fix: restore metrics

AztecProtocol · Apr 18, 2024 · 9d880ed · 9d880ed
1 parent 11afa89
commit 9d880ed
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 9 deletions.
diff --git a/yarn-project/prover-client/src/orchestrator/orchestrator.ts b/yarn-project/prover-client/src/orchestrator/orchestrator.ts
@@ -13,6 +13,7 @@ import {
   type ProvingResult,
   type ProvingTicket,
 } from '@aztec/circuit-types/interfaces';
+import { type CircuitSimulationStats } from '@aztec/circuit-types/stats';
 import {
   type BaseOrMergeRollupPublicInputs,
   BaseParityInputs,
@@ -33,6 +34,7 @@ import { padArrayEnd } from '@aztec/foundation/collection';
 import { createDebugLogger } from '@aztec/foundation/log';
 import { promiseWithResolvers } from '@aztec/foundation/promise';
 import { type Tuple } from '@aztec/foundation/serialize';
+import { Timer } from '@aztec/foundation/timer';
 import { type MerkleTreeOperations } from '@aztec/world-state';
 
 import { inspect } from 'util';
@@ -312,7 +314,24 @@ export class ProvingOrchestrator {
     // and reject the proving job whilst keeping the event loop free of rejections
     const safeJob = async () => {
       try {
+        const timer = new Timer();
         const [publicInputs, proof] = await this.queue.prove(request);
+        const duration = timer.ms();
+
+        const inputSize = 'toBuffer' in request.inputs ? request.inputs.toBuffer().length : 0;
+        const outputSize = 'toBuffer' in publicInputs ? publicInputs.toBuffer().length : 0;
+        const circuitName = this.getCircuitNameFromRequest(request);
+        const stats: CircuitSimulationStats | undefined = circuitName
+          ? {
+              eventName: 'circuit-simulation',
+              circuitName,
+              duration,
+              inputSize,
+              outputSize,
+            }
+          : undefined;
+
+        logger.debug(`Simulated ${ProvingRequestType[request.type]} circuit duration=${duration}ms`, stats);
 
         if (!provingState?.verifyState()) {
           logger.debug(`State no longer valid, discarding result of job type ${ProvingRequestType[request.type]}`);
@@ -330,6 +349,43 @@ export class ProvingOrchestrator {
     setImmediate(safeJob);
   }
 
+  private getCircuitNameFromRequest(request: ProvingRequest): CircuitSimulationStats['circuitName'] | null {
+    switch (request.type) {
+      case ProvingRequestType.PUBLIC_VM:
+        return null;
+      case ProvingRequestType.PUBLIC_KERNEL_NON_TAIL:
+        switch (request.kernelType) {
+          case PublicKernelType.SETUP:
+            return 'public-kernel-setup';
+          case PublicKernelType.APP_LOGIC:
+            return 'public-kernel-app-logic';
+          case PublicKernelType.TEARDOWN:
+            return 'public-kernel-teardown';
+          default:
+            return null;
+        }
+      case ProvingRequestType.PUBLIC_KERNEL_TAIL:
+        switch (request.kernelType) {
+          case PublicKernelType.TAIL:
+            return 'public-kernel-tail';
+          default:
+            return null;
+        }
+      case ProvingRequestType.BASE_ROLLUP:
+        return 'base-rollup';
+      case ProvingRequestType.MERGE_ROLLUP:
+        return 'merge-rollup';
+      case ProvingRequestType.ROOT_ROLLUP:
+        return 'root-rollup';
+      case ProvingRequestType.BASE_PARITY:
+        return 'base-parity';
+      case ProvingRequestType.ROOT_PARITY:
+        return 'root-parity';
+      default:
+        return null;
+    }
+  }
+
   // Updates the merkle trees for a transaction. The first enqueued job for a transaction
   private async prepareBaseRollupInputs(
     provingState: ProvingState | undefined,

diff --git a/yarn-project/prover-client/src/orchestrator/tx-proving-state.ts b/yarn-project/prover-client/src/orchestrator/tx-proving-state.ts
@@ -44,11 +44,13 @@ export class TxProvingState {
         kernelRequest.type === PublicKernelType.TAIL
           ? {
               type: ProvingRequestType.PUBLIC_KERNEL_TAIL,
-              inputs: kernelRequest,
+              kernelType: kernelRequest.type,
+              inputs: kernelRequest.inputs,
             }
           : {
               type: ProvingRequestType.PUBLIC_KERNEL_NON_TAIL,
-              inputs: kernelRequest,
+              kernelType: kernelRequest.type,
+              inputs: kernelRequest.inputs,
             };
       const publicFunction: PublicFunction = {
         vmProof: undefined,

diff --git a/yarn-project/prover-client/src/prover-pool/circuit-prover-agent.ts b/yarn-project/prover-client/src/prover-pool/circuit-prover-agent.ts
@@ -53,18 +53,25 @@ export class CircuitProverAgent implements ProvingAgent {
     this.runningPromise = undefined;
   }
 
-  private work({ type, inputs }: ProvingRequest): Promise<ProvingRequestResult<typeof type>> {
+  private work(request: ProvingRequest): Promise<ProvingRequestResult<typeof type>> {
+    const { type, inputs } = request;
     switch (type) {
       case ProvingRequestType.PUBLIC_VM: {
         return Promise.resolve([{}, makeEmptyProof()] as const);
       }
 
       case ProvingRequestType.PUBLIC_KERNEL_NON_TAIL: {
-        return this.prover.getPublicKernelProof(inputs);
+        return this.prover.getPublicKernelProof({
+          type: request.kernelType,
+          inputs,
+        });
       }
 
       case ProvingRequestType.PUBLIC_KERNEL_TAIL: {
-        return this.prover.getPublicTailProof(inputs);
+        return this.prover.getPublicTailProof({
+          type: request.kernelType,
+          inputs,
+        });
       }
 
       case ProvingRequestType.BASE_ROLLUP: {

diff --git a/yarn-project/prover-client/src/prover-pool/proving-request.ts b/yarn-project/prover-client/src/prover-pool/proving-request.ts
@@ -30,15 +30,18 @@ export enum ProvingRequestType {
 export type ProvingRequest =
   | {
       type: ProvingRequestType.PUBLIC_VM;
-      inputs: unknown;
+      // prefer object over unknown so that we can run "in" checks, e.g. `'toBuffer' in request.inputs`
+      inputs: object;
     }
   | {
       type: ProvingRequestType.PUBLIC_KERNEL_NON_TAIL;
-      inputs: PublicKernelNonTailRequest;
+      kernelType: PublicKernelNonTailRequest['type'];
+      inputs: PublicKernelNonTailRequest['inputs'];
     }
   | {
       type: ProvingRequestType.PUBLIC_KERNEL_TAIL;
-      inputs: PublicKernelTailRequest;
+      kernelType: PublicKernelTailRequest['type'];
+      inputs: PublicKernelTailRequest['inputs'];
     }
   | {
       type: ProvingRequestType.BASE_PARITY;
@@ -62,7 +65,7 @@ export type ProvingRequest =
     };
 
 export type ProvingRequestPublicInputs = {
-  [ProvingRequestType.PUBLIC_VM]: unknown;
+  [ProvingRequestType.PUBLIC_VM]: object;
 
   [ProvingRequestType.PUBLIC_KERNEL_NON_TAIL]: PublicKernelCircuitPublicInputs;
   [ProvingRequestType.PUBLIC_KERNEL_TAIL]: KernelCircuitPublicInputs;