feat: GPU prover (#66)

* feat: use GPU for proving

* chore: update Dockerfile

* chore: cleanup

* chore: update params.zip

* chore: update circom_proofgen.sh

* chore: update

Relayer.Dockerfile

@@ -1,5 +1,5 @@
 # Use the base image
-FROM us-central1-docker.pkg.dev/zkairdrop/ether-email-auth/relayer-base:v1
+FROM bisht13/relayer-base
 
 # Copy the project files
 COPY packages/relayer /relayer/packages/relayer

packages/prover/Dockerfile

@@ -1,49 +1,66 @@
-FROM python:3.10
+FROM nvidia/cuda:12.4.0-devel-ubuntu22.04
 
 RUN apt-get update && apt-get upgrade -y 
 # Update the package list and install necessary dependencies
 RUN apt-get update && \
-    apt install -y cmake build-essential pkg-config libssl-dev libgmp-dev libsodium-dev nasm git awscli gcc nodejs npm 
+    DEBIAN_FRONTEND=noninteractive apt install -y --no-install-recommends \
+    cmake \
+    build-essential \
+    pkg-config \
+    libssl-dev \
+    libgmp-dev \
+    libffi-dev \
+    libsodium-dev \
+    nasm \
+    git \
+    awscli \
+    gcc \
+    nodejs \
+    npm \
+    curl \
+    m4 \
+    python3 \
+    python3-pip \
+    python3-dev \
+    wget \
+    software-properties-common \
+    unzip \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set Python 3 as the default python version
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python3 1 \
+    && update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1
 
 # Node install
 RUN npm install -g n 
-RUN n 18
+RUN n 22
 RUN npm install -g yarn snarkjs
-RUN git clone -b feat/body-parsing-circuit https://github.com/zkemail/ether-email-auth.git
+
+RUN git clone -b feat/gpu https://github.com/zkemail/ether-email-auth.git
 WORKDIR /ether-email-auth/packages/prover
 RUN pip install -r requirements.txt
 RUN cp ./circom_proofgen.sh /root
 WORKDIR /root
 RUN ls /root
-# RUN mkdir params
-# RUN cp /email-wallet/packages/prover/params/account_creation.wasm /root/params
-# RUN cp /email-wallet/packages/prover/params/account_init.wasm /root/params
-# RUN cp /email-wallet/packages/prover/params/account_transport.wasm /root/params
-# RUN cp /email-wallet/packages/prover/params/claim.wasm /root/params
-# RUN cp /email-wallet/packages/prover/params/email_sender.wasm /root/params
 RUN mkdir params
 WORKDIR /root/params
-RUN gdown "https://drive.google.com/uc?id=1XDPFIL5YK8JzLGoTjmHLXO9zMDjSQcJH"
+RUN gdown "https://drive.google.com/uc?id=1l3mNqFYv-YZc2efFlphFUkoaCnGCxFtE"
 RUN unzip params.zip
 RUN mv params/* /root/params
 WORKDIR /root
 RUN ls params
-# RUN mv build params
-# RUN curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-account-creation/contributions/emailwallet-account-creation_00019.zkey --output ./params/account_creation.zkey
-# RUN curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-account-init/contributions/emailwallet-account-init_00007.zkey --output ./params/account_init.zkey
-# RUN curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-account-transport/contributions/emailwallet-account-transport_00005.zkey --output ./params/account_transport.zkey
-# RUN curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-claim/contributions/emailwallet-claim_00006.zkey --output ./params/claim.zkey
-# RUN curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-email-sender/contributions/emailwallet-email-sender_00006.zkey --output ./params/email_sender.zkey
 RUN chmod +x circom_proofgen.sh
 RUN mkdir build
 
-RUN git clone https://github.com/iden3/rapidsnark-old.git rapidsnark
+RUN git clone https://github.com/Orbiter-Finance/rapidsnark.git rapidsnark
 WORKDIR /root/rapidsnark
 RUN yarn
 RUN git submodule init
 RUN git submodule update
-RUN npx task createFieldSources
-RUN npx task buildPistache
-RUN npx task buildProver
-RUN chmod +x build/prover
+RUN ./build_gmp.sh host
+RUN mkdir build_prover
+WORKDIR /root/rapidsnark/build_prover
+RUN cmake .. -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=../package -DNVML_LIBRARY=/usr/local/cuda-12.4/targets/x86_64-linux/lib/stubs/libnvidia-ml.so
+RUN make -j$(nproc) && make install
+RUN chmod +x ../package/bin/prover_cuda
 WORKDIR /root

packages/prover/circom_proofgen.sh

@@ -21,36 +21,24 @@ public_path="${buildDir}/rapidsnark_public_${circuitName}_${nonce}.json"
 cd "${SCRIPT_DIR}"
 echo "entered zk email path: ${SCRIPT_DIR}"
 
-echo "NODE_OPTIONS='--max-old-space-size=644000' snarkjs wc "${paramsDir}/${circuitName}.wasm" "${input_path}" "${witness_path}""
-NODE_OPTIONS='--max-old-space-size=644000' snarkjs wc "${paramsDir}/${circuitName}.wasm" "${input_path}" "${witness_path}"  | tee /dev/stderr
+${paramsDir}/${circuitName}_cpp/${circuitName} "${input_path}" "${witness_path}" | tee /dev/stderr
 status_jswitgen=$?
-echo "✓ Finished witness generation with js! ${status_jswitgen}"
-
-# TODO: Get C-based witness gen to work
-# echo "/${build_dir}/${CIRCUIT_NAME}_cpp/${CIRCUIT_NAME} ${input_wallet_path} ${witness_path}"
-# "/${build_dir}/${CIRCUIT_NAME}_cpp/${CIRCUIT_NAME}" "${input_wallet_path}" "${witness_path}"
-# status_c_wit=$?
-
-# echo "Finished C witness gen! Status: ${status_c_wit}"
-# if [ $status_c_wit -ne 0 ]; then
-#     echo "C based witness gen failed with status (might be on machine specs diff than compilation): ${status_c_wit}"
-#     exit 1
-# fi
+echo "✓ Finished witness generation with cpp! ${status_jswitgen}"
 
 if [ $isLocal = 1 ]; then
     # DEFAULT SNARKJS PROVER (SLOW)
     NODE_OPTIONS='--max-old-space-size=644000' snarkjs groth16 prove "${paramsDir}/${circuitName}.zkey" "${witness_path}" "${proof_path}" "${public_path}"
     status_prover=$?
     echo "✓ Finished slow proofgen! Status: ${status_prover}"
 else
-    # RAPIDSNARK PROVER (10x FASTER)
-    echo "ldd ${SCRIPT_DIR}/rapidsnark/build/prover"
-    ldd "${SCRIPT_DIR}/rapidsnark/build/prover"
+    # RAPIDSNARK PROVER GPU
+    echo "ldd ${SCRIPT_DIR}/rapidsnark/package/bin/prover_cuda"
+    ldd "${SCRIPT_DIR}/rapidsnark/package/bin/prover_cuda"
     status_lld=$?
     echo "✓ lld prover dependencies present! ${status_lld}"
 
-    echo "${SCRIPT_DIR}/rapidsnark/build/prover ${paramsDir}/${circuitName}.zkey ${witness_path} ${proof_path} ${public_path}"
-    "${SCRIPT_DIR}/rapidsnark/build/prover" "${paramsDir}/${circuitName}.zkey" "${witness_path}" "${proof_path}" "${public_path}"  | tee /dev/stderr
+    echo "${SCRIPT_DIR}/rapidsnark/package/bin/prover_cuda ${paramsDir}/${circuitName}.zkey ${witness_path} ${proof_path} ${public_path}"
+    "${SCRIPT_DIR}/rapidsnark/package/bin/prover_cuda" "${paramsDir}/${circuitName}.zkey" "${witness_path}" "${proof_path}" "${public_path}"  | tee /dev/stderr
     status_prover=$?
     echo "✓ Finished rapid proofgen! Status: ${status_prover}"
 fi

packages/prover/core.py

@@ -3,7 +3,7 @@
 import json
 import logging
 
-logger = logging.getLogger(__name__)
+# logger = logging.getLogger(__name__)
 
 
 def gen_email_auth_proof(nonce: str, is_local: bool, input: dict) -> dict:
@@ -19,18 +19,28 @@ def gen_email_auth_proof(nonce: str, is_local: bool, input: dict) -> dict:
 
 
 def store_input(circuit_name: str, nonce: str, json_data: dict):
+    print("Storing input")
     cur_dir = get_cur_dir()
+    print(f"Current dir: {cur_dir}")
     build_dir = os.path.join(cur_dir, "build")
     # check if build_dir exists
     if not os.path.exists(build_dir):
         os.makedirs(build_dir)
 
+    print(f"Build dir: {build_dir}")
     json_file_path = os.path.join(
         build_dir, "input_" + circuit_name + "_" + nonce + ".json"
     )
-    logger.info(f"Store user input to {json_file_path}")
+    print(f"Json file path: {json_file_path}")
+    print(f"Json data: {json_data}")
+    print(f"Json data type: {type(json_data)}")
+    # logger.info(f"Store user input to {json_file_path}")
     with open(json_file_path, "w") as json_file:
         json_file.write(json_data)
+    # Read the file back
+    with open(json_file_path, "r") as json_file:
+        print(json_file.read())
+    print("Stored input")
 
 
 def load_proof(circuit_name: str, nonce: str) -> dict:
@@ -39,7 +49,7 @@ def load_proof(circuit_name: str, nonce: str) -> dict:
     json_file_path = os.path.join(
         build_dir, "rapidsnark_proof_" + circuit_name + "_" + nonce + ".json"
     )
-    logger.info(f"Loading proof from {json_file_path}")
+    # logger.info(f"Loading proof from {json_file_path}")
     with open(json_file_path, "r") as json_file:
         return json.loads(json_file.read())
 
@@ -50,7 +60,7 @@ def load_pub_signals(circuit_name: str, nonce: str) -> dict:
     json_file_path = os.path.join(
         build_dir, "rapidsnark_public_" + circuit_name + "_" + nonce + ".json"
     )
-    logger.info(f"Loading public signals from {json_file_path}")
+    # logger.info(f"Loading public signals from {json_file_path}")
     with open(json_file_path, "r") as json_file:
         return json.loads(json_file.read())
 
@@ -59,9 +69,9 @@ def gen_proof(circuit_name: str, nonce: str, is_local: bool):
     is_local_int: int = 1 if is_local else 0
     cur_dir = get_cur_dir()
     params_dir = os.path.join(cur_dir, "params")
-    logger.info(f"Params dir: {params_dir}")
+    # logger.info(f"Params dir: {params_dir}")
     build_dir = os.path.join(cur_dir, "build")
-    logger.info(f"Build dir: {build_dir}")
+    # logger.info(f"Build dir: {build_dir}")
     result = subprocess.run(
         [
             os.path.join(cur_dir, "circom_proofgen.sh"),
@@ -72,9 +82,9 @@ def gen_proof(circuit_name: str, nonce: str, is_local: bool):
             str(is_local_int),
         ]
     )
-    logger.info(f"Proof generation result: {result.returncode}")
-    if result.stderr is not None:
-        logger.error(result.stderr)
+    # logger.info(f"Proof generation result: {result.returncode}")
+    # if result.stderr is not None:
+    #     logger.error(result.stderr)
     print(result.stdout)
     print(result.stderr)
 

packages/prover/local_setup.sh

@@ -6,11 +6,6 @@ mkdir -p build
 npm install -g snarkjs@latest
 pip install -r requirements.txt
 mkdir build && cd build
-gdown "https://drive.google.com/uc?id=1XDPFIL5YK8JzLGoTjmHLXO9zMDjSQcJH"
+gdown "https://drive.google.com/uc?id=1l3mNqFYv-YZc2efFlphFUkoaCnGCxFtE"
 unzip params.zip
-# curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-account-creation/contributions/emailwallet-account-creation_00019.zkey --output /root/params/account_creation.zkey
-# curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-account-init/contributions/emailwallet-account-init_00007.zkey --output /root/params/account_init.zkey
-# curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-account-transport/contributions/emailwallet-account-transport_00005.zkey --output /root/params/account_transport.zkey
-# curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-claim/contributions/emailwallet-claim_00006.zkey --output /root/params/claim.zkey
-# curl https://email-wallet-trusted-setup-ceremony-pse-p0tion-production.s3.eu-central-1.amazonaws.com/circuits/emailwallet-email-sender/contributions/emailwallet-email-sender_00006.zkey --output /root/params/email_sender.zkey
 chmod +x circom_proofgen.sh

packages/prover/modal_server.py

@@ -5,7 +5,7 @@
 from google.cloud.logging_v2.handlers import setup_logging
 from google.oauth2 import service_account
 
-app = modal.App("email-auth-prover-v1.3.0")
+app = modal.App("email-auth-prover-v1.4.0")
 
 image = modal.Image.from_dockerfile("Dockerfile")
 
@@ -15,7 +15,8 @@
     mounts=[
         modal.Mount.from_local_python_packages("core"),
     ],
-    cpu=14,
+    cpu=16,
+    gpu="any",
     secrets=[modal.Secret.from_name("gc-ether-email-auth-prover")],
 )
 @modal.wsgi_app()
@@ -45,17 +46,17 @@ def prove_email_auth():
         print("prove_email_auth")
         req = request.get_json()
         input = req["input"]
-        logger = logging.getLogger(__name__)
-        logger.info(req)
+        # logger = logging.getLogger(__name__)
+        # logger.info(req)
         print(req)
         nonce = random.randint(
             0,
             sys.maxsize,
         )
-        logger.info(nonce)
+        # logger.info(nonce)
         print(nonce)
         proof = gen_email_auth_proof(str(nonce), False, input)
-        logger.info(proof)
+        # logger.info(proof)
         print(proof)
         return jsonify(proof)