feat: dockerized proving server + cli

.github/workflows/test.yml

@@ -39,11 +39,11 @@ jobs:
         run: |
           export PROVIDER_URI=${{ secrets.PROVIDER_URI_SEPOLIA }}
           mkdir data
-          cargo run --example keccak -- --input sdk/data/keccak_input.json --config sdk/data/keccak_config.json --aggregate --auto-config-aggregation keygen
-          cargo run --example keccak -- --input sdk/data/keccak_input.json --config sdk/data/keccak_config.json --aggregate run
-          cargo run --example rlc -- -c sdk/data/rlc_config.json keygen
-          cargo run --example rlc -- --input sdk/data/rlc_input.json run
-          cargo run --example account_age -- -k 15 keygen
-          cargo run --example account_age -- --input sdk/data/account_age_input.json run
-          cargo run --example quickstart -- -k 15 keygen
-          cargo run --example quickstart -- --input sdk/data/quickstart_input.json run
+          cargo run --example keccak -- run --input sdk/data/keccak_input.json --config sdk/data/keccak_config.json --aggregate --auto-config-aggregation keygen
+          cargo run --example keccak -- run --input sdk/data/keccak_input.json --config sdk/data/keccak_config.json --aggregate prove
+          cargo run --example rlc -- run -c sdk/data/rlc_config.json keygen
+          cargo run --example rlc -- run --input sdk/data/rlc_input.json prove
+          cargo run --example account_age -- run -k 15 keygen
+          cargo run --example account_age -- run --input sdk/data/account_age_input.json prove
+          cargo run --example quickstart -- run -k 15 keygen
+          cargo run --example quickstart -- run --input sdk/data/quickstart_input.json prove

.gitignore

@@ -9,3 +9,5 @@ Cargo.lock
 params
 .DS_Store
 debug/
+.cargo
+Dockerfile.cpu

circuit/src/run/aggregation.rs

