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 run
+          cargo run --example rlc -- run -c sdk/data/rlc_config.json keygen
+          cargo run --example rlc -- run --input sdk/data/rlc_input.json run
+          cargo run --example account_age -- run -k 15 keygen
+          cargo run --example account_age -- run --input sdk/data/account_age_input.json run
+          cargo run --example quickstart -- run -k 15 keygen
+          cargo run --example quickstart -- run --input sdk/data/quickstart_input.json run

.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

@@ -8,4 +8,76 @@ See [./examples/account_age.rs](./examples/account_age.rs) for an example Axiom
 cargo run --example account_age -- --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`, `prove`, `keygen`, or `run`.
+
+
+## CLI 
+
+```Usage: account_age <COMMAND>
+
+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
+  run     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/bin.rs

@@ -0,0 +1,53 @@
+use clap::{Parser, Subcommand};
+
+use crate::{
+    cli::types::AxiomCircuitRunnerOptions, dockerize::DockerizeCmd,
+    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, $I);
+        #[tokio::main]
+        async fn main() {
+            env_logger::init();
+            let cli = <$crate::bin::Cli as clap::Parser>::parse();
+            match cli.command {
+                $crate::bin::Commands::Serve(args) => {
+                    let _ = server(args).await;
+                }
+                $crate::bin::Commands::Run(args) => {
+                    let thread = std::thread::spawn(|| {
+                        $crate::cli::run_cli_on_scaffold::<$A, $I>(args);
+                    });
+                    thread.join().unwrap();
+                }
+                $crate::bin::Commands::Dockerize(args) => {
+                    let env_args: Vec<String> = std::env::args().collect();
+                    $crate::dockerize::gen_dockerfile(env_args, args);
+                }
+            }
+        }
+    };
+}

sdk/src/cmd/mod.rs → sdk/src/cli/mod.rs

@@ -45,8 +45,9 @@ use crate::{
 pub fn run_cli_on_scaffold<
     A: AxiomCircuitScaffold<Http, Fr>,
     I: Into<A::InputValue> + DeserializeOwned,
->() {
-    let cli = AxiomCircuitRunnerOptions::parse();
+>(
+    cli: AxiomCircuitRunnerOptions,
+) {
     match cli.command {
         SnarkCmd::Mock | SnarkCmd::Run => {
             if cli.input_path.is_none() {
@@ -82,11 +83,9 @@ pub fn run_cli_on_scaffold<
         .provider
         .unwrap_or_else(|| env::var("PROVIDER_URI").expect("The `provider` argument is required for the selected command. Either pass it as an argument or set the `PROVIDER_URI` environment variable."));
     let provider = Provider::<Http>::try_from(provider_uri).unwrap();
-    let data_path = cli.data_path.unwrap_or_else(|| PathBuf::from("data"));
     let srs_path = cli
         .srs
         .unwrap_or_else(|| dirs::home_dir().unwrap().join(".axiom/srs/challenge_0085"));
-    let circuit_name = cli.name.unwrap_or_else(|| "circuit".to_string());
 
     let mut max_user_outputs = USER_MAX_OUTPUTS;
     let mut max_subqueries = USER_MAX_SUBQUERIES;
@@ -164,7 +163,7 @@ pub fn run_cli_on_scaffold<
         SnarkCmd::Keygen => {
             let srs = read_srs_from_dir_or_install(&srs_path, runner.k() as u32);
             let (vk, pk, pinning) = keygen(&mut runner, &srs);
-            write_keygen_output(&vk, &pk, &pinning, data_path.clone());
+            write_keygen_output(&vk, &pk, &pinning, cli.data_path.clone());
             let metadata = if should_aggregate {
                 if input.is_none() {
                     panic!("The `input` argument is required for keygen with aggregation.");
@@ -192,7 +191,7 @@ pub fn run_cli_on_scaffold<
                 );
                 let agg_params = agg_keygen_output.2.params;
                 let agg_vk = agg_keygen_output.0.clone();
-                write_agg_keygen_output(agg_keygen_output, data_path.clone());
+                write_agg_keygen_output(agg_keygen_output, cli.data_path.clone());
                 get_agg_axiom_client_circuit_metadata(
                     &runner,
                     &agg_kzg_params,
@@ -205,31 +204,34 @@ pub fn run_cli_on_scaffold<
             };
             write_metadata(
                 metadata,
-                data_path.join(PathBuf::from(format!("{}.json", circuit_name))),
+                cli.data_path
+                    .join(PathBuf::from(format!("{}.json", cli.name))),
             );
         }
         SnarkCmd::Run => {
-            let metadata =
-                read_metadata(data_path.join(PathBuf::from(format!("{}.json", circuit_name))));
+            let metadata = read_metadata(
+                cli.data_path
+                    .join(PathBuf::from(format!("{}.json", cli.name))),
+            );
             let circuit_id = metadata.circuit_id.clone();
-            let (pk, pinning) = read_pk_and_pinning(data_path.clone(), circuit_id, &runner);
+            let (pk, pinning) = read_pk_and_pinning(cli.data_path.clone(), circuit_id, &runner);
             let prover =
                 AxiomCircuit::<Fr, Http, A>::prover(provider, pinning.clone()).use_inputs(input);
             let srs = read_srs_from_dir_or_install(&srs_path, prover.k() as u32);
             let inner_output = run(prover, &pk, &srs);
             let output = if should_aggregate {
                 let agg_circuit_id = metadata.agg_circuit_id.expect("No aggregation circuit ID");
                 let (agg_pk, agg_pinning) =
-                    read_agg_pk_and_pinning::<A::CoreParams>(data_path.clone(), agg_circuit_id);
+                    read_agg_pk_and_pinning::<A::CoreParams>(cli.data_path.clone(), agg_circuit_id);
                 let agg_srs = read_srs_from_dir_or_install(&srs_path, agg_pinning.params.degree);
-                agg_circuit_run(agg_pinning, inner_output, agg_pk, &agg_srs)
+                agg_circuit_run(agg_pinning, inner_output, &agg_pk, &agg_srs)
             } else {
                 inner_output
             };
             write_output(
                 output,
-                data_path.join(PathBuf::from("output.snark")),
-                data_path.join(PathBuf::from("output.json")),
+                cli.data_path.join(PathBuf::from("output.snark")),
+                cli.data_path.join(PathBuf::from("output.json")),
             );
         }
     }
@@ -242,5 +244,6 @@ where
     A::Input<AssignedValue<Fr>>: Debug,
 {
     env_logger::init();
-    run_cli_on_scaffold::<AxiomCompute<A>, A::LogicInput>();
+    let cli = AxiomCircuitRunnerOptions::parse();
+    run_cli_on_scaffold::<AxiomCompute<A>, A::LogicInput>(cli);
 }

sdk/src/cli/readme.md

@@ -0,0 +1,26 @@
+## Usage
+
+```
+Run keygen and real/mock proving
+
+Usage: <BIN NAME> [OPTIONS] <COMMAND>
+
+Commands:
+  mock    Run the mock prover
+  keygen  Generate new proving & verifying keys
+  run     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
+```