@aztec/prover-client 0.69.1 → 0.71.0

package/src/proving_broker/factory.ts

@@ -1,7 +1,6 @@
-import { type ProverBrokerConfig } from '@aztec/circuit-types';
-import { AztecLmdbStore } from '@aztec/kv-store/lmdb';
 import { type TelemetryClient } from '@aztec/telemetry-client';
 
+import { type ProverBrokerConfig } from './config.js';
 import { ProvingBroker } from './proving_broker.js';
 import { InMemoryBrokerDatabase } from './proving_broker_database/memory.js';
 import { KVBrokerDatabase } from './proving_broker_database/persisted.js';
@@ -10,15 +9,17 @@ export async function createAndStartProvingBroker(
   config: ProverBrokerConfig,
   client: TelemetryClient,
 ): Promise<ProvingBroker> {
-  const database = config.proverBrokerDataDirectory
-    ? new KVBrokerDatabase(AztecLmdbStore.open(config.proverBrokerDataDirectory), client)
-    : new InMemoryBrokerDatabase();
+  const database = config.dataDirectory ? await KVBrokerDatabase.new(config, client) : new InMemoryBrokerDatabase();
 
-  const broker = new ProvingBroker(database, client, {
-    jobTimeoutMs: config.proverBrokerJobTimeoutMs,
-    maxRetries: config.proverBrokerJobMaxRetries,
-    timeoutIntervalMs: config.proverBrokerPollIntervalMs,
-  });
+  const broker = new ProvingBroker(
+    database,
+    {
+      jobTimeoutMs: config.proverBrokerJobTimeoutMs,
+      maxRetries: config.proverBrokerJobMaxRetries,
+      timeoutIntervalMs: config.proverBrokerPollIntervalMs,
+    },
+    client,
+  );
 
   await broker.start();
   return broker;

package/src/proving_broker/fixtures.ts

@@ -0,0 +1,14 @@
+import { type ProofUri, type ProvingJobId, ProvingRequestType, makeProvingJobId } from '@aztec/circuit-types';
+import { randomBytes } from '@aztec/foundation/crypto';
+
+export function makeRandomProvingJobId(epochNumber?: number): ProvingJobId {
+  return makeProvingJobId(epochNumber ?? 1, ProvingRequestType.BASE_PARITY, randomBytes(8).toString('hex'));
+}
+
+export function makeInputsUri(): ProofUri {
+  return randomBytes(8).toString('hex') as ProofUri;
+}
+
+export function makeOutputsUri(): ProofUri {
+  return randomBytes(8).toString('hex') as ProofUri;
+}

package/src/proving_broker/index.ts

@@ -4,5 +4,6 @@ export * from './rpc.js';
 export * from './proving_broker_database.js';
 export * from './proving_broker_database/memory.js';
 export * from './proving_broker_database/persisted.js';
-export * from './proof_store.js';
+export * from './proof_store/index.js';
 export * from './factory.js';
+export * from './config.js';

package/src/proving_broker/proof_store/factory.ts

@@ -0,0 +1,42 @@
+import { createLogger } from '@aztec/foundation/log';
+
+import { GoogleCloudStorageProofStore } from './gcs_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) {
+    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://'.`);
+  }
+}

package/src/proving_broker/proof_store/gcs_proof_store.ts

@@ -0,0 +1,72 @@
+import {
+  type ProofUri,
+  type ProvingJobId,
+  type ProvingJobInputs,
+  type ProvingJobInputsMap,
+  type ProvingJobResult,
+  type ProvingJobResultsMap,
+  ProvingRequestType,
+  getProvingJobInputClassFor,
+} from '@aztec/circuit-types';
+
+import { Storage } from '@google-cloud/storage';
+import { join } from 'path';
+
+import { type ProofStore } from './proof_store.js';
+
+const INPUTS_PATH = 'inputs';
+
+export class GoogleCloudStorageProofStore implements ProofStore {
+  private readonly storage: Storage;
+
+  constructor(private readonly bucketName: string, private readonly path: string) {
+    this.storage = new Storage();
+  }
+
+  public async saveProofInput<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    inputs: ProvingJobInputsMap[T],
+  ): Promise<ProofUri> {
+    const path = join(this.path, INPUTS_PATH, ProvingRequestType[type], id);
+    const file = this.storage.bucket(this.bucketName).file(path);
+    await file.save(inputs.toBuffer());
+    return file.cloudStorageURI.toString() as ProofUri;
+  }
+
+  saveProofOutput<T extends ProvingRequestType>(
+    _id: ProvingJobId,
+    _type: T,
+    _result: ProvingJobResultsMap[T],
+  ): Promise<ProofUri> {
+    throw new Error('Not implemented');
+  }
+
+  public async getProofInput(uri: ProofUri): Promise<ProvingJobInputs> {
+    try {
+      const url = new URL(uri);
+      const bucket = this.storage.bucket(url.host);
+      const path = url.pathname.replace(/^\/+/, '');
+      const file = bucket.file(path);
+      if (!(await file.exists())) {
+        throw new Error(`File at ${uri} does not exist`);
+      }
+
+      const typeString = path.split('/').at(-2);
+      const type = typeString ? ProvingRequestType[typeString as keyof typeof ProvingRequestType] : undefined;
+      if (type === undefined) {
+        throw new Error(`Unrecognized proof type ${type} in path ${path}`);
+      }
+
+      const contents = await file.download();
+      const inputs = getProvingJobInputClassFor(type).fromBuffer(contents[0]);
+      return { inputs, type } as ProvingJobInputs;
+    } catch (err) {
+      throw new Error(`Error getting proof input at ${uri}: ${err}`);
+    }
+  }
+
+  getProofOutput(_uri: ProofUri): Promise<ProvingJobResult> {
+    throw new Error('Not implemented');
+  }
+}

