feat(zk_toolbox): Allow running docker images for provers

.github/workflows/build-docker-from-tag.yml

@@ -23,7 +23,7 @@ concurrency: docker-build
 jobs:
   setup:
     name: Setup
-    runs-on: [ubuntu-latest]
+    runs-on: [ ubuntu-latest ]
     outputs:
       image_tag_suffix: ${{ steps.set.outputs.image_tag_suffix }}
       prover_fri_gpu_key_id: ${{ steps.extract-prover-fri-setup-key-ids.outputs.gpu_short_commit_sha }}
@@ -48,7 +48,7 @@ jobs:
 
   build-push-core-images:
     name: Build and push image
-    needs: [setup]
+    needs: [ setup ]
     uses: ./.github/workflows/build-core-template.yml
     if: contains(github.ref_name, 'core')
     secrets:
@@ -60,7 +60,7 @@ jobs:
 
   build-push-tee-prover-images:
     name: Build and push images
-    needs: [setup]
+    needs: [ setup ]
     uses: ./.github/workflows/build-tee-prover-template.yml
     if: contains(github.ref_name, 'core')
     secrets:
@@ -72,7 +72,7 @@ jobs:
 
   build-push-contract-verifier:
     name: Build and push image
-    needs: [setup]
+    needs: [ setup ]
     uses: ./.github/workflows/build-contract-verifier-template.yml
     if: contains(github.ref_name, 'contract_verifier')
     secrets:
@@ -83,20 +83,20 @@ jobs:
 
   build-push-prover-images:
     name: Build and push image
-    needs: [setup]
+    needs: [ setup ]
     uses: ./.github/workflows/build-prover-template.yml
     if: contains(github.ref_name, 'prover')
     with:
       image_tag_suffix: ${{ needs.setup.outputs.image_tag_suffix }}
       ERA_BELLMAN_CUDA_RELEASE: ${{ vars.ERA_BELLMAN_CUDA_RELEASE }}
-      CUDA_ARCH: "60;70;75;89"
+      CUDA_ARCH: "60;70;75;80;89"
     secrets:
       DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }}
       DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}
 
   build-push-witness-generator-image-avx512:
     name: Build and push image
-    needs: [setup]
+    needs: [ setup ]
     uses: ./.github/workflows/build-witness-generator-template.yml
     if: contains(github.ref_name, 'prover')
     with:
@@ -110,7 +110,7 @@ jobs:
 
   build-gar-prover-fri-gpu:
     name: Build GAR prover FRI GPU
-    needs: [setup, build-push-prover-images]
+    needs: [ setup, build-push-prover-images ]
     uses: ./.github/workflows/build-prover-fri-gpu-gar.yml
     if: contains(github.ref_name, 'prover')
     with:

docker/proof-fri-gpu-compressor/Dockerfile

@@ -4,8 +4,7 @@ FROM nvidia/cuda:12.2.0-devel-ubuntu22.04 as builder
 ARG DEBIAN_FRONTEND=noninteractive
 
 ARG CUDA_ARCH=89
-ARG A100_CUDA_ARCH=80
-ENV CUDAARCHS=${CUDA_ARCH};${A100_CUDA_ARCH}
+ENV CUDAARCHS=${CUDA_ARCH}
 
 RUN apt-get update && apt-get install -y curl clang openssl libssl-dev gcc g++ git \
     pkg-config build-essential libclang-dev && \

zk_toolbox/crates/zk_inception/README.md

@@ -504,12 +504,19 @@ Run prover
   Possible values: `gateway`, `witness-generator`, `witness-vector-generator`, `prover`, `compressor`,
   `prover-job-monitor`
 
+- `--docker` - Whether to run image of the component instead of binary.
+
+  Possible values: `true`, `false`
+
 - `--round <ROUND>`
 
   Possible values: `all-rounds`, `basic-circuits`, `leaf-aggregation`, `node-aggregation`, `recursion-tip`, `scheduler`
 
 - `--threads <THREADS>`
 
+- `--max-allocation <MAX_ALLOCATION_IN_BYTES>` - in case you are running prover component, the value limits maximum
+  memory allocation of it in bytes.
+
 ## `zk_inception prover init-bellman-cuda`
 
 Initialize bellman-cuda

zk_toolbox/crates/zk_inception/src/commands/prover/args/init.rs

