Skip to content

Commit

Permalink
feat(prover): Add cpu-target for witness_generator compilation (#2687)
Browse files Browse the repository at this point in the history
## What ❔

<!-- What are the changes this PR brings about? -->
<!-- Example: This PR adds a PR template to the repo. -->
<!-- (For bigger PRs adding more context is appreciated) -->

## Why ❔

<!-- Why are these changes done? What goal do they contribute to? What
are the principles behind them? -->
<!-- Example: PR templates ensure PR reviewers, observers, and future
iterators are in context about the evolution of repos. -->

## Checklist

<!-- Check your PR fulfills the following items. -->
<!-- For draft PRs check the boxes as you complete them. -->

- [ ] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [ ] Tests for the changes have been added / updated.
- [ ] Documentation comments have been added / updated.
- [ ] Code has been formatted via `zk fmt` and `zk lint`.

---------

Co-authored-by: Yury Akudovich <[email protected]>
  • Loading branch information
0xVolosnikov and yorik committed Aug 21, 2024
1 parent 161ebdc commit 75851ff
Show file tree
Hide file tree
Showing 5 changed files with 243 additions and 0 deletions.
14 changes: 14 additions & 0 deletions .github/workflows/build-docker-from-tag.yml
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,20 @@ jobs:
DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }}
DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}

build-push-witness-generator-image-avx512:
name: Build and push image
needs: [setup]
uses: ./.github/workflows/build-witness-generator-template.yml
if: contains(github.ref_name, 'prover')
with:
image_tag_suffix: ${{ needs.setup.outputs.image_tag_suffix }}-avx512
ERA_BELLMAN_CUDA_RELEASE: ${{ vars.ERA_BELLMAN_CUDA_RELEASE }}
CUDA_ARCH: "60;70;75;89"
WITNESS_GENERATOR_RUST_FLAGS: "-Ctarget_feature=+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl"
secrets:
DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }}
DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}

build-gar-prover-fri-gpu:
name: Build GAR prover FRI GPU
needs: [setup, build-push-prover-images]
Expand Down
206 changes: 206 additions & 0 deletions .github/workflows/build-witness-generator-template.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,206 @@
name: Build witness generator image with custom compiler flags
on:
workflow_call:
secrets:
DOCKERHUB_USER:
description: "DOCKERHUB_USER"
required: true
DOCKERHUB_TOKEN:
description: "DOCKERHUB_TOKEN"
required: true
inputs:
ERA_BELLMAN_CUDA_RELEASE:
description: "ERA_BELLMAN_CUDA_RELEASE"
type: string
required: true
image_tag_suffix:
description: "Optional suffix to override tag name generation"
type: string
required: false
action:
description: "Action with docker image"
type: string
default: "push"
required: false
is_pr_from_fork:
description: "Indicates whether the workflow is invoked from a PR created from fork"
type: boolean
default: false
required: false
CUDA_ARCH:
description: "CUDA Arch to build"
type: string
default: "89"
required: false
WITNESS_GENERATOR_RUST_FLAGS:
description: "Rust flags for witness_generator compilation"
type: string
default: ""
required: false
outputs:
protocol_version:
description: "Protocol version of the binary"
value: ${{ jobs.build-images.outputs.protocol_version }}

jobs:
build-images:
name: Build and Push Docker Images
env:
IMAGE_TAG_SUFFIX: ${{ inputs.image_tag_suffix }}
RUNNER_COMPOSE_FILE: "docker-compose-runner-nightly.yml"
ERA_BELLMAN_CUDA_RELEASE: ${{ inputs.ERA_BELLMAN_CUDA_RELEASE }}
CUDA_ARCH: ${{ inputs.CUDA_ARCH }}
WITNESS_GENERATOR_RUST_FLAGS: ${{ inputs.WITNESS_GENERATOR_RUST_FLAGS }}
runs-on: [ matterlabs-ci-runner-c3d ]
strategy:
matrix:
component:
- witness-generator
outputs:
protocol_version: ${{ steps.protocolversion.outputs.protocol_version }}
steps:
- uses: actions/checkout@a5ac7e51b41094c92402da3b24376905380afc29 # v4
with:
submodules: "recursive"