package/src/proving_broker/proof_store/index.ts

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

package/src/proving_broker/{proof_store.ts → proof_store/inline_proof_store.ts}

@@ -10,50 +10,7 @@ import {
 import { jsonParseWithSchema, jsonStringify } from '@aztec/foundation/json-rpc';
 import { type ZodFor } from '@aztec/foundation/schemas';
 
-/**
- * A database for storing proof inputs and outputs.
- */
-export interface ProofStore {
-  /**
-   * Save a proof input to the database.
-   * @param jobId - The ID of the job the proof input is associated with.
-   * @param type - The type of the proving request.
-   * @param inputs - The proof input to save.
-   * @returns The URI of the saved proof input.
-   */
-  saveProofInput<T extends ProvingRequestType>(
-    jobId: ProvingJobId,
-    type: T,
-    inputs: ProvingJobInputsMap[T],
-  ): Promise<ProofUri>;
-
-  /**
-   * Save a proof output to the database.
-   * @param jobId - The ID of the job the proof input is associated with.
-   * @param type - The type of the proving request.
-   * @param result - The proof output to save.
-   * @returns The URI of the saved proof output.
-   */
-  saveProofOutput<T extends ProvingRequestType>(
-    id: ProvingJobId,
-    type: T,
-    result: ProvingJobResultsMap[T],
-  ): Promise<ProofUri>;
-
-  /**
-   * Retrieve a proof input from the database.
-   * @param uri - The URI of the proof input to retrieve.
-   * @returns The proof input.
-   */
-  getProofInput(uri: ProofUri): Promise<ProvingJobInputs>;
-
-  /**
-   * Retrieve a proof output from the database.
-   * @param uri - The URI of the proof output to retrieve.
-   * @returns The proof output.
-   */
-  getProofOutput(uri: ProofUri): Promise<ProvingJobResult>;
-}
+import { type ProofStore } from './proof_store.js';
 
 // use an ASCII encoded data uri https://datatracker.ietf.org/doc/html/rfc2397#section-2
 // we do this to avoid double encoding to base64 (since the inputs already serialize to a base64 string)

package/src/proving_broker/proof_store/proof_store.ts

@@ -0,0 +1,54 @@
+import {
+  type ProofUri,
+  type ProvingJobId,
+  type ProvingJobInputs,
+  type ProvingJobInputsMap,
+  type ProvingJobResult,
+  type ProvingJobResultsMap,
+  type ProvingRequestType,
+} from '@aztec/circuit-types';
+
+/**
+ * A database for storing proof inputs and outputs.
+ */
+export interface ProofStore {
+  /**
+   * Save a proof input to the database.
+   * @param jobId - The ID of the job the proof input is associated with.
+   * @param type - The type of the proving request.
+   * @param inputs - The proof input to save.
+   * @returns The URI of the saved proof input.
+   */
+  saveProofInput<T extends ProvingRequestType>(
+    jobId: ProvingJobId,
+    type: T,
+    inputs: ProvingJobInputsMap[T],
+  ): Promise<ProofUri>;
+
+  /**
+   * Save a proof output to the database.
+   * @param jobId - The ID of the job the proof input is associated with.
+   * @param type - The type of the proving request.
+   * @param result - The proof output to save.
+   * @returns The URI of the saved proof output.
+   */
+  saveProofOutput<T extends ProvingRequestType>(
+    id: ProvingJobId,
+    type: T,
+    result: ProvingJobResultsMap[T],
+  ): Promise<ProofUri>;
+
+  /**
+   * Retrieve a proof input from the database.
+   * @param uri - The URI of the proof input to retrieve.
+   * @returns The proof input.
+   */
+  getProofInput(uri: ProofUri): Promise<ProvingJobInputs>;
+
+  /**
+   * Retrieve a proof output from the database.
+   * @param uri - The URI of the proof output to retrieve.
+   * @returns The proof output.
+   */
+  getProofOutput(uri: ProofUri): Promise<ProvingJobResult>;
+}

package/src/proving_broker/proving_agent.ts

@@ -12,9 +12,15 @@ 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, type Traceable, type Tracer, trackSpan } from '@aztec/telemetry-client';
-
-import { type ProofStore } from './proof_store.js';
+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';
 
@@ -36,12 +42,12 @@ export class ProvingAgent implements Traceable {
     private proofStore: ProofStore,
     /** The prover implementation to defer jobs to */
     private circuitProver: ServerCircuitProver,
-    /** A telemetry client through which to emit metrics */
-    client: TelemetryClient,
     /** Optional list of allowed proof types to build */
     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'),
   ) {
     this.tracer = client.getTracer('ProvingAgent');