Merge pull request risc0#86 from taikoxyz/input-cache

feat: Input cache & remove unused components

host/config/config.json

@@ -1,12 +1,19 @@
 {
-    "bind": "0.0.0.0:8080",
-    "cache": "/tmp/sgx",
-    "log_path": "./log/raiko",
-    "concurrency_limit": 10,
-    "max_log_days": 7,
-    "log_level": "info",
     "network": "taiko_a7",
+    "prover": "0x70997970C51812dc3A010C7d01b50e0d17dc79C8",
+    "graffiti": "8008500000000000000000000000000000000000000000000000000000000000",
+    "proof_type": "sgx",
     "sgx": {
-        "instance_id": 0
+        "instance_id": 456,
+        "setup": false,
+        "bootstrap": false,
+        "prove": true,
+        "input_path": null
+    },
+    "risc0": {
+        "bonsai": true,
+        "snark": true,
+        "profile": false,
+        "execution_po2": 20
     }
 }

host/src/prover/execution.rs → host/src/execution.rs

@@ -14,21 +14,24 @@ use serde::{Deserialize, Serialize};
 use tracing::{info, warn};
 
 use super::error::Result;
-use crate::{
-    host::host::preflight,
-    memory,
-    prover::{error::HostError, request::ProofRequest},
-};
-
-pub async fn execute<D: Prover>(config: &serde_json::Value) -> Result<Proof> {
-    println!("- {:?}", config);
+use crate::{error::HostError, memory, preflight::preflight, request::ProofRequest};
 
+pub async fn execute<D: Prover>(
+    config: &serde_json::Value,
+    cached_input: Option<GuestInput>,
+) -> Result<(GuestInput, Proof)> {
     // Generate the input
-    memory::reset_stats();
-    let measurement = Measurement::start("Generating input...", false);
-    let input = prepare_input(config).await?;
-    measurement.stop_with("=> Input generated");
-    memory::print_stats("Input generation peak memory used: ");
+    let input = if let Some(cached_input) = cached_input {
+        println!("Using cached input");
+        cached_input
+    } else {
+        memory::reset_stats();
+        let measurement = Measurement::start("Generating input...", false);
+        let input = prepare_input(config).await?;
+        measurement.stop_with("=> Input generated");
+        memory::print_stats("Input generation peak memory used: ");
+        input
+    };
 
     // 2. Test run the block
     memory::reset_stats();
@@ -58,8 +61,9 @@ pub async fn execute<D: Prover>(config: &serde_json::Value) -> Result<Proof> {
     // Prove
     memory::reset_stats();
     let measurement = Measurement::start("Generating proof...", false);
-    let res = D::run(input, output, config)
+    let res = D::run(input.clone(), output, config)
         .await
+        .map(|proof| (input, proof))
         .map_err(|e| HostError::GuestError(e.to_string()));
     measurement.stop_with("=> Proof generated");
     memory::print_stats("Prover peak memory used: ");

host/src/main.rs

@@ -16,44 +16,30 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-pub mod host;
-mod metrics;
-mod prover;
+pub mod error;
+pub mod execution;
+pub mod preflight;
+pub mod provider_db;
+pub mod request;
+pub mod server;
 
 use std::{alloc, fmt::Debug, fs::File, io::BufReader, path::PathBuf};
 
 use anyhow::Result;
 use cap::Cap;
-use prover::server::serve;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
+use server::serve;
 use structopt::StructOpt;
+use tracing_appender::{
+    non_blocking::WorkerGuard,
+    rolling::{Builder, Rotation},
+};
+use tracing_subscriber::FmtSubscriber;
 
 #[global_allocator]
 static ALLOCATOR: Cap<alloc::System> = Cap::new(alloc::System, usize::max_value());
 
-mod memory {
-    use crate::ALLOCATOR;
-
-    pub(crate) fn reset_stats() {
-        ALLOCATOR.reset_stats();
-    }
-
-    pub(crate) fn get_max_allocated() -> usize {
-        ALLOCATOR.max_allocated()
-    }
-
-    pub(crate) fn print_stats(title: &str) {
-        let max_memory = get_max_allocated();
-        println!(
-            "{}{}.{} MB",
-            title,
-            max_memory / 1000000,
-            max_memory % 1000000
-        );
-    }
-}
-
 #[derive(StructOpt, Default, Clone, Serialize, Deserialize, Debug)]
 #[serde(default)]
 pub struct Opt {
@@ -62,51 +48,57 @@ pub struct Opt {
     /// [default: 0.0.0.0:8080]
     address: String,
 
-    #[structopt(long, require_equals = true, default_value = "/tmp")]
-    /// Use a local directory as a cache for RPC calls. Accepts a custom directory.
-    cache: PathBuf,
+    #[structopt(long, require_equals = true, default_value = "16")]
+    /// Limit the max number of in-flight requests
+    concurrency_limit: usize,
 
     #[structopt(long, require_equals = true)]
     log_path: Option<PathBuf>,
 
-    #[structopt(long, require_equals = true, default_value = "1000")]
-    proof_cache: usize,
-
-    #[structopt(long, require_equals = true, default_value = "10")]
-    concurrency_limit: usize,
-
-    #[structopt(long, require_equals = true, default_value = "taiko_a7")]
-    network: String,
-
     #[structopt(long, require_equals = true, default_value = "7")]
-    max_log_days: usize,
+    max_log: usize,
 
-    #[structopt(long, require_equals = true, default_value = "20")]
-    // WARNING: must be larger than concurrency_limit
-    max_caches: usize,
+    #[structopt(long, require_equals = true, default_value = "host/config/config.json")]
+    /// Path to a config file that includes sufficent json args to request
+    /// a proof of specified type. Curl json-rpc overrides its contents
+    config_path: PathBuf,
 
     #[structopt(long, require_equals = true)]
-    config_path: Option<PathBuf>,
+    /// Use a local directory as a cache for input. Accepts a custom directory.
+    cache_path: Option<PathBuf>,
 
     #[structopt(long, require_equals = true, env = "RUST_LOG", default_value = "info")]
+    /// Set the log level
     log_level: String,
 }
 
 #[tokio::main]
 async fn main() -> Result<()> {
+    let opt = Opt::from_args();
     let config = get_config(None).unwrap();
-    let opt = Opt::deserialize(&config).unwrap();
-    println!("Start config: {:?}", opt);
+    println!("Start config:\n{:#?}", config);
+    println!("Args:\n{:#?}", opt);
 
-    let subscriber_builder = tracing_subscriber::FmtSubscriber::builder()
-        .with_env_filter(&opt.log_level)
+    let _guard = subscribe_log(&opt.log_path, &opt.log_level, opt.max_log);
+
+    serve(opt).await?;
+    Ok(())
+}
+
+fn subscribe_log(
+    log_path: &Option<PathBuf>,
+    log_level: &String,
+    max_log: usize,
+) -> Option<WorkerGuard> {
+    let subscriber_builder = FmtSubscriber::builder()
+        .with_env_filter(log_level)
         .with_test_writer();
-    let _guard = match opt.log_path {
+    match log_path {
         Some(ref log_path) => {
-            let file_appender = tracing_appender::rolling::Builder::new()
-                .rotation(tracing_appender::rolling::Rotation::DAILY)
+            let file_appender = Builder::new()
+                .rotation(Rotation::DAILY)
                 .filename_prefix("raiko.log")
-                .max_log_files(opt.max_log_days)
+                .max_log_files(max_log)
                 .build(log_path)
                 .expect("initializing rolling file appender failed");
             let (non_blocking, _guard) = tracing_appender::non_blocking(file_appender);
@@ -119,24 +111,19 @@ async fn main() -> Result<()> {
             tracing::subscriber::set_global_default(subscriber).unwrap();
             None
         }
-    };
-
-    serve(opt).await?;
-    Ok(())
+    }
 }
 
 /// Gets the config going through all possible sources
-pub fn get_config(request_config: Option<Value>) -> Result<Value> {
+fn get_config(request_config: Option<Value>) -> Result<Value> {
     let mut config = Value::default();
     let opt = Opt::from_args();
 
     // Config file has the lowest preference
-    if let Some(config_path) = &opt.config_path {
-        let file = File::open(config_path)?;
-        let reader = BufReader::new(file);
-        let file_config: Value = serde_json::from_reader(reader)?;
-        merge(&mut config, &file_config);
-    };
+    let file = File::open(&opt.config_path)?;
+    let reader = BufReader::new(file);
+    let file_config: Value = serde_json::from_reader(reader)?;
+    merge(&mut config, &file_config);
 
     // Command line values have higher preference
     let cli_config = serde_json::to_value(&opt)?;
@@ -161,3 +148,25 @@ fn merge(a: &mut Value, b: &Value) {
         (a, b) => *a = b.clone(),
     }
 }
+
+mod memory {
+    use crate::ALLOCATOR;
+
+    pub(crate) fn reset_stats() {
+        ALLOCATOR.reset_stats();
+    }
+
+    pub(crate) fn get_max_allocated() -> usize {
+        ALLOCATOR.max_allocated()
+    }
+
+    pub(crate) fn print_stats(title: &str) {
+        let max_memory = get_max_allocated();
+        println!(
+            "{}{}.{} MB",
+            title,
+            max_memory / 1000000,
+            max_memory % 1000000
+        );
+    }
+}