@@ -188,7 +188,7 @@ pub struct ProverInitArgsFinal {
     pub proof_store: ProofStorageConfig,
     pub public_store: Option<ProofStorageConfig>,
     pub setup_key_config: SetupKeyConfig,
-    pub bellman_cuda_config: InitBellmanCudaArgs,
+    pub bellman_cuda_config: Option<InitBellmanCudaArgs>,
     pub cloud_type: CloudConnectionMode,
     pub database_config: Option<ProverDatabaseConfig>,
 }
@@ -460,8 +460,20 @@ impl ProverInitArgs {
         })
     }
 
-    fn fill_bellman_cuda_values_with_prompt(&self) -> anyhow::Result<InitBellmanCudaArgs> {
-        self.bellman_cuda_config.clone().fill_values_with_prompt()
+    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()?,
+            ))
+        } else {
+            Ok(None)
+        }
     }
 
     fn get_cloud_type_with_prompt(&self) -> CloudConnectionMode {

zk_toolbox/crates/zk_inception/src/commands/prover/args/run.rs

@@ -1,8 +1,22 @@
+use anyhow::anyhow;
 use clap::{Parser, ValueEnum};
 use common::{Prompt, PromptSelect};
+use config::ChainConfig;
 use strum::{EnumIter, IntoEnumIterator};
 
-use crate::messages::{MSG_ROUND_SELECT_PROMPT, MSG_RUN_COMPONENT_PROMPT, MSG_THREADS_PROMPT};
+use crate::{
+    consts::{
+        COMPRESSOR_BINARY_NAME, COMPRESSOR_DOCKER_IMAGE, PROVER_BINARY_NAME, PROVER_DOCKER_IMAGE,
+        PROVER_GATEWAY_BINARY_NAME, PROVER_GATEWAY_DOCKER_IMAGE, PROVER_JOB_MONITOR_BINARY_NAME,
+        PROVER_JOB_MONITOR_DOCKER_IMAGE, WITNESS_GENERATOR_BINARY_NAME,
+        WITNESS_GENERATOR_DOCKER_IMAGE, WITNESS_VECTOR_GENERATOR_BINARY_NAME,
+        WITNESS_VECTOR_GENERATOR_DOCKER_IMAGE,
+    },
+    messages::{
+        MSG_ROUND_SELECT_PROMPT, MSG_RUN_COMPONENT_PROMPT, MSG_THREADS_PROMPT,
+        MSG_WITNESS_GENERATOR_ROUND_ERR,
+    },
+};
 
 #[derive(Debug, Clone, Parser, Default)]
 pub struct ProverRunArgs {
@@ -12,6 +26,10 @@ pub struct ProverRunArgs {
     pub witness_generator_args: WitnessGeneratorArgs,
     #[clap(flatten)]
     pub witness_vector_generator_args: WitnessVectorGeneratorArgs,
+    #[clap(flatten)]
+    pub fri_prover_args: FriProverRunArgs,
+    #[clap(long)]
+    pub docker: Option<bool>,
 }
 
 #[derive(
@@ -32,6 +50,108 @@ pub enum ProverComponent {
     ProverJobMonitor,
 }
 
+impl ProverComponent {
+    pub fn image_name(&self) -> &'static str {
+        match self {
+            Self::Gateway => PROVER_GATEWAY_DOCKER_IMAGE,
+            Self::WitnessGenerator => WITNESS_GENERATOR_DOCKER_IMAGE,
+            Self::WitnessVectorGenerator => WITNESS_VECTOR_GENERATOR_DOCKER_IMAGE,
+            Self::Prover => PROVER_DOCKER_IMAGE,
+            Self::Compressor => COMPRESSOR_DOCKER_IMAGE,
+            Self::ProverJobMonitor => PROVER_JOB_MONITOR_DOCKER_IMAGE,
+        }
+    }
+
+    pub fn binary_name(&self) -> &'static str {
+        match self {
+            Self::Gateway => PROVER_GATEWAY_BINARY_NAME,
+            Self::WitnessGenerator => WITNESS_GENERATOR_BINARY_NAME,
+            Self::WitnessVectorGenerator => WITNESS_VECTOR_GENERATOR_BINARY_NAME,
+            Self::Prover => PROVER_BINARY_NAME,
+            Self::Compressor => COMPRESSOR_BINARY_NAME,
+            Self::ProverJobMonitor => PROVER_JOB_MONITOR_BINARY_NAME,
+        }
+    }
+
+    pub fn get_application_args(&self, in_docker: bool) -> anyhow::Result<Vec<String>> {
+        let mut application_args = vec![];
+
+        if self == &Self::Prover || self == &Self::Compressor {
+            if in_docker {
+                application_args.push("--gpus=all".to_string());
+            } else {
+                application_args.push("--features=gpu".to_string());
+            }
+        }
+
+        Ok(application_args)
+    }
+
+    pub fn get_additional_args(
+        &self,
+        in_docker: bool,
+        args: ProverRunArgs,
+        chain: &ChainConfig,
+    ) -> anyhow::Result<Vec<String>> {
+        let mut additional_args = vec![];
+        if in_docker {
+            additional_args.push("--config-path=/configs/general.yaml".to_string());
+            additional_args.push("--secrets-path=/configs/secrets.yaml".to_string());
+        } else {
+            let general_config = chain
+                .path_to_general_config()
+                .into_os_string()
+                .into_string()
+                .map_err(|_| anyhow!("Failed to convert path to string"))?;
+            let secrets_config = chain
+                .path_to_secrets_config()
+                .into_os_string()
+                .into_string()
+                .map_err(|_| anyhow!("Failed to convert path to string"))?;
+
+            additional_args.push(format!("--config-path={}", general_config));
+            additional_args.push(format!("--secrets-path={}", secrets_config));
+        }
+
+        match self {
+            Self::WitnessGenerator => {
+                additional_args.push(
+                    match args
+                        .witness_generator_args
+                        .round
+                        .expect(MSG_WITNESS_GENERATOR_ROUND_ERR)
+                    {
+                        WitnessGeneratorRound::AllRounds => "--all_rounds",
+                        WitnessGeneratorRound::BasicCircuits => "--round=basic_circuits",
+                        WitnessGeneratorRound::LeafAggregation => "--round=leaf_aggregation",
+                        WitnessGeneratorRound::NodeAggregation => "--round=node_aggregation",
+                        WitnessGeneratorRound::RecursionTip => "--round=recursion_tip",
+                        WitnessGeneratorRound::Scheduler => "--round=scheduler",
+                    }
+                    .to_string(),
+                );
+            }
+            Self::WitnessVectorGenerator => {
+                additional_args.push(format!(
+                    "--threads={}",
+                    args.witness_vector_generator_args.threads.unwrap_or(1)
+                ));
+            }
+            Self::Prover => {
+                if args.fri_prover_args.max_allocation.is_some() {
+                    additional_args.push(format!(
+                        "--max-allocation={}",
+                        args.fri_prover_args.max_allocation.unwrap()
+                    ));
+                };
+            }
+            _ => {}
+        };
+
+        Ok(additional_args)
+    }
+}
+
 #[derive(Debug, Clone, Parser, Default)]
 pub struct WitnessGeneratorArgs {
     #[clap(long)]
@@ -76,8 +196,15 @@ impl WitnessVectorGeneratorArgs {
     }
 }
 