- name: setup-env
run: |
echo ZKSYNC_HOME=$(pwd) >> $GITHUB_ENV
echo CI=1 >> $GITHUB_ENV
echo $(pwd)/bin >> $GITHUB_PATH
echo CI=1 >> .env
echo IN_DOCKER=1 >> .env
- name: start-services
run: |
echo "IMAGE_TAG_SUFFIX=${{ env.IMAGE_TAG_SUFFIX }}" >> .env
mkdir -p ./volumes/postgres
run_retried docker compose pull zk postgres
docker compose up -d zk postgres
ci_run sccache --start-server
- name: init
run: |
ci_run git config --global --add safe.directory /usr/src/zksync
ci_run git config --global --add safe.directory /usr/src/zksync/contracts/system-contracts
ci_run git config --global --add safe.directory /usr/src/zksync/contracts
ci_run zk
- name: download CRS for GPU compressor
if: matrix.component == 'proof-fri-gpu-compressor'
run: |
ci_run run_retried curl -LO https://storage.googleapis.com/matterlabs-setup-keys-us/setup-keys/setup_2\^24.key
- name: login to Docker registries
if: github.event_name != 'pull_request' && (github.ref == 'refs/heads/main' || startsWith(github.ref, 'refs/tags/'))
run: |
ci_run docker login -u ${{ secrets.DOCKERHUB_USER }} -p ${{ secrets.DOCKERHUB_TOKEN }}
ci_run gcloud auth configure-docker us-docker.pkg.dev -q
# We need to run this only when ERA_BELLMAN_CUDA_RELEASE is not available
# In our case it happens only when PR is created from fork
- name: Wait for runner IP to be not rate-limited against GH API
if: inputs.is_pr_from_fork == true
run: |
api_endpoint="https://api.github.com/users/zksync-era-bot"
wait_time=60
max_retries=60
retry_count=0
while [[ $retry_count -lt $max_retries ]]; do
response=$(run_retried curl -s -w "%{http_code}" -o temp.json "$api_endpoint")
http_code=$(echo "$response" | tail -n1)
if [[ "$http_code" == "200" ]]; then
echo "Request successful. Not rate-limited."
cat temp.json
rm temp.json
exit 0
elif [[ "$http_code" == "403" ]]; then
rate_limit_exceeded=$(jq -r '.message' temp.json | grep -i "API rate limit exceeded")
if [[ -n "$rate_limit_exceeded" ]]; then
retry_count=$((retry_count+1))
echo "API rate limit exceeded. Retry $retry_count of $max_retries. Retrying in $wait_time seconds..."
sleep $wait_time
else
echo "Request failed with HTTP status $http_code."
cat temp.json
rm temp.json
exit 1
fi
else
echo "Request failed with HTTP status $http_code."
cat temp.json
rm temp.json
exit 1
fi
done
echo "Reached the maximum number of retries ($max_retries). Exiting."
rm temp.json
exit 1
- name: protocol-version
id: protocolversion
# TODO: use -C flag, when it will become stable.
shell: bash
run: |
ci_run bash -c "cd prover && cargo build --release --bin prover_version"
PPV=$(ci_run prover/target/release/prover_version)
echo Protocol version is ${PPV}
echo "protocol_version=${PPV}" >> $GITHUB_OUTPUT
echo "PROTOCOL_VERSION=${PPV}" >> $GITHUB_ENV
- name: setup-rust-flags-env
if: matrix.component == 'witness-generator'
run: |
echo RUST_FLAGS="${{ env.WITNESS_GENERATOR_RUST_FLAGS }}" >> $GITHUB_ENV
- name: update-images
env:
DOCKER_ACTION: ${{ inputs.action }}
COMPONENT: ${{ matrix.component }}
run: |
PASSED_ENV_VARS="ERA_BELLMAN_CUDA_RELEASE,CUDA_ARCH,PROTOCOL_VERSION,RUST_FLAGS" \
ci_run zk docker $DOCKER_ACTION $COMPONENT
- name: Show sccache stats
if: always()
run: |
ci_run sccache --show-stats || true
ci_run cat /tmp/sccache_log.txt || true
copy-images:
name: Copy images between docker registries
needs: build-images
env:
IMAGE_TAG_SUFFIX: ${{ inputs.image_tag_suffix }}
PROTOCOL_VERSION: ${{ needs.build-images.outputs.protocol_version }}
runs-on: matterlabs-ci-runner
if: ${{ inputs.action == 'push' }}
strategy:
matrix:
component:
- witness-vector-generator
steps:
- name: Set up Docker Buildx
uses: docker/setup-buildx-action@v3

