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

package/src/proving_broker/factory.ts

@@ -6,9 +6,10 @@ import { InMemoryBrokerDatabase } from './proving_broker_database/memory.js';
 import { KVBrokerDatabase } from './proving_broker_database/persisted.js';
 
 export async function createAndStartProvingBroker(
-  config: ProverBrokerConfig,
+  _config: ProverBrokerConfig,
   client: TelemetryClient,
 ): Promise<ProvingBroker> {
+  const config = { ..._config, dataStoreMapSizeKb: _config.proverBrokerStoreMapSizeKb ?? _config.dataStoreMapSizeKb };
   const database = config.dataDirectory ? await KVBrokerDatabase.new(config, client) : new InMemoryBrokerDatabase();
 
   const broker = new ProvingBroker(database, config, client);

package/src/proving_broker/fixtures.ts

@@ -1,9 +1,14 @@
-import { randomBytes } from '@aztec/foundation/crypto';
+import { EpochNumber } from '@aztec/foundation/branded-types';
+import { randomBytes } from '@aztec/foundation/crypto/random';
 import { type ProofUri, type ProvingJobId, makeProvingJobId } from '@aztec/stdlib/interfaces/server';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
 
-export function makeRandomProvingJobId(epochNumber?: number): ProvingJobId {
-  return makeProvingJobId(epochNumber ?? 1, ProvingRequestType.BASE_PARITY, randomBytes(8).toString('hex'));
+export function makeRandomProvingJobId(epochNumber?: EpochNumber): ProvingJobId {
+  return makeProvingJobId(
+    epochNumber ?? EpochNumber(1),
+    ProvingRequestType.PARITY_BASE,
+    randomBytes(8).toString('hex'),
+  );
 }
 
 export function makeInputsUri(): ProofUri {

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/proof_store/factory.ts

@@ -1,42 +1,20 @@
 import { createLogger } from '@aztec/foundation/log';
+import { createFileStore } from '@aztec/stdlib/file-store';
 
-import { GoogleCloudStorageProofStore } from './gcs_proof_store.js';
+import { FileStoreProofStore } from './file_store_proof_store.js';
 import { InlineProofStore } from './inline_proof_store.js';
 import type { ProofStore } from './proof_store.js';
 
-export function createProofStore(config: string | undefined, logger = createLogger('prover-client:proof-store')) {
-  if (config === undefined) {
+export async function createProofStore(
+  config: string | undefined,
+  logger = createLogger('prover-client:proof-store'),
+): Promise<ProofStore> {
+  if (!config) {
     logger.info('Creating inline proof store');
     return new InlineProofStore();
-  } else if (config.startsWith('gs://')) {
-    try {
-      const url = new URL(config);
-      const bucket = url.host;
-      const path = url.pathname.replace(/^\/+/, '');
-      logger.info(`Creating google cloud proof store at ${bucket}`, { bucket, path });
-      return new GoogleCloudStorageProofStore(bucket, path);
-    } catch (err) {
-      throw new Error(
-        `Invalid google cloud proof store definition: '${config}'. Supported values are 'gs://bucket-name/path/to/store'.`,
-      );
-    }
-  } else {
-    throw new Error(`Unknown proof store config: '${config}'. Supported values are 'gs://bucket-name/path/to/store'.`);
   }
-}
 
-export function createProofStoreForUri(
-  uri: string,
-  logger = createLogger('prover-client:proof-store'),
-): Pick<ProofStore, 'getProofInput' | 'getProofOutput'> {
-  if (uri.startsWith('data://')) {
-    return createProofStore(undefined, logger);
-  } else if (uri.startsWith('gs://')) {
-    const url = new URL(uri);
-    const basePath = url.pathname.replace(/^\/+/, '').split('/').slice(0, -3);
-    url.pathname = basePath.join('/');
-    return createProofStore(uri, logger);
-  } else {
-    throw new Error(`Unknown proof store config: '${uri}'. Supported protocols are 'data://' and 'gs://'.`);
-  }
+  const fileStore = await createFileStore(config, logger);
+  logger.info(`Creating file store proof store at ${config}`);
+  return new FileStoreProofStore(fileStore);
 }

package/src/proving_broker/proof_store/file_store_proof_store.ts

@@ -0,0 +1,78 @@
+import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';
+import type { FileStore } from '@aztec/stdlib/file-store';
+import {
+  type ProofUri,
+  type ProvingJobId,
+  type ProvingJobInputs,
+  type ProvingJobInputsMap,
+  ProvingJobResult,
+  type ProvingJobResultsMap,
+  getProvingJobInputClassFor,
+} from '@aztec/stdlib/interfaces/server';
+import { ProvingRequestType } from '@aztec/stdlib/proofs';
+
+import type { ProofStore } from './proof_store.js';
+
+const INPUTS_PATH = 'inputs';
+const OUTPUTS_PATH = 'outputs';
+
+/**
+ * A proof store implementation backed by a generic FileStore.
+ * Supports any storage backend (GCS, S3, local filesystem) via the FileStore abstraction.
+ */
+export class FileStoreProofStore implements ProofStore {
+  constructor(private readonly fileStore: FileStore) {}
+
+  async saveProofInput<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    inputs: ProvingJobInputsMap[T],
+  ): Promise<ProofUri> {
+    const path = `${INPUTS_PATH}/${ProvingRequestType[type]}/${id}`;
+    const uri = await this.fileStore.save(path, inputs.toBuffer());
+    return uri as ProofUri;
+  }
+
+  async saveProofOutput<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    result: ProvingJobResultsMap[T],
+  ): Promise<ProofUri> {
+    const jobResult = { type, result } as ProvingJobResult;
+    const json = jsonStringify(jobResult);
+    const path = `${OUTPUTS_PATH}/${ProvingRequestType[type]}/${id}.json`;
+    const uri = await this.fileStore.save(path, Buffer.from(json, 'utf-8'));
+    return uri as ProofUri;
+  }
+
+  async getProofInput(uri: ProofUri): Promise<ProvingJobInputs> {
+    try {
+      const buffer = await this.fileStore.read(uri);
+      const type = this.extractTypeFromUri(uri);
+      const inputs = getProvingJobInputClassFor(type).fromBuffer(buffer);
+      return { inputs, type } as ProvingJobInputs;
+    } catch (err) {
+      throw new Error(`Error getting proof input at ${uri}: ${err}`);
+    }
+  }
+
+  async getProofOutput(uri: ProofUri): Promise<ProvingJobResult> {
+    try {
+      const buffer = await this.fileStore.read(uri);
+      return jsonParseWithSchema(buffer.toString('utf-8'), ProvingJobResult);
+    } catch (err) {
+      throw new Error(`Error getting proof output at ${uri}: ${err}`);
+    }
+  }
+
+  private extractTypeFromUri(uri: string): ProvingRequestType {
+    const url = new URL(uri);
+    const pathParts = url.pathname.split('/').filter(Boolean);
+    const typeString = pathParts.at(-2);
+    const type = typeString ? ProvingRequestType[typeString as keyof typeof ProvingRequestType] : undefined;
+    if (type === undefined) {
+      throw new Error(`Unrecognized proof type ${typeString} in URI ${uri}`);
+    }
+    return type;
+  }
+}

package/src/proving_broker/proof_store/index.ts

@@ -1,3 +1,4 @@
 export * from './proof_store.js';
 export * from './inline_proof_store.js';
 export * from './factory.js';
+export * from './file_store_proof_store.js';

package/src/proving_broker/proof_store/inline_proof_store.ts

@@ -52,7 +52,7 @@ export class InlineProofStore implements ProofStore {
     return (PREFIX + SEPARATOR + encoded) as ProofUri;
   }
 
-  private decode<T>(uri: ProofUri, schema: ZodFor<T>): Promise<T> {
+  private decode<T>(uri: ProofUri, schema: ZodFor<T>): T {
     const [prefix, data] = uri.split(SEPARATOR);
     if (prefix !== PREFIX) {
       throw new Error('Invalid proof input URI: ' + prefix);

package/src/proving_broker/proving_agent.ts

@@ -1,10 +1,11 @@
-import { createLogger } from '@aztec/foundation/log';
+import { AbortError } from '@aztec/foundation/error';
+import { type Logger, type LoggerBindings, 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,
@@ -12,28 +13,17 @@ import type {
   ServerCircuitProver,
 } from '@aztec/stdlib/interfaces/server';
 import { ProvingRequestType } from '@aztec/stdlib/proofs';
-import {
-  type TelemetryClient,
-  type Traceable,
-  type Tracer,
-  getTelemetryClient,
-  trackSpan,
-} from '@aztec/telemetry-client';
 
 import type { ProofStore } from './proof_store/index.js';
-import { ProvingAgentInstrumentation } from './proving_agent_instrumentation.js';
 import { ProvingJobController, ProvingJobControllerStatus } from './proving_job_controller.js';
 
 /**
  * A helper class that encapsulates a circuit prover and connects it to a job source.
  */
-export class ProvingAgent implements Traceable {
+export class ProvingAgent {
   private currentJobController?: ProvingJobController;
   private runningPromise: RunningPromise;
-  private instrumentation: ProvingAgentInstrumentation;
-  private idleTimer: Timer | undefined;
-
-  public readonly tracer: Tracer;
+  private log: Logger;
 
   constructor(
     /** The source of proving jobs */
@@ -46,12 +36,9 @@ export class ProvingAgent implements Traceable {
     private proofAllowList: Array<ProvingRequestType> = [],
     /** How long to wait between jobs */
     private pollIntervalMs = 1000,
-    /** A telemetry client through which to emit metrics */
-    client: TelemetryClient = getTelemetryClient(),
-    private log = createLogger('prover-client:proving-agent'),
+    bindings?: LoggerBindings,
   ) {
-    this.tracer = client.getTracer('ProvingAgent');
-    this.instrumentation = new ProvingAgentInstrumentation(client);
+    this.log = createLogger('prover-client:proving-agent', bindings);
     this.runningPromise = new RunningPromise(this.work.bind(this), this.log, this.pollIntervalMs);
   }
 
@@ -64,7 +51,6 @@ export class ProvingAgent implements Traceable {
   }
 
   public start(): void {
-    this.idleTimer = new Timer();
     this.runningPromise.start();
   }
 
@@ -73,41 +59,77 @@ export class ProvingAgent implements Traceable {
     await this.runningPromise.stop();
   }
 
-  @trackSpan('ProvingAgent.safeWork')
+  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' };
+  }
+
   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 +138,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 +156,12 @@ 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();
+      },
+      this.log.getBindings(),
     );
 
     if (abortedProofJobId) {
@@ -154,28 +181,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;
+  }
 }