feat(prover): merging key generation into a single binary (matter-lab…

…s#1101) ## What ❔ * Merging 3 different key generation steps into a single binary * Added 'clap' support with arguments to select what keys should be created. * Also started printing out the md5sum of the generated setup keys, for easier debugging. ## Why ❔ * Having it in one binary is cleaner, and will allow new things (like checking the validity of all the keys) in the next steps. * There will be a separate PR to update our nighly/manual workflows.
danzid-tech · Feb 16, 2024 · 6de8b84 · 6de8b84
1 parent 98f87cd
commit 6de8b84
Show file tree

Hide file tree

Showing 10 changed files with 367 additions and 102 deletions.
diff --git a/checks-config/era.dic b/checks-config/era.dic
@@ -896,3 +896,4 @@ downcasting
 parameterized
 reimplementation
 composability
+md5
diff --git a/infrastructure/zk/src/prover_setup.ts b/infrastructure/zk/src/prover_setup.ts
@@ -176,38 +176,21 @@ async function setupArtifactsMode() {
     }
 }
 
-async function generateSetupDataForBaseLayer(proverType: ProverType) {
-    await generateSetupData(true, proverType);
-}
-
-async function generateSetupDataForRecursiveLayers(proverType: ProverType) {
-    await generateSetupData(false, proverType);
-}
-
-async function generateSetupData(isBaseLayer: boolean, proverType: ProverType) {
+async function generateAllSetupData(proverType: ProverType) {
     const currentEnv = env.get();
 
     const proverKeysDir = `${process.env.ZKSYNC_HOME}/etc/hyperchains/prover-keys/${currentEnv}/${proverType}/`;
     fs.mkdirSync(proverKeysDir, { recursive: true });
     const proverDir = `${process.env.ZKSYNC_HOME}/prover`;
     process.chdir(proverDir);
-    const range = isBaseLayer ? 13 : 15;
     const gpuFeatureFlag = proverType == ProverType.GPU ? '--features "gpu"' : '';
-    for (let i = 1; i <= range; i++) {
-        const spawnCommand = `zk f cargo run ${gpuFeatureFlag} --release --bin zksync_setup_data_generator_fri -- --numeric-circuit ${i} ${
-            isBaseLayer ? '--is_base_layer' : ''
-        }`;
-        await utils.spawn(spawnCommand);
-    }
+    const command = proverType == ProverType.GPU ? 'generate-sk-gpu' : 'generate-sk';
+    const spawnCommand = `zk f cargo run ${gpuFeatureFlag} --release --bin key_generator -- ${command} all`;
+    await utils.spawn(spawnCommand);
 
     process.chdir(process.env.ZKSYNC_HOME as string);
 }
 
-async function generateAllSetupData(proverType: ProverType) {
-    await generateSetupDataForBaseLayer(proverType);
-    await generateSetupDataForRecursiveLayers(proverType);
-}
-
 async function downloadDefaultSetupKeys(proverType: ProverType, region: string) {
     const proverKeysUrls = require(`${process.env.ZKSYNC_HOME}/prover/setup-data-${proverType}-keys.json`);
     const currentEnv = env.get();

diff --git a/prover/Cargo.lock b/prover/Cargo.lock
diff --git a/prover/setup.sh b/prover/setup.sh
@@ -2,8 +2,10 @@
 # This script sets up the necessary data needed by the CPU/GPU FRI prover to be used locally.
 
 GPU_FLAG=""
+GENERATE_SK_COMMAND="generate-sk"
 if [ "$1" = "gpu" ]; then
     GPU_FLAG='--features gpu'
+    GENERATE_SK_COMMAND="generate-sk-gpu"
 fi
 
 if [[ -z "${ZKSYNC_HOME}" ]]; then
@@ -21,24 +23,32 @@ rm ../etc/env/base/fri_proof_compressor.toml.backup
 
 zk config compile dev
 
+
+SHOULD_GENERATE_KEYS=false
+
 for i in {1..13}
 do
     if ! [ -f vk_setup_data_generator_server_fri/data/setup_basic_${i}_data.bin ]; then
-        zk f cargo run $GPU_FLAG --release --bin zksync_setup_data_generator_fri -- --numeric-circuit $i --is_base_layer
+        SHOULD_GENERATE_KEYS=true        
     fi
 done
 
 if ! [ -f vk_setup_data_generator_server_fri/data/setup_scheduler_data.bin ]; then
-    zk f cargo run $GPU_FLAG --release --bin zksync_setup_data_generator_fri -- --numeric-circuit 1
+    SHOULD_GENERATE_KEYS=true            
 fi
 
 if ! [ -f vk_setup_data_generator_server_fri/data/setup_node_data.bin ]; then
-    zk f cargo run $GPU_FLAG --release --bin zksync_setup_data_generator_fri -- --numeric-circuit 2
+    SHOULD_GENERATE_KEYS=true        
 fi
 
 for i in {3..15}
 do
     if ! [ -f vk_setup_data_generator_server_fri/data/setup_leaf_${i}_data.bin ]; then
-        zk f cargo run $GPU_FLAG --release --bin zksync_setup_data_generator_fri -- --numeric-circuit $i
+        SHOULD_GENERATE_KEYS=true        
     fi
 done
+
+
+if [ "$SHOULD_GENERATE_KEYS" = "true" ]; then
+    zk f cargo run $GPU_FLAG --release --bin key_generator -- $GENERATE_SK_COMMAND all
+fi
diff --git a/prover/vk_setup_data_generator_server_fri/Cargo.toml b/prover/vk_setup_data_generator_server_fri/Cargo.toml
@@ -5,20 +5,13 @@ edition = "2021"
 
 
 [[bin]]
-name = "zksync_vk_generator_fri"
+name = "key_generator"
 path = "src/main.rs"
 
-[[bin]]
-name = "zksync_setup_data_generator_fri"
-path = "src/setup_data_generator.rs"
-
 [lib]
 name = "zksync_vk_setup_data_server_fri"
 path = "src/lib.rs"
 
-[[bin]]
-name = "zksync_commitment_generator_fri"
-path = "src/commitment_generator.rs"
 
 [dependencies]
 vlog = { path = "../../core/lib/vlog" }
@@ -35,6 +28,7 @@ zksync_config = { path = "../../core/lib/config" }
 zksync_env_config = { path = "../../core/lib/env_config" }
 
 anyhow = "1.0"
+clap = { version = "4.4.6", features = ["derive"] }
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter"] }
 serde_json = "1.0"
@@ -45,6 +39,7 @@ bincode = "1"
 structopt = "0.3.26"
 once_cell = "1.8.0"
 toml_edit = "0.14.4"
+md5 = "0.7.0"
 
 [dev-dependencies]
 proptest = "1.2.0"

diff --git a/prover/vk_setup_data_generator_server_fri/README.md b/prover/vk_setup_data_generator_server_fri/README.md
@@ -1,22 +1,52 @@
-# Setup data and VK generator and server
+# Circuit key generator
+
+Tool for generating different type of circuit related keys:
+
+- verification keys
+- setup keys (for CPU and GPU)
+- commitments
+
+## Verification keys
+
+The current set of verification keys is committed under 'data/' directory. If you want to refresh it (for example after
+circuit changes), first please make sure that you have a CRS file (used for SNARK key), and then you can run:
+
+```shell
+CRS_FILE=yyy ZKSYNC_HOME=xxx cargo run --release --bin key_generator generate-vk
+```
+
+### CRS FILE
 
 The SNARK VK generation requires the `CRS_FILE` environment variable to be present and point to the correct file. The
-file can be downloaded from the following
-[link](https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2^26.key) its also present in dir after
-zk init keys/setup/setup_2^26.key
+file can be downloaded (around 1GB) from the following
+[link](https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2^24.key) its also present in dir after
+zk init keys/setup/setup_2^24.key
+
+## Commitments
+
+Commitments is basically a 'hash' of the verification keys, that is used in our configuration, and also in Verifier.sol.
 
-## generating setup-data for specific circuit type
+You can run it with `dry-run`, to see the results, or set `dry-run` to false to update the config in
+`etc/env/base/contracts.toml`.
 
-`zk f cargo +nightly-2023-08-21 run --release --bin zksync_setup_data_generator_fri -- --numeric-circuit 1 --is_base_layer`
+```shell
+ZKSYNC_HOME=xxx cargo run --release --bin key_generator update-commitments --dry-run=true
+```
 
-## generating GPU setup-data for specific circuit type
+## Setup keys
 
-`zk f cargo +nightly-2023-08-21 run --features "gpu" --release --bin zksync_setup_data_generator_fri -- --numeric-circuit 1 --is_base_layer`
+Setup keys are used when you run the actual prover. They are around 15GB for each circuit type, and we have different
+setup keys for GPU vs CPU prover.
 
-## Generating VK's
+For example, the command below will generate the setup keys for the basic circuit of type 3.
 
-`cargo +nightly-2023-08-21 run --release --bin zksync_vk_generator_fri`
+```shell
+ZKSYNC_HOME=xxx cargo run --release --bin key_generator generate-sk basic 3
+```
 
-## generating VK commitment for existing VK's
+And command below will generate all the GPU keys (notice that we have to build with 'gpu' feature enabled, as this adds
+additional dependencies).
 
-`cargo +nightly-2023-08-21 run --release --bin zksync_commitment_generator_fri`
+```shell
+ZKSYNC_HOME=xxx cargo run --feature=gpu --release --bin key_generator generate-sk-gpu all
+```
diff --git a/prover/vk_setup_data_generator_server_fri/src/commitment_generator.rs b/prover/vk_setup_data_generator_server_fri/src/commitment_generator.rs
@@ -4,22 +4,23 @@ use zksync_vk_setup_data_server_fri::{
     vk_commitment_helper::{get_toml_formatted_value, read_contract_toml, write_contract_toml},
 };
 
-fn main() -> anyhow::Result<()> {
-    tracing::info!("Starting commitment generation!");
-    read_and_update_contract_toml()
-}
-
-fn read_and_update_contract_toml() -> anyhow::Result<()> {
+pub fn read_and_update_contract_toml(dryrun: bool) -> anyhow::Result<()> {
     let mut contract_doc = read_contract_toml().context("read_contract_toml()")?;
     let vk_commitments = generate_commitments().context("generate_commitments()")?;
+
     contract_doc["contracts"]["FRI_RECURSION_LEAF_LEVEL_VK_HASH"] =
         get_toml_formatted_value(vk_commitments.leaf);
     contract_doc["contracts"]["FRI_RECURSION_NODE_LEVEL_VK_HASH"] =
         get_toml_formatted_value(vk_commitments.node);
     contract_doc["contracts"]["FRI_RECURSION_SCHEDULER_LEVEL_VK_HASH"] =
         get_toml_formatted_value(vk_commitments.scheduler);
     tracing::info!("Updated toml content: {:?}", contract_doc.to_string());
-    write_contract_toml(contract_doc).context("write_contract_toml")
+    if !dryrun {
+        write_contract_toml(contract_doc).context("write_contract_toml")?;
+    } else {
+        tracing::warn!("DRY RUN - Not updating the file.")
+    }
+    Ok(())
 }
 
 #[cfg(test)]
@@ -28,6 +29,6 @@ mod test {
 
     #[test]
     fn test_read_and_update_contract_toml() {
-        read_and_update_contract_toml().unwrap();
+        read_and_update_contract_toml(true).unwrap();
     }
 }
diff --git a/prover/vk_setup_data_generator_server_fri/src/commitment_utils.rs b/prover/vk_setup_data_generator_server_fri/src/commitment_utils.rs
@@ -22,6 +22,7 @@ lazy_static! {
     static ref COMMITMENTS: Lazy<L1VerifierConfig> = Lazy::new(|| { circuit_commitments().unwrap() });
 }
 
+#[derive(Debug)]
 pub struct VkCommitments {
     pub leaf: String,
     pub node: String,