- name: Login to us-central1 GAR
run: |
gcloud auth print-access-token --lifetime=7200 --impersonate-service-account=gha-ci-runners@matterlabs-infra.iam.gserviceaccount.com | docker login -u oauth2accesstoken --password-stdin https://us-docker.pkg.dev
- name: Login and push to Asia GAR
run: |
gcloud auth print-access-token --lifetime=7200 --impersonate-service-account=gha-ci-runners@matterlabs-infra.iam.gserviceaccount.com | docker login -u oauth2accesstoken --password-stdin https://asia-docker.pkg.dev
docker buildx imagetools create \
--tag asia-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ needs.build-images.outputs.protocol_version }}-${{ inputs.image_tag_suffix }} \
us-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ needs.build-images.outputs.protocol_version }}-${{ inputs.image_tag_suffix }}
- name: Login and push to Europe GAR
run: |
gcloud auth print-access-token --lifetime=7200 --impersonate-service-account=gha-ci-runners@matterlabs-infra.iam.gserviceaccount.com | docker login -u oauth2accesstoken --password-stdin https://europe-docker.pkg.dev
docker buildx imagetools create \
--tag europe-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ needs.build-images.outputs.protocol_version }}-${{ inputs.image_tag_suffix }} \
us-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ needs.build-images.outputs.protocol_version }}-${{ inputs.image_tag_suffix }}
15 changes: 15 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,21 @@ jobs:
DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }}
DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}

build-witness-generator-image-avx512:
name: Build prover images with avx512 instructions
needs: changed_files
if: ${{ (needs.changed_files.outputs.prover == 'true' || needs.changed_files.outputs.all == 'true') && !contains(github.ref_name, 'release-please--branches') }}
uses: ./.github/workflows/build-witness-generator-template.yml
with:
image_tag_suffix: ${{ needs.setup.outputs.image_tag_suffix }}-avx512
action: "build"
ERA_BELLMAN_CUDA_RELEASE: ${{ vars.ERA_BELLMAN_CUDA_RELEASE }}
is_pr_from_fork: ${{ github.event.pull_request.head.repo.fork == true }}
WITNESS_GENERATOR_RUST_FLAGS: "-Ctarget_feature=+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl"
secrets:
DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }}
DOCKERHUB_TOKEN: ${{ secrets.DOCKERHUB_TOKEN }}

ci-success:
name: Github Status Check
runs-on: ubuntu-latest
Expand Down
2 changes: 2 additions & 0 deletions docker/witness-generator/Dockerfile
Original file line number Diff line number Diff line change
@@ -1,6 +1,8 @@
FROM matterlabs/zksync-build-base:latest AS builder

ARG DEBIAN_FRONTEND=noninteractive
ARG RUST_FLAGS=""
ENV RUSTFLAGS=${RUST_FLAGS}

WORKDIR /usr/src/zksync
COPY . .
Expand Down
6 changes: 6 additions & 0 deletions infrastructure/zk/src/docker.ts
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,12 @@ async function _build(image: string, tagList: string[], dockerOrg: string, platf
const cudaArch = process.env.CUDA_ARCH;
buildArgs += `--build-arg CUDA_ARCH='${cudaArch}' `;
}
if (image === 'witness-generator') {
const rustFlags = process.env.RUST_FLAGS;
if (rustFlags) {
buildArgs += `--build-arg RUST_FLAGS='${rustFlags}' `;
}
}
buildArgs += extraArgs;

const buildCommand =
Expand Down

0 comments on commit 75851ff

Please sign in to comment.