Merge branch 'refs/heads/main' into afo/prover-docker-setup

# Conflicts: # zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs # zk_toolbox/crates/zk_inception/src/commands/prover/init.rs
matter-labs · Sep 9, 2024 · 957cff9 · 957cff9
2 parents a7bba25 + 0a9e096
commit 957cff9
Show file tree

Hide file tree

Showing 14 changed files with 236 additions and 145 deletions.
diff --git a/core/bin/zksync_tee_prover/src/config.rs b/core/bin/zksync_tee_prover/src/config.rs
@@ -22,11 +22,21 @@ pub(crate) struct TeeProverConfig {
     pub max_retries: usize,
     /// Initial back-off interval when retrying recovery on a retriable error. Each subsequent retry interval
     /// will be multiplied by [`Self.retry_backoff_multiplier`].
-    pub initial_retry_backoff: Duration,
+    pub initial_retry_backoff_sec: u64,
     /// Multiplier for the back-off interval when retrying recovery on a retriable error.
     pub retry_backoff_multiplier: f32,
     /// Maximum back-off interval when retrying recovery on a retriable error.
-    pub max_backoff: Duration,
+    pub max_backoff_sec: u64,
+}
+
+impl TeeProverConfig {
+    pub fn initial_retry_backoff(&self) -> Duration {
+        Duration::from_secs(self.initial_retry_backoff_sec)
+    }
+
+    pub fn max_backoff(&self) -> Duration {
+        Duration::from_secs(self.max_backoff_sec)
+    }
 }
 
 impl FromEnv for TeeProverConfig {
@@ -39,9 +49,9 @@ impl FromEnv for TeeProverConfig {
     /// export TEE_PROVER_TEE_TYPE="sgx"
     /// export TEE_PROVER_API_URL="http://127.0.0.1:3320"
     /// export TEE_PROVER_MAX_RETRIES=10
-    /// export TEE_PROVER_INITIAL_RETRY_BACKOFF=1
+    /// export TEE_PROVER_INITIAL_RETRY_BACKOFF_SEC=1
     /// export TEE_PROVER_RETRY_BACKOFF_MULTIPLIER=2.0
-    /// export TEE_PROVER_MAX_BACKOFF=128
+    /// export TEE_PROVER_MAX_BACKOFF_SEC=128
     /// ```
     fn from_env() -> anyhow::Result<Self> {
         let config: Self = envy::prefixed("TEE_PROVER_").from_env()?;

diff --git a/core/bin/zksync_tee_prover/src/tee_prover.rs b/core/bin/zksync_tee_prover/src/tee_prover.rs
@@ -129,7 +129,7 @@ impl Task for TeeProver {
             .await?;
 
         let mut retries = 1;
-        let mut backoff = config.initial_retry_backoff;
+        let mut backoff = config.initial_retry_backoff();
         let mut observer = METRICS.job_waiting_time.start();
 
         loop {
@@ -141,7 +141,7 @@ impl Task for TeeProver {
             let need_to_sleep = match result {
                 Ok(batch_number) => {
                     retries = 1;
-                    backoff = config.initial_retry_backoff;
+                    backoff = config.initial_retry_backoff();
                     if let Some(batch_number) = batch_number {
                         observer.observe();
                         observer = METRICS.job_waiting_time.start();
@@ -162,7 +162,7 @@ impl Task for TeeProver {
                     retries += 1;
                     backoff = std::cmp::min(
                         backoff.mul_f32(config.retry_backoff_multiplier),
-                        config.max_backoff,
+                        config.max_backoff(),
                     );
                     true
                 }

diff --git a/core/node/state_keeper/src/keeper.rs b/core/node/state_keeper/src/keeper.rs
@@ -14,6 +14,7 @@ use zksync_multivm::{
     },
     utils::StorageWritesDeduplicator,
 };
+use zksync_shared_metrics::{TxStage, APP_METRICS};
 use zksync_state::{OwnedStorage, ReadStorageFactory};
 use zksync_types::{
     block::L2BlockExecutionData, l2::TransactionType, protocol_upgrade::ProtocolUpgradeTx,
@@ -463,6 +464,9 @@ impl ZkSyncStateKeeper {
                     .with_context(|| format!("failed re-executing transaction {:?}", tx.hash()))?;
                 let result = TxExecutionResult::new(result, &tx);
 
+                APP_METRICS.processed_txs[&TxStage::StateKeeper].inc();
+                APP_METRICS.processed_l1_txs[&TxStage::StateKeeper].inc_by(tx.is_l1().into());
+
                 let TxExecutionResult::Success {
                     tx_result,
                     tx_metrics,
@@ -742,6 +746,9 @@ impl ZkSyncStateKeeper {
         let exec_result = TxExecutionResult::new(exec_result, &tx);
         latency.observe();
 
+        APP_METRICS.processed_txs[&TxStage::StateKeeper].inc();
+        APP_METRICS.processed_l1_txs[&TxStage::StateKeeper].inc_by(tx.is_l1().into());
+
         let latency = KEEPER_METRICS.determine_seal_resolution.start();
         // All of `TxExecutionResult::BootloaderOutOfGasForTx`,
         // `Halt::NotEnoughGasProvided` correspond to out-of-gas errors but of different nature.

diff --git a/etc/nix/container-tee_prover.nix b/etc/nix/container-tee_prover.nix
@@ -33,9 +33,9 @@ nixsgxLib.mkSGXContainer {
       env = {
         TEE_PROVER_API_URL.passthrough = true;
         TEE_PROVER_MAX_RETRIES.passthrough = true;
-        TEE_PROVER_INITIAL_RETRY_BACKOFF_SECONDS.passthrough = true;
+        TEE_PROVER_INITIAL_RETRY_BACKOFF_SEC.passthrough = true;
         TEE_PROVER_RETRY_BACKOFF_MULTIPLIER.passthrough = true;
-        TEE_PROVER_MAX_BACKOFF_SECONDS.passthrough = true;
+        TEE_PROVER_MAX_BACKOFF_SEC.passthrough = true;
         API_PROMETHEUS_LISTENER_PORT.passthrough = true;
         API_PROMETHEUS_PUSHGATEWAY_URL.passthrough = true;
         API_PROMETHEUS_PUSH_INTERVAL_MS.passthrough = true;

diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/args/compressor_keys.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/args/compressor_keys.rs
@@ -0,0 +1,22 @@
+use clap::Parser;
+use common::Prompt;
+
+use crate::messages::MSG_SETUP_COMPRESSOR_KEY_PATH_PROMPT;
+
+#[derive(Debug, Clone, Parser, Default)]
+pub struct CompressorKeysArgs {
+    #[clap(long)]
+    pub path: Option<String>,
+}
+
+impl CompressorKeysArgs {
+    pub fn fill_values_with_prompt(self, default: &str) -> CompressorKeysArgs {
+        let path = self.path.unwrap_or_else(|| {
+            Prompt::new(MSG_SETUP_COMPRESSOR_KEY_PATH_PROMPT)
+                .default(default)
+                .ask()
+        });
+
+        CompressorKeysArgs { path: Some(path) }
+    }
+}
diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs
@@ -8,7 +8,10 @@ use url::Url;
 use xshell::Shell;
 use zksync_config::configs::fri_prover::CloudConnectionMode;
 
-use super::init_bellman_cuda::InitBellmanCudaArgs;
+use super::{
+    compressor_keys::CompressorKeysArgs, init_bellman_cuda::InitBellmanCudaArgs,
+    setup_keys::SetupKeysArgs,
+};
 use crate::{
     commands::prover::gcs::get_project_ids,
     consts::{DEFAULT_CREDENTIALS_FILE, DEFAULT_PROOF_STORE_DIR},
@@ -18,25 +21,24 @@ use crate::{
         MSG_CREATE_GCS_BUCKET_LOCATION_PROMPT, MSG_CREATE_GCS_BUCKET_NAME_PROMTP,
         MSG_CREATE_GCS_BUCKET_PROJECT_ID_NO_PROJECTS_PROMPT,
         MSG_CREATE_GCS_BUCKET_PROJECT_ID_PROMPT, MSG_CREATE_GCS_BUCKET_PROMPT,
-        MSG_DOWNLOAD_SETUP_KEY_PROMPT, MSG_GETTING_PROOF_STORE_CONFIG,
-        MSG_GETTING_PUBLIC_STORE_CONFIG, MSG_PROOF_STORE_CONFIG_PROMPT, MSG_PROOF_STORE_DIR_PROMPT,
+        MSG_DOWNLOAD_SETUP_COMPRESSOR_KEY_PROMPT, MSG_GETTING_PROOF_STORE_CONFIG,
+        MSG_GETTING_PUBLIC_STORE_CONFIG, MSG_INITIALIZE_BELLMAN_CUDA_PROMPT,
+        MSG_PROOF_STORE_CONFIG_PROMPT, MSG_PROOF_STORE_DIR_PROMPT,
         MSG_PROOF_STORE_GCS_BUCKET_BASE_URL_ERR, MSG_PROOF_STORE_GCS_BUCKET_BASE_URL_PROMPT,
         MSG_PROOF_STORE_GCS_CREDENTIALS_FILE_PROMPT, MSG_PROVER_DB_NAME_HELP,
-        MSG_PROVER_DB_URL_HELP, MSG_SAVE_TO_PUBLIC_BUCKET_PROMPT, MSG_SETUP_KEY_PATH_PROMPT,
+        MSG_PROVER_DB_URL_HELP, MSG_SAVE_TO_PUBLIC_BUCKET_PROMPT, MSG_SETUP_KEYS_PROMPT,
         MSG_USE_DEFAULT_DATABASES_HELP,
     },
 };
 
-#[derive(Debug, Clone, Serialize, Deserialize, Parser, Default)]
+#[derive(Debug, Clone, Parser, Default)]
 pub struct ProverInitArgs {
     // Proof store object
     #[clap(long)]
     pub proof_store_dir: Option<String>,
     #[clap(flatten)]
-    #[serde(flatten)]
     pub proof_store_gcs_config: ProofStorageGCSTmp,
     #[clap(flatten)]
-    #[serde(flatten)]
     pub create_gcs_bucket_config: ProofStorageGCSCreateBucketTmp,
 
     // Public store object
@@ -45,20 +47,25 @@ pub struct ProverInitArgs {
     #[clap(long)]
     pub public_store_dir: Option<String>,
     #[clap(flatten)]
-    #[serde(flatten)]
     pub public_store_gcs_config: PublicStorageGCSTmp,
     #[clap(flatten)]
-    #[serde(flatten)]
     pub public_create_gcs_bucket_config: PublicStorageGCSCreateBucketTmp,
 
     // Bellman cuda
     #[clap(flatten)]
-    #[serde(flatten)]
     pub bellman_cuda_config: InitBellmanCudaArgs,
+    #[clap(long, default_missing_value = "true", num_args = 0..=1)]
+    pub bellman_cuda: Option<bool>,
+
+    #[clap(long, default_missing_value = "true", num_args = 0..=1)]
+    pub setup_compressor_keys: Option<bool>,
+    #[clap(flatten)]
+    pub compressor_keys_args: CompressorKeysArgs,
 
     #[clap(flatten)]
-    #[serde(flatten)]
-    pub setup_key_config: SetupKeyConfigTmp,
+    pub setup_keys_args: SetupKeysArgs,
+    #[clap(long, default_missing_value = "true", num_args = 0..=1)]
+    pub setup_keys: Option<bool>,
 
     #[clap(long)]
     pub setup_database: Option<bool>,
@@ -137,7 +144,7 @@ pub struct PublicStorageGCSCreateBucketTmp {
 }
 
 #[derive(Clone, Debug, Serialize, Deserialize, Parser, Default)]
-pub struct SetupKeyConfigTmp {
+pub struct SetupCompressorKeyConfigTmp {
     #[clap(long)]
     pub download_key: Option<bool>,
     #[clap(long)]
@@ -171,12 +178,6 @@ pub enum ProofStorageConfig {
     GCSCreateBucket(ProofStorageGCSCreateBucket),
 }
 
-#[derive(Debug, Clone)]
-pub struct SetupKeyConfig {
-    pub download_key: bool,
-    pub setup_key_path: String,
-}
-
 #[derive(Debug, Clone)]
 pub struct ProverDatabaseConfig {
     pub database_config: DatabaseConfig,
@@ -187,7 +188,8 @@ pub struct ProverDatabaseConfig {
 pub struct ProverInitArgsFinal {
     pub proof_store: ProofStorageConfig,
     pub public_store: Option<ProofStorageConfig>,
-    pub setup_key_config: SetupKeyConfig,
+    pub compressor_key_args: Option<CompressorKeysArgs>,
+    pub setup_keys: Option<SetupKeysArgs>,
     pub bellman_cuda_config: Option<InitBellmanCudaArgs>,
     pub cloud_type: CloudConnectionMode,
     pub database_config: Option<ProverDatabaseConfig>,
@@ -197,20 +199,23 @@ impl ProverInitArgs {
     pub(crate) fn fill_values_with_prompt(
         &self,
         shell: &Shell,
-        setup_key_path: &str,
+        default_compressor_key_path: &str,
         chain_config: &ChainConfig,
     ) -> anyhow::Result<ProverInitArgsFinal> {
         let proof_store = self.fill_proof_storage_values_with_prompt(shell)?;
         let public_store = self.fill_public_storage_values_with_prompt(shell)?;
-        let setup_key_config = self.fill_setup_key_values_with_prompt(setup_key_path);
-        let bellman_cuda_config = self.fill_bellman_cuda_values_with_prompt()?;
+        let compressor_key_args =
+            self.fill_setup_compressor_key_values_with_prompt(default_compressor_key_path);
+        let bellman_cuda_config = self.fill_bellman_cuda_values_with_prompt();
         let cloud_type = self.get_cloud_type_with_prompt();
         let database_config = self.fill_database_values_with_prompt(chain_config);
+        let setup_keys = self.fill_setup_keys_values_with_prompt();
 
         Ok(ProverInitArgsFinal {
             proof_store,
             public_store,
-            setup_key_config,
+            compressor_key_args,
+            setup_keys,
             bellman_cuda_config,
             cloud_type,
             database_config,
@@ -336,29 +341,38 @@ impl ProverInitArgs {
         }
     }
 
-    fn fill_setup_key_values_with_prompt(&self, setup_key_path: &str) -> SetupKeyConfig {
-        let download_key = self
-            .clone()
-            .setup_key_config
-            .download_key
-            .unwrap_or_else(|| {
-                PromptConfirm::new(MSG_DOWNLOAD_SETUP_KEY_PROMPT)
-                    .default(true)
-                    .ask()
-            });
-        let setup_key_path = self
-            .clone()
-            .setup_key_config
-            .setup_key_path
-            .unwrap_or_else(|| {
-                Prompt::new(MSG_SETUP_KEY_PATH_PROMPT)
-                    .default(setup_key_path)
-                    .ask()
-            });
+    fn fill_setup_compressor_key_values_with_prompt(
+        &self,
+        default_path: &str,
+    ) -> Option<CompressorKeysArgs> {
+        let download_key = self.clone().setup_compressor_keys.unwrap_or_else(|| {
+            PromptConfirm::new(MSG_DOWNLOAD_SETUP_COMPRESSOR_KEY_PROMPT)
+                .default(false)
+                .ask()
+        });
 
-        SetupKeyConfig {
-            download_key,
-            setup_key_path,
+        if download_key {
+            Some(
+                self.compressor_keys_args
+                    .clone()
+                    .fill_values_with_prompt(default_path),
+            )
+        } else {
+            None
+        }
+    }
+
+    fn fill_setup_keys_values_with_prompt(&self) -> Option<SetupKeysArgs> {
+        let args = self.setup_keys_args.clone();
+
+        if self.setup_keys.unwrap_or_else(|| {
+            PromptConfirm::new(MSG_SETUP_KEYS_PROMPT)
+                .default(false)
+                .ask()
+        }) {
+            Some(args)
+        } else {
+            None
         }
     }
 
@@ -460,19 +474,16 @@ impl ProverInitArgs {
         })
     }
 
-    fn fill_bellman_cuda_values_with_prompt(&self) -> anyhow::Result<Option<InitBellmanCudaArgs>> {
-        let init_bellman_cuda = if self.bellman_cuda_config.bellman_cuda_dir.is_none() {
-            PromptConfirm::new("Do you want to setup bellman-cuda? (You can do it later by running `zk_inception prover init-bellman-cuda`)").default(false).ask()
-        } else {
-            true
-        };
-
-        if init_bellman_cuda {
-            Ok(Some(
-                self.bellman_cuda_config.clone().fill_values_with_prompt()?,
-            ))
+    fn fill_bellman_cuda_values_with_prompt(&self) -> Option<InitBellmanCudaArgs> {
+        let args = self.bellman_cuda_config.clone();
+        if self.bellman_cuda.unwrap_or_else(|| {
+            PromptConfirm::new(MSG_INITIALIZE_BELLMAN_CUDA_PROMPT)
+                .default(false)
+                .ask()
+        }) {
+            Some(args)
         } else {
-            Ok(None)
+            None
         }
     }
 

diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/args/init_bellman_cuda.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/args/init_bellman_cuda.rs
@@ -30,7 +30,7 @@ impl std::fmt::Display for BellmanCudaPathSelection {
 }
 
 impl InitBellmanCudaArgs {
-    pub fn fill_values_with_prompt(self) -> anyhow::Result<InitBellmanCudaArgs> {
+    pub fn fill_values_with_prompt(self) -> InitBellmanCudaArgs {
         let bellman_cuda_dir = self.bellman_cuda_dir.unwrap_or_else(|| {
             match PromptSelect::new(
                 MSG_BELLMAN_CUDA_ORIGIN_SELECT,
@@ -43,8 +43,8 @@ impl InitBellmanCudaArgs {
             }
         });
 
-        Ok(InitBellmanCudaArgs {
+        InitBellmanCudaArgs {
             bellman_cuda_dir: Some(bellman_cuda_dir),
-        })
+        }
     }
 }
diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/args/mod.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/args/mod.rs
@@ -1,3 +1,4 @@
+pub mod compressor_keys;
 pub mod init;
 pub mod init_bellman_cuda;
 pub mod run;

diff --git a/zk_toolbox/crates/zk_inception/src/commands/prover/args/setup_keys.rs b/zk_toolbox/crates/zk_inception/src/commands/prover/args/setup_keys.rs
@@ -2,7 +2,7 @@ use clap::{Parser, ValueEnum};
 use common::PromptSelect;
 use strum::{EnumIter, IntoEnumIterator};
 
-use crate::messages::{MSG_SETUP_KEYS_DOWNLOAD_HELP, MSG_SETUP_KEYS_REGION_PROMPT};
+use crate::messages::{MSG_SETUP_KEYS_DOWNLOAD_SELECTION_PROMPT, MSG_SETUP_KEYS_REGION_PROMPT};
 
 #[derive(Debug, Clone, Parser, Default)]
 pub struct SetupKeysArgs {
@@ -33,9 +33,9 @@ pub enum Region {
 
 impl SetupKeysArgs {
     pub fn fill_values_with_prompt(self) -> SetupKeysArgsFinal {
-        let mode = self
-            .mode
-            .unwrap_or_else(|| PromptSelect::new(MSG_SETUP_KEYS_DOWNLOAD_HELP, Mode::iter()).ask());
+        let mode = self.mode.unwrap_or_else(|| {
+            PromptSelect::new(MSG_SETUP_KEYS_DOWNLOAD_SELECTION_PROMPT, Mode::iter()).ask()
+        });
 
         if mode == Mode::Download {
             let region = self.region.unwrap_or_else(|| {