@@ -75,7 +75,7 @@ pub fn agg_circuit_prove<CoreParams>(
 pub fn agg_circuit_run<CoreParams>(
     agg_circuit_pinning: AggregationCircuitPinning<CoreParams>,
     inner_output: AxiomV2CircuitOutput,
-    pk: ProvingKey<G1Affine>,
+    pk: &ProvingKey<G1Affine>,
     params: &ParamsKZG<Bn256>,
 ) -> AxiomV2CircuitOutput {
     let circuit = create_aggregation_circuit(
@@ -85,7 +85,7 @@ pub fn agg_circuit_run<CoreParams>(
     );
     let circuit = circuit.use_break_points(agg_circuit_pinning.break_points);
     let agg_circuit_params = circuit.builder.config_params.clone();
-    let agg_snark = gen_snark_shplonk(params, &pk, circuit, None::<&str>);
+    let agg_snark = gen_snark_shplonk(params, pk, circuit, None::<&str>);
     let compute_query = build_axiom_v2_compute_query(
         agg_snark.clone(),
         AxiomCircuitParams::Base(agg_circuit_params.clone()),

circuit/src/tests/keccak.rs

@@ -183,7 +183,7 @@ pub fn test_compute_query<S: AxiomCircuitScaffold<Http, Fr>>(_circuit: S) {
         &agg_kzg_params,
         false,
     );
-    let final_output = agg_circuit_run(agg_pinning, output.clone(), agg_pk, &agg_kzg_params);
+    let final_output = agg_circuit_run(agg_pinning, output.clone(), &agg_pk, &agg_kzg_params);
     let circuit = create_aggregation_circuit(
         agg_circuit_params,
         output.snark.clone(),

sdk/Cargo.toml

@@ -25,3 +25,4 @@ bincode = "1.3.3"
 rocket = { version = "0.5", features = ["json"] }
 dirs = "5.0.1"
 reqwest = {version = "0.12.3", features = ["blocking"]}
+tokio = "1.37.0"

sdk/examples/account_age.rs

@@ -2,7 +2,7 @@ use std::fmt::Debug;
 
 use axiom_sdk::{
     axiom::{AxiomAPI, AxiomComputeFn, AxiomComputeInput, AxiomResult},
-    cmd::run_cli,
+    axiom_main,
     ethers::types::Address,
     halo2_base::{
         gates::{GateInstructions, RangeInstructions},
@@ -48,8 +48,4 @@ impl AxiomComputeFn for AccountAgeInput {
     }
 }
 
-// axiom_compute_prover_server!(AccountAgeInput);
-
-fn main() {
-    run_cli::<AccountAgeInput>();
-}
+axiom_main!(AccountAgeInput);

sdk/examples/keccak.rs

@@ -2,7 +2,7 @@ use std::fmt::Debug;
 
 use axiom_sdk::{
     axiom::{AxiomAPI, AxiomComputeFn, AxiomComputeInput, AxiomResult},
-    cmd::run_cli,
+    axiom_main,
     halo2_base::AssignedValue,
     Fr,
 };
@@ -34,8 +34,4 @@ impl AxiomComputeFn for KeccakInput {
     }
 }
 
-fn main() {
-    run_cli::<KeccakInput>();
-}
-
-// axiom_compute_prover_server!(KeccakInput);
+axiom_main!(KeccakInput);

sdk/examples/quickstart.rs

@@ -2,7 +2,7 @@ use std::{fmt::Debug, str::FromStr};
 
 use axiom_sdk::{
     axiom::{AxiomAPI, AxiomComputeFn, AxiomComputeInput, AxiomResult},
-    cmd::run_cli,
+    axiom_main,
     ethers::types::{Address, H256},
     halo2_base::AssignedValue,
     subquery::{AccountField, HeaderField, TxField},
@@ -101,6 +101,4 @@ impl AxiomComputeFn for QuickstartInput {
     }
 }
 
-fn main() {
-    run_cli::<QuickstartInput>();
-}
+axiom_main!(QuickstartInput);

sdk/examples/rlc.rs

@@ -3,7 +3,7 @@ use std::fmt::Debug;
 use axiom_circuit::axiom_eth::rlc::circuit::builder::RlcCircuitBuilder;
 use axiom_sdk::{
     axiom::{AxiomAPI, AxiomComputeFn, AxiomComputeInput, AxiomResult},
-    cmd::run_cli,
+    axiom_main,
     halo2_base::{
         gates::{GateInstructions, RangeChip, RangeInstructions},
         AssignedValue,
@@ -53,6 +53,4 @@ impl AxiomComputeFn for RlcInput {
     }
 }
 
-fn main() {
-    run_cli::<RlcInput>();
-}
+axiom_main!(RlcInput);

sdk/readme.md

@@ -5,7 +5,79 @@
 See [./examples/account_age.rs](./examples/account_age.rs) for an example Axiom compute circuit. To run the `account_age` circuit:
 
 ```
-cargo run --example account_age -- --input data/account_age_input.json -k 12 -p <PROVIDER_URI> <CMD>
+cargo run --example account_age -- run --input data/account_age_input.json -k 12 -p <PROVIDER_URI> <CMD>
 ```
 
-where `PROVIDER_URI` is a JSON-RPC URL, and `CMD` is `mock`, `prove`, `keygen`, or `run`.
+where `PROVIDER_URI` is a JSON-RPC URL, and `CMD` is `mock`, `keygen`, or `prove`.
+
+
+## CLI 
+
+```Usage: account_age <COMMAND> <SUBCOMMAND ARGS> <SUBCOMMAND>
+
+Commands:
+  serve  Run a circuit proving server
+  run    Run keygen and real/mock proving
+  help   Print this message or the help of the given subcommand(s)
+
+Options:
+  -h, --help     Print help
+  -V, --version  Print version
+```
+
+### Run
+
+```
+Run keygen and real/mock proving
+
+Usage: account_age run [OPTIONS] <COMMAND>
+
+Commands:
+  mock    Run the mock prover
+  keygen  Generate new proving & verifying keys
+  prove     Generate an Axiom compute query
+  help    Print this message or the help of the given subcommand(s)
+
+Options:
+  -k, --degree <DEGREE>          To determine the size of your circuit (12..25)
+  -p, --provider <PROVIDER>      JSON RPC provider URI
+  -i, --input <INPUT_PATH>       JSON inputs to feed into your circuit
+  -n, --name <NAME>              Name of the output metadata file [default: circuit]
+  -d, --data-path <DATA_PATH>    For saving build artifacts [default: data]
+  -c, --config <CONFIG>          For specifying custom circuit parameters
+      --srs <SRS>                For specifying custom KZG params directory [default: params]
+      --aggregate                Whether to aggregate the output (defaults to false)
+      --auto-config-aggregation  Whether to aggregate the output (defaults to false)
+  -h, --help                     Print help
+  -V, --version                  Print version
+```
+
+
+### Serve
+
+```
+Run a circuit proving server
+
+Usage: account_age serve [OPTIONS]
+
+Options:
+  -d, --data-path <DATA_PATH>  For loading build artifacts [default: data]
+  -c, --name <CIRCUIT_NAME>    Name of the circuit metadata file [default: circuit]
+  -p, --provider <PROVIDER>    JSON RPC provider URI
+      --srs <SRS_PATH>         For specifying custom KZG params directory (defaults to `params`) [default: params]
+  -h, --help                   Print help
+  -V, --version                Print version
+```
+
+### Dockerize
+
+```
+To generate a Dockerfile for running the circuit binary
+
+Usage: account_age dockerize [OPTIONS]
+
+Options:
+  -o, --output-path <OUTPUT_PATH>  For loading build artifacts [default: Dockerfile.cpu]
+  -h, --help                       Print help
+  -V, --version                    Print version
+```

sdk/src/cli.rs

@@ -0,0 +1,53 @@
+use clap::{Parser, Subcommand};
+
+use crate::{
+    dockerize::DockerizeCmd, run::types::AxiomCircuitRunnerOptions,
+    server::types::AxiomComputeServerCmd,
+};
+
+#[derive(Parser)]
+#[command(author, version, about, long_about = None)]
+pub struct Cli {
+    #[command(subcommand)]
+    pub command: Commands,
+}
+
+#[derive(Subcommand)]
+pub enum Commands {
+    /// Run a circuit proving server
+    Serve(AxiomComputeServerCmd),
+    /// Run keygen and real/mock proving
+    Run(AxiomCircuitRunnerOptions),
+    /// To generate a Dockerfile for running the circuit binary
+    Dockerize(DockerizeCmd),
+}
+
+#[macro_export]
+macro_rules! axiom_main {
+    ($A:ty) => {
+        axiom_main!($crate::axiom::AxiomCompute<$A>, $A);
+    };
+    ($A:ty, $I: ty) => {
+        $crate::axiom_compute_prover_server!($A);
+        #[tokio::main]
+        async fn main() {
+            env_logger::init();
+            let cli = <$crate::cli::Cli as clap::Parser>::parse();
+            match cli.command {
+                $crate::cli::Commands::Serve(args) => {
+                    let _ = server(args).await;
+                }
+                $crate::cli::Commands::Run(args) => {
+                    let thread = std::thread::spawn(|| {
+                        $crate::run::run_cli_on_scaffold::<$A, $I>(args);
+                    });
+                    thread.join().unwrap();
+                }
+                $crate::cli::Commands::Dockerize(args) => {
+                    let env_args: Vec<String> = std::env::args().collect();
+                    $crate::dockerize::gen_dockerfile(env_args, args);
+                }
+            }
+        }
+    };
+}

sdk/src/dockerize/Dockerfile.template.cpu

@@ -0,0 +1,27 @@
+FROM rust:1.74 as builder
+
+RUN apt-get update && apt-get install -y \
+    git \
+    openssh-client \
+    && rm -rf /var/lib/apt/lists/*
+
+RUN apt-get update -y
+RUN git config --global url."https://github.com/".insteadOf "[email protected]:"
+
+WORKDIR /code
+COPY . .
+
+ENV JEMALLOC_SYS_WITH_MALLOC_CONF="background_thread:true,metadata_thp:always,dirty_decay_ms:1000000,muzzy_decay_ms:1000000,abort_conf:true"
+ENV RUST_LOG=info
+
+RUN {{build_command}}
+
+FROM debian:stable-slim
+WORKDIR /code
+RUN apt update \
+    && apt install -y openssl ca-certificates \
+    && apt clean \
+    && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/*
+
+COPY --from=builder {{target_path}} ./bin
+EXPOSE 8000

sdk/src/dockerize/mod.rs

@@ -0,0 +1,57 @@
+use std::{fs::File, io::Write};
+
+use clap::{command, Parser};
+
+const TEMPLATE: &str = include_str!(concat!(
+    env!("CARGO_MANIFEST_DIR"),
+    "/src/dockerize/Dockerfile.template.cpu"
+));
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+pub struct DockerizeCmd {
+    #[arg(
+        short,
+        long = "output-path",
+        help = "For loading build artifacts",
+        default_value = "Dockerfile.cpu"
+    )]
+    /// The path to load build artifacts from
+    pub output_path: String,
+}
+
+pub fn gen_dockerfile(args: Vec<String>, cmd_args: DockerizeCmd) {
+    let bin_path = args[0].clone();
+    let bin_path_parts: Vec<&str> = bin_path.split(std::path::MAIN_SEPARATOR).collect();
+    assert_eq!(
+        bin_path_parts[0], "target",
+        "The dockerize command only supports the `target` build directory"
+    );
+    if !(bin_path_parts[1] == "debug" || bin_path_parts[1] == "release") {
+        panic!("The dockerize command only supports `debug` or `release` build profiles");
+    }
+    let is_example = bin_path_parts[2] == "examples";
+    let bin_name = if is_example {
+        bin_path_parts[3]
+    } else {
+        bin_path_parts[2]
+    };
+    let build_flag = if is_example { "--example" } else { "--bin" };
+    let build_string = format!("cargo build {} {} --release", build_flag, bin_name);
+    let target_path = if bin_path_parts[1] == "debug" {
+        bin_path.replace("debug", "release")
+    } else {
+        bin_path.clone()
+    };
+    let docker_target_path = format!("/code/{}", target_path);
+    let modified_template = TEMPLATE
+        .replace("{{build_command}}", &build_string)
+        .replace("{{target_path}}", &docker_target_path);
+
+    let mut file = File::create(&cmd_args.output_path).expect("Failed to create Dockerfile");
+    file.write_all(modified_template.as_bytes())
+        .expect("Failed to write Dockerfile");
+    log::info!("Dockerfile written to {}", &cmd_args.output_path);
+    log::info!("To build the Dockerfile, run:");
+    log::info!("docker build -f {} .", &cmd_args.output_path)
+}

sdk/src/dockerize/readme.md

@@ -0,0 +1,12 @@
+## Usage
+
+```
+To generate a Dockerfile for running the circuit binary
+
+Usage: <BIN NAME> dockerize [OPTIONS]
+
+Options:
+  -o, --output-path <OUTPUT_PATH>  For loading build artifacts [default: Dockerfile.cpu]
+  -h, --help                       Print help
+  -V, --version                    Print version
+```

sdk/src/lib.rs

@@ -116,6 +116,7 @@
 
 #![allow(incomplete_features)]
 #![feature(associated_type_defaults)]
+#![feature(async_fn_in_trait)]
 mod api;
 
 // pub(crate) mod utils;
@@ -134,13 +135,19 @@ pub mod axiom {
         compute::{AxiomCompute, AxiomComputeFn, AxiomComputeInput, AxiomResult},
     };
 }
-/// Contains a CLI for running any Axiom Compute function (any struct that implements the `AxiomComputeFn` trait)
-pub mod cmd;
 /// Contains the traits and types required to implement an Axiom Compute function (re-exported from the `axiom` module)
 pub(crate) mod compute;
+/// Contains a CLI for running any Axiom Compute function (any struct that implements the `AxiomComputeFn` trait)
+pub mod run;
+/// Contains a web server for running any Axiom Compute function (any struct that implements the `AxiomComputeFn` trait)
+pub mod server;
 /// Module with all subquery types and builders
 pub mod subquery;
 /// Re-export ethers-rs
 pub use ethers;
+/// Run either the proving/keygen CLI or the server from the same binary
+pub mod cli;
+/// Util to create a Docker image for a circuit
+pub mod dockerize;
 /// Module with utility functions for running the CLI
 pub mod utils;