diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 2936995..aac95de 100644 --- a/.github/workflows/test.yml +++ b/.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 diff --git a/.gitignore b/.gitignore index 0f4e008..baaf3d0 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,5 @@ Cargo.lock params .DS_Store debug/ +.cargo +Dockerfile.cpu diff --git a/circuit/src/run/aggregation.rs b/circuit/src/run/aggregation.rs index 54a3260..02fc8a5 100644 --- a/circuit/src/run/aggregation.rs +++ b/circuit/src/run/aggregation.rs @@ -75,7 +75,7 @@ pub fn agg_circuit_prove( pub fn agg_circuit_run( agg_circuit_pinning: AggregationCircuitPinning, inner_output: AxiomV2CircuitOutput, - pk: ProvingKey, + pk: &ProvingKey, params: &ParamsKZG, ) -> AxiomV2CircuitOutput { let circuit = create_aggregation_circuit( @@ -85,7 +85,7 @@ pub fn agg_circuit_run( ); 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()), diff --git a/circuit/src/tests/keccak.rs b/circuit/src/tests/keccak.rs index 60ebcdd..217ea5b 100644 --- a/circuit/src/tests/keccak.rs +++ b/circuit/src/tests/keccak.rs @@ -183,7 +183,7 @@ pub fn test_compute_query>(_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(), diff --git a/sdk/Cargo.toml b/sdk/Cargo.toml index 15b68eb..59529cb 100644 --- a/sdk/Cargo.toml +++ b/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" diff --git a/sdk/examples/account_age.rs b/sdk/examples/account_age.rs index a56bfd1..76c2d08 100644 --- a/sdk/examples/account_age.rs +++ b/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::(); -} +axiom_main!(AccountAgeInput); diff --git a/sdk/examples/keccak.rs b/sdk/examples/keccak.rs index 2063da8..03b7083 100644 --- a/sdk/examples/keccak.rs +++ b/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::(); -} - -// axiom_compute_prover_server!(KeccakInput); +axiom_main!(KeccakInput); diff --git a/sdk/examples/quickstart.rs b/sdk/examples/quickstart.rs index 0a5acd8..752f5b1 100644 --- a/sdk/examples/quickstart.rs +++ b/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::(); -} +axiom_main!(QuickstartInput); diff --git a/sdk/examples/rlc.rs b/sdk/examples/rlc.rs index cce80ec..5cb077a 100644 --- a/sdk/examples/rlc.rs +++ b/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::(); -} +axiom_main!(RlcInput); diff --git a/sdk/readme.md b/sdk/readme.md index 97c9d12..59005fc 100644 --- a/sdk/readme.md +++ b/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 +cargo run --example account_age -- run --input data/account_age_input.json -k 12 -p ``` -where `PROVIDER_URI` is a JSON-RPC URL, and `CMD` is `mock`, `prove`, `keygen`, or `run`. \ No newline at end of file +where `PROVIDER_URI` is a JSON-RPC URL, and `CMD` is `mock`, `keygen`, or `prove`. + + +## CLI + +```Usage: account_age + +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] + +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 To determine the size of your circuit (12..25) + -p, --provider JSON RPC provider URI + -i, --input JSON inputs to feed into your circuit + -n, --name Name of the output metadata file [default: circuit] + -d, --data-path For saving build artifacts [default: data] + -c, --config For specifying custom circuit parameters + --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 For loading build artifacts [default: data] + -c, --name Name of the circuit metadata file [default: circuit] + -p, --provider JSON RPC provider URI + --srs 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 For loading build artifacts [default: Dockerfile.cpu] + -h, --help Print help + -V, --version Print version +``` \ No newline at end of file diff --git a/sdk/src/cli.rs b/sdk/src/cli.rs new file mode 100644 index 0000000..368b39f --- /dev/null +++ b/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 = std::env::args().collect(); + $crate::dockerize::gen_dockerfile(env_args, args); + } + } + } + }; +} diff --git a/sdk/src/dockerize/Dockerfile.template.cpu b/sdk/src/dockerize/Dockerfile.template.cpu new file mode 100644 index 0000000..fd05599 --- /dev/null +++ b/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 "git@github.com:" + +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 diff --git a/sdk/src/dockerize/mod.rs b/sdk/src/dockerize/mod.rs new file mode 100644 index 0000000..784bbcf --- /dev/null +++ b/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, 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) +} diff --git a/sdk/src/dockerize/readme.md b/sdk/src/dockerize/readme.md new file mode 100644 index 0000000..044bb09 --- /dev/null +++ b/sdk/src/dockerize/readme.md @@ -0,0 +1,12 @@ +## Usage + +``` +To generate a Dockerfile for running the circuit binary + +Usage: dockerize [OPTIONS] + +Options: + -o, --output-path For loading build artifacts [default: Dockerfile.cpu] + -h, --help Print help + -V, --version Print version +``` \ No newline at end of file diff --git a/sdk/src/lib.rs b/sdk/src/lib.rs index b46f989..d113e98 100644 --- a/sdk/src/lib.rs +++ b/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; diff --git a/sdk/src/cmd/mod.rs b/sdk/src/run/mod.rs similarity index 88% rename from sdk/src/cmd/mod.rs rename to sdk/src/run/mod.rs index 97f0e4a..7b22c98 100644 --- a/sdk/src/cmd/mod.rs +++ b/sdk/src/run/mod.rs @@ -45,10 +45,11 @@ use crate::{ pub fn run_cli_on_scaffold< A: AxiomCircuitScaffold, I: Into + DeserializeOwned, ->() { - let cli = AxiomCircuitRunnerOptions::parse(); +>( + cli: AxiomCircuitRunnerOptions, +) { match cli.command { - SnarkCmd::Mock | SnarkCmd::Run => { + SnarkCmd::Mock | SnarkCmd::Prove => { if cli.input_path.is_none() { panic!("The `input_path` argument is required for the selected command."); } @@ -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::::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; @@ -131,7 +130,7 @@ pub fn run_cli_on_scaffold< AxiomCircuitParams::Base(base_params) } } else { - let degree = if cli.command == SnarkCmd::Run { + let degree = if cli.command == SnarkCmd::Prove { // The k will be read from the pinning file instead 12 } else { @@ -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,14 +204,17 @@ 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)))); + SnarkCmd::Prove => { + 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::::prover(provider, pinning.clone()).use_inputs(input); let srs = read_srs_from_dir_or_install(&srs_path, prover.k() as u32); @@ -220,16 +222,16 @@ pub fn run_cli_on_scaffold< 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::(data_path.clone(), agg_circuit_id); + read_agg_pk_and_pinning::(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>: Debug, { env_logger::init(); - run_cli_on_scaffold::, A::LogicInput>(); + let cli = AxiomCircuitRunnerOptions::parse(); + run_cli_on_scaffold::, A::LogicInput>(cli); } diff --git a/sdk/src/run/readme.md b/sdk/src/run/readme.md new file mode 100644 index 0000000..5b9931c --- /dev/null +++ b/sdk/src/run/readme.md @@ -0,0 +1,26 @@ +## Usage + +``` +Run keygen and real/mock proving + +Usage: [OPTIONS] + +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 To determine the size of your circuit (12..25) + -p, --provider JSON RPC provider URI + -i, --input JSON inputs to feed into your circuit + -n, --name Name of the output metadata file [default: circuit] + -d, --data-path For saving build artifacts [default: data] + -c, --config For specifying custom circuit parameters + --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 +``` \ No newline at end of file diff --git a/sdk/src/cmd/types.rs b/sdk/src/run/types.rs similarity index 86% rename from sdk/src/cmd/types.rs rename to sdk/src/run/types.rs index eea258b..d14be33 100644 --- a/sdk/src/cmd/types.rs +++ b/sdk/src/run/types.rs @@ -13,7 +13,7 @@ pub enum SnarkCmd { /// Generate new proving & verifying keys Keygen, /// Generate an Axiom compute query - Run, + Prove, } #[derive(Serialize, Deserialize, Clone, Debug)] @@ -56,31 +56,34 @@ pub struct AxiomCircuitRunnerOptions { /// The JSON inputs to feed into your circuit pub input_path: Option, - #[arg(short, long = "name", help = "Name of the output metadata file")] + #[arg( + short, + long = "name", + help = "Name of the output metadata file", + default_value = "circuit" + )] /// Name of the output metadata file - pub name: Option, + pub name: String, #[arg( short, long = "data-path", - help = "For saving build artifacts (optional)" + help = "For saving build artifacts", + default_value = "data" )] /// The path to save build artifacts - pub data_path: Option, + pub data_path: PathBuf, //Advanced options #[arg( short = 'c', long = "config", - help = "For specifying custom circuit parameters (optional)" + help = "For specifying custom circuit parameters" )] /// The path to a custom circuit configuration pub config: Option, - #[arg( - long = "srs", - help = "For specifying custom KZG params directory (defaults to `params`)" - )] + #[arg(long = "srs", help = "For specifying custom KZG params directory")] /// The path to the KZG params folder pub srs: Option, diff --git a/sdk/src/server/mod.rs b/sdk/src/server/mod.rs new file mode 100644 index 0000000..ac461fb --- /dev/null +++ b/sdk/src/server/mod.rs @@ -0,0 +1,317 @@ +use std::{ + collections::hash_map::DefaultHasher, + env, + fmt::Debug, + hash::{Hash, Hasher}, + path::PathBuf, +}; + +use axiom_circuit::{ + axiom_eth::{ + halo2_base::AssignedValue, halo2curves::bn256::Fr, + utils::build_utils::keygen::read_srs_from_dir, + }, + run::{aggregation::agg_circuit_run, inner::run}, + scaffold::{AxiomCircuit, AxiomCircuitScaffold}, + types::AxiomV2CircuitOutput, +}; +use ethers::providers::{Http, Provider}; +use rocket::State; +use serde::de::DeserializeOwned; +use tokio::runtime::Runtime; + +use self::types::{ + AggregationCircuitCtx, AxiomComputeCircuitCtx, AxiomComputeCtx, AxiomComputeJobStatus, + AxiomComputeManager, AxiomComputeServerCmd, +}; +use crate::{ + compute::{AxiomCompute, AxiomComputeFn}, + utils::io::{read_agg_pk_and_pinning, read_metadata, read_pinning, read_pk}, +}; + +pub mod types; + +pub trait AxiomClientProvingServer: AxiomCircuitScaffold { + type Context: Clone = (); + type ServerPayload = (); + type RequestInput: DeserializeOwned; + + fn construct_context(params: Self::CoreParams) -> Self::Context; + async fn process_input( + ctx: Self::Context, + input: Self::RequestInput, + provider: Provider, + ) -> Result<(Self::InputValue, Self::ServerPayload), String>; + #[allow(unused_variables)] + fn callback( + ctx: Self::Context, + payload: Self::ServerPayload, + output: AxiomV2CircuitOutput, + ) -> Result<(), ()> { + Ok(()) + } +} + +impl AxiomClientProvingServer for AxiomCompute +where + A::Input: Default + Debug, + A::Input>: Debug, +{ + type Context = (); + type RequestInput = A::LogicInput; + fn construct_context(_: A::CoreParams) -> Self::Context { + () + } + async fn process_input( + _ctx: Self::Context, + input: Self::RequestInput, + _provider: Provider, + ) -> Result<(A::Input, ()), String> { + Ok((input.into(), ())) + } +} + +pub async fn add_job(ctx: &State, job: String) -> u64 { + let mut hasher = DefaultHasher::new(); + job.hash(&mut hasher); + let job_id = hasher.finish(); + if ctx.inputs.lock().unwrap().contains_key(&job_id) { + return job_id; + } + ctx.job_queue.lock().unwrap().push(job_id); + ctx.inputs.lock().unwrap().insert(job_id, job); + ctx.job_status + .lock() + .unwrap() + .insert(job_id, AxiomComputeJobStatus::Received); + job_id +} + +pub fn prover_loop( + manager: AxiomComputeManager, + ctx: AxiomComputeCtx, + mut shutdown: tokio::sync::mpsc::Receiver<()>, +) { + loop { + let job = { + let mut queue = manager.job_queue.lock().unwrap(); + queue.pop() + }; + + if let Some(job) = job { + let inputs = { + let inputs_lock = manager.inputs.lock().unwrap(); + inputs_lock.clone() + }; + let raw_input = inputs.get(&job).unwrap(); + let input: A::RequestInput = serde_json::from_str(raw_input).unwrap(); + let rt = Runtime::new().unwrap(); + let processed_input = rt.block_on(async { + A::process_input(ctx.server_ctx.clone(), input, ctx.provider.clone()).await + }); + if processed_input.is_err() { + manager + .job_status + .lock() + .unwrap() + .insert(job, AxiomComputeJobStatus::ErrorInputPrep); + continue; + } + let (processed_input, payload) = processed_input.unwrap(); + let runner = AxiomCircuit::::prover( + ctx.provider.clone(), + ctx.child.pinning.clone(), + ) + .use_inputs(Some(processed_input)); + let scaffold_output = runner.scaffold_output(); + manager + .job_status + .lock() + .unwrap() + .insert(job, AxiomComputeJobStatus::DataQueryReady); + manager + .data_query + .lock() + .unwrap() + .insert(job, scaffold_output); + let inner_output = run(runner, &ctx.child.pk, &ctx.child.params); + manager + .job_status + .lock() + .unwrap() + .insert(job, AxiomComputeJobStatus::InnerOutputReady); + let output = if ctx.agg.is_some() { + let agg_ctx = ctx.agg.as_ref().unwrap(); + agg_circuit_run( + agg_ctx.pinning.clone(), + inner_output, + &agg_ctx.pk, + &agg_ctx.params, + ) + } else { + inner_output + }; + manager.outputs.lock().unwrap().insert(job, output.clone()); + manager + .job_status + .lock() + .unwrap() + .insert(job, AxiomComputeJobStatus::OutputReady); + let callback = A::callback(ctx.server_ctx.clone(), payload, output); + let callback_status = if callback.is_ok() { + AxiomComputeJobStatus::CallbackSuccess + } else { + AxiomComputeJobStatus::ErrorCallback + }; + manager + .job_status + .lock() + .unwrap() + .insert(job, callback_status); + } else { + if shutdown.try_recv().is_ok() { + break; + } + std::thread::sleep(std::time::Duration::from_secs(1)); + } + } +} + +pub fn initialize( + options: AxiomComputeServerCmd, +) -> AxiomComputeCtx { + let data_path = PathBuf::from(options.data_path); + let srs_path = PathBuf::from(options.srs_path); + let metadata = + read_metadata(data_path.join(PathBuf::from(format!("{}.json", options.circuit_name)))); + let provider_uri = options + .provider + .unwrap_or_else(|| env::var("PROVIDER_URI").expect("The `provider` argument is required. Either pass it as an argument or set the `PROVIDER_URI` environment variable.")); + let provider = Provider::::try_from(provider_uri).unwrap(); + let circuit_id = metadata.circuit_id.clone(); + let pinning = read_pinning(data_path.join(format!("{}.pinning", circuit_id))); + let runner = AxiomCircuit::::new(provider.clone(), pinning.clone().params) + .use_pinning(pinning.clone()); + let pk = read_pk(data_path.join(format!("{}.pk", circuit_id)), &runner); + let params = read_srs_from_dir(&srs_path, runner.k() as u32).expect("Unable to read SRS"); + + let agg = metadata.agg_circuit_id.map(|agg_circuit_id| { + let (agg_pk, agg_pinning) = read_agg_pk_and_pinning(data_path.clone(), agg_circuit_id); + let agg_params = + read_srs_from_dir(&srs_path, agg_pinning.params.degree).expect("Unable to read SRS"); + AggregationCircuitCtx { + pk: agg_pk, + pinning: agg_pinning, + params: agg_params, + } + }); + + let server_ctx = A::construct_context(pinning.clone().core_params); + + AxiomComputeCtx { + child: AxiomComputeCircuitCtx { + pk, + pinning, + params, + }, + agg, + provider, + server_ctx, + } +} +#[rocket::get("/job_status/")] +pub async fn get_job_status( + id: u64, + ctx: &rocket::State, +) -> (rocket::http::Status, String) { + let job_status = ctx.job_status.lock().unwrap(); + match job_status.get(&id) { + Some(status) => (rocket::http::Status::Ok, format!("{:?}", status)), + None => (rocket::http::Status::NotFound, "Job not found".to_string()), + } +} + +#[rocket::get("/data_query/")] +pub async fn get_data_query( + id: u64, + ctx: &rocket::State, +) -> (rocket::http::Status, String) { + let data_query = ctx.data_query.lock().unwrap(); + match data_query.get(&id) { + Some(query) => ( + rocket::http::Status::Ok, + serde_json::to_string(query).unwrap(), + ), + None => ( + rocket::http::Status::NotFound, + "Data query not found".to_string(), + ), + } +} + +#[rocket::get("/circuit_output/")] +pub async fn get_circuit_output( + id: u64, + ctx: &rocket::State, +) -> (rocket::http::Status, String) { + let outputs = ctx.outputs.lock().unwrap(); + match outputs.get(&id) { + Some(output) => ( + rocket::http::Status::Ok, + serde_json::to_string(output).unwrap(), + ), + None => ( + rocket::http::Status::NotFound, + "Circuit output not found".to_string(), + ), + } +} + +#[macro_export] +macro_rules! axiom_compute_prover_server { + ($A:ty) => { + #[rocket::post("/start_proving_job", format = "json", data = "")] + pub async fn start_proving_job( + req: rocket::serde::json::Json< + <$A as $crate::server::AxiomClientProvingServer>::RequestInput, + >, + ctx: &rocket::State<$crate::server::types::AxiomComputeManager>, + ) -> String { + let input = serde_json::to_string(&req.into_inner()).unwrap(); + let id = $crate::server::add_job(ctx, input).await; + id.to_string() + } + + async fn server( + options: $crate::server::types::AxiomComputeServerCmd, + ) -> Result<(), rocket::Error> { + let job_queue: $crate::server::types::AxiomComputeManager = Default::default(); + let (shutdown_tx, shutdown_rx) = tokio::sync::mpsc::channel(1); + let worker_manager = job_queue.clone(); + std::thread::spawn(|| { + let ctx = $crate::server::initialize::<$A>(options); + $crate::server::prover_loop::<$A>(worker_manager, ctx, shutdown_rx); + }); + rocket::build() + .configure(rocket::Config { + address: std::str::FromStr::from_str("0.0.0.0").unwrap(), + port: 8000, + ..rocket::Config::default() + }) + .manage(job_queue) + .mount( + "/", + rocket::routes![ + start_proving_job, + $crate::server::get_job_status, + $crate::server::get_data_query, + $crate::server::get_circuit_output + ], + ) + .launch() + .await?; + + Ok(()) + } + }; +} diff --git a/sdk/src/server/readme.md b/sdk/src/server/readme.md new file mode 100644 index 0000000..e6ffd05 --- /dev/null +++ b/sdk/src/server/readme.md @@ -0,0 +1,37 @@ +## Usage + +``` +Run a circuit proving server + +Usage: [OPTIONS] + +Options: + -d, --data-path For loading build artifacts [default: data] + -c, --name Name of the circuit metadata file [default: circuit] + -p, --provider JSON RPC provider URI + --srs For specifying custom KZG params directory (defaults to `params`) [default: params] + -h, --help Print help + -V, --version Print version +``` + +## Routes + +### Start Proving Job +- Type: POST +- Path: `/start_proving_job` +- Description: Accepts input data in JSON format to initiate a proving job and returns a job ID + +### Job Status +- Type: GET +- Path: `/job_status/` +- Description: Retrieves the status of a job by its ID + +### Data Query +- Type: GET +- Path: `/data_query/` +- Description: Fetches the data query associated with a job ID as soon as it is ready (before the full circuit output is ready) + +### Circuit Output +- Type: GET +- Path: `/circuit_output/` +- Description: Obtains the circuit output for a given job ID diff --git a/sdk/src/server/types.rs b/sdk/src/server/types.rs new file mode 100644 index 0000000..e3016cd --- /dev/null +++ b/sdk/src/server/types.rs @@ -0,0 +1,96 @@ +use std::{ + collections::HashMap, + sync::{Arc, Mutex}, +}; + +use axiom_circuit::{ + axiom_eth::{ + halo2_proofs::{plonk::ProvingKey, poly::kzg::commitment::ParamsKZG}, + halo2curves::bn256::{Bn256, G1Affine}, + }, + types::{ + AggregationCircuitPinning, AxiomCircuitPinning, AxiomV2CircuitOutput, AxiomV2DataAndResults, + }, +}; +use clap::Parser; +use ethers::providers::{Http, Provider}; +use serde::Serialize; + +use super::AxiomClientProvingServer; + +#[derive(Clone, Debug)] +pub struct AxiomComputeCircuitCtx { + pub pk: ProvingKey, + pub pinning: AxiomCircuitPinning, + pub params: ParamsKZG, +} + +#[derive(Clone, Debug)] +pub struct AggregationCircuitCtx { + pub pk: ProvingKey, + pub pinning: AggregationCircuitPinning, + pub params: ParamsKZG, +} + +#[derive(Clone, Debug)] +pub struct AxiomComputeCtx { + pub child: AxiomComputeCircuitCtx, + pub agg: Option>, + pub provider: Provider, + pub server_ctx: A::Context, +} + +#[derive(Clone, Debug, Serialize)] +pub enum AxiomComputeJobStatus { + Received, + DataQueryReady, + InnerOutputReady, + OutputReady, + CallbackSuccess, + ErrorInputPrep, + ErrorCallback, + ErrorInnerCircuit, +} + +#[derive(Clone, Debug, Default)] +pub struct AxiomComputeManager { + pub job_queue: Arc>>, + pub inputs: Arc>>, + pub job_status: Arc>>, + pub data_query: Arc>>, + pub outputs: Arc>>, +} + +#[derive(Parser, Debug)] +#[command(author, version, about, long_about = None)] +pub struct AxiomComputeServerCmd { + #[arg( + short, + long = "data-path", + help = "For loading build artifacts", + default_value = "data" + )] + /// The path to load build artifacts from + pub data_path: String, + + #[arg( + short, + long = "name", + help = "Name of the circuit metadata file", + default_value = "circuit" + )] + /// Name of the circuit metadata file + pub circuit_name: String, + + #[arg(short = 'p', long = "provider", help = "JSON RPC provider URI")] + /// The JSON RPC provider URI + pub provider: Option, + + #[arg( + long = "srs", + help = "For specifying custom KZG params directory (defaults to `params`)", + default_value = "params" + )] + /// The path to the KZG params folder + pub srs_path: String, +}