+#[derive(Debug, Clone, Parser, Default)]
+pub struct FriProverRunArgs {
+    /// Memory allocation limit in bytes (for prover component)
+    #[clap(long)]
+    pub max_allocation: Option<usize>,
+}
+
 impl ProverRunArgs {
-    pub fn fill_values_with_prompt(&self) -> anyhow::Result<ProverRunArgs> {
+    pub fn fill_values_with_prompt(self) -> anyhow::Result<ProverRunArgs> {
         let component = self.component.unwrap_or_else(|| {
             PromptSelect::new(MSG_RUN_COMPONENT_PROMPT, ProverComponent::iter()).ask()
         });
@@ -90,10 +217,18 @@ impl ProverRunArgs {
             .witness_vector_generator_args
             .fill_values_with_prompt(component)?;
 
+        let docker = self.docker.unwrap_or_else(|| {
+            Prompt::new("Do you want to run Docker image for the component?")
+                .default("false")
+                .ask()
+        });
+
         Ok(ProverRunArgs {
             component: Some(component),
             witness_generator_args,
             witness_vector_generator_args,
+            fri_prover_args: self.fri_prover_args,
+            docker: Some(docker),
         })
     }
 }

zk_toolbox/crates/zk_inception/src/commands/prover/init.rs

@@ -86,7 +86,9 @@ pub(crate) async fn run(args: ProverInitArgs, shell: &Shell) -> anyhow::Result<(
 
     chain_config.save_general_config(&general_config)?;
 
-    init_bellman_cuda(shell, args.bellman_cuda_config).await?;
+    if let Some(bellman_cuda_config) = args.bellman_cuda_config {
+        init_bellman_cuda(shell, bellman_cuda_config).await?;
+    }
 
     if let Some(prover_db) = &args.database_config {
         let spinner = Spinner::new(MSG_INITIALIZING_DATABASES_SPINNER);