From 00badcaa0eb0557704560c66ee0e81791f9275ef Mon Sep 17 00:00:00 2001 From: Yury Akudovich Date: Tue, 3 Dec 2024 10:48:19 +0100 Subject: [PATCH 1/3] feat(prover): Add cuobjdump output after prover builds (#3354) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What ❔ Add CUDA cuobjdump output on binary after build. ## Why ❔ To make sure that we build in all the needed GPU support. ## Checklist - [x] 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. - [x] Code has been formatted via `zkstack dev fmt` and `zkstack dev lint`. ref ZKD-2085 --- docker/circuit-prover-gpu/Dockerfile | 3 +++ docker/proof-fri-gpu-compressor/Dockerfile | 3 +++ docker/prover-gpu-fri/Dockerfile | 3 +++ 3 files changed, 9 insertions(+) diff --git a/docker/circuit-prover-gpu/Dockerfile b/docker/circuit-prover-gpu/Dockerfile index a5396e31ca88..852f42d705ec 100644 --- a/docker/circuit-prover-gpu/Dockerfile +++ b/docker/circuit-prover-gpu/Dockerfile @@ -45,6 +45,9 @@ COPY . . RUN cd prover && cargo build --release --bin zksync_circuit_prover +# Output build in CUDA architectures for debugging purposes. +RUN cuobjdump /usr/src/zksync/prover/target/release/zksync_circuit_prover + FROM nvidia/cuda:12.4.0-runtime-ubuntu22.04 RUN apt-get update && apt-get install -y curl libpq5 ca-certificates && rm -rf /var/lib/apt/lists/* diff --git a/docker/proof-fri-gpu-compressor/Dockerfile b/docker/proof-fri-gpu-compressor/Dockerfile index e744787c8259..6157bca0ac13 100644 --- a/docker/proof-fri-gpu-compressor/Dockerfile +++ b/docker/proof-fri-gpu-compressor/Dockerfile @@ -50,6 +50,9 @@ RUN cd prover && \ RUN cd prover && BELLMAN_CUDA_DIR=$PWD/bellman-cuda cargo build --features "gpu" --release --bin zksync_proof_fri_compressor +# Output build in CUDA architectures for debugging purposes. +RUN cuobjdump /usr/src/zksync/prover/target/release/zksync_proof_fri_compressor + FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 RUN apt-get update && apt-get install -y curl libpq5 ca-certificates && rm -rf /var/lib/apt/lists/* diff --git a/docker/prover-gpu-fri/Dockerfile b/docker/prover-gpu-fri/Dockerfile index 2a680a49c5de..db497fc5d500 100644 --- a/docker/prover-gpu-fri/Dockerfile +++ b/docker/prover-gpu-fri/Dockerfile @@ -45,6 +45,9 @@ COPY . . RUN cd prover && cargo build --release --features "gpu" --bin zksync_prover_fri +# Output build in CUDA architectures for debugging purposes. +RUN cuobjdump /usr/src/zksync/prover/target/release/zksync_prover_fri + FROM nvidia/cuda:12.2.0-runtime-ubuntu22.04 RUN apt-get update && apt-get install -y curl libpq5 ca-certificates && rm -rf /var/lib/apt/lists/* From de4b56dc133608cbfc48441a60950536cc3d3bf0 Mon Sep 17 00:00:00 2001 From: Yury Akudovich Date: Tue, 3 Dec 2024 17:32:58 +0100 Subject: [PATCH 2/3] ci: Remove unused CUDA_ARCH 60,70 (#3357) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What ❔ Remove unneeded CUDA_ARCH 60 and 70. Move needed CUDA_ARCH to default, so actual list will be in a single place after old CI cleanup. Remove CUDA_ARCH from witness-generator build, it isn't used there. ## Why ❔ To decrease build time and binary size. ## Checklist - [x] 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. - [x] Code has been formatted via `zkstack dev fmt` and `zkstack dev lint`. ref ZKD-2123 --- .github/workflows/build-docker-from-tag.yml | 16 +++++++--------- .github/workflows/build-prover-template.yml | 2 +- .../build-witness-generator-template.yml | 13 +++---------- .github/workflows/new-build-prover-template.yml | 7 +++++-- .../new-build-witness-generator-template.yml | 10 ++-------- .github/workflows/release-test-stage.yml | 2 -- 6 files changed, 18 insertions(+), 32 deletions(-) diff --git a/.github/workflows/build-docker-from-tag.yml b/.github/workflows/build-docker-from-tag.yml index b3f442ff4662..e48539c90738 100644 --- a/.github/workflows/build-docker-from-tag.yml +++ b/.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/new-build-core-template.yml if: contains(github.ref_name, 'core') secrets: @@ -61,7 +61,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: @@ -73,7 +73,7 @@ jobs: build-push-contract-verifier: name: Build and push image - needs: [ setup ] + needs: [setup] uses: ./.github/workflows/new-build-contract-verifier-template.yml if: contains(github.ref_name, 'contract_verifier') secrets: @@ -85,13 +85,12 @@ jobs: build-push-prover-images: name: Build and push image - needs: [ setup ] + needs: [setup] uses: ./.github/workflows/new-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;80;89" action: "push" secrets: DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }} @@ -99,13 +98,12 @@ jobs: build-push-witness-generator-image-avx512: name: Build and push image - needs: [ setup ] + needs: [setup] uses: ./.github/workflows/new-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;80;89" WITNESS_GENERATOR_RUST_FLAGS: "-Ctarget_feature=+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl" action: "push" secrets: @@ -114,7 +112,7 @@ jobs: build-gar-prover-fri-gpu-and-circuit-prover-gpu-gar: 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-and-circuit-prover-gpu-gar.yml if: contains(github.ref_name, 'prover') with: diff --git a/.github/workflows/build-prover-template.yml b/.github/workflows/build-prover-template.yml index 91de5dd51ecf..762ec496943c 100644 --- a/.github/workflows/build-prover-template.yml +++ b/.github/workflows/build-prover-template.yml @@ -30,7 +30,7 @@ on: CUDA_ARCH: description: "CUDA Arch to build" type: string - default: "89" + default: "75;80;89" required: false outputs: protocol_version: diff --git a/.github/workflows/build-witness-generator-template.yml b/.github/workflows/build-witness-generator-template.yml index d9493f97cae1..95053b89d3d8 100644 --- a/.github/workflows/build-witness-generator-template.yml +++ b/.github/workflows/build-witness-generator-template.yml @@ -27,11 +27,6 @@ on: 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 @@ -49,10 +44,9 @@ jobs: 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 }} ZKSYNC_USE_CUDA_STUBS: true - runs-on: [ matterlabs-ci-runner-c3d ] + runs-on: [matterlabs-ci-runner-c3d] strategy: matrix: component: @@ -91,7 +85,6 @@ jobs: 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: | @@ -162,11 +155,11 @@ jobs: DOCKER_ACTION: ${{ inputs.action }} COMPONENT: ${{ matrix.component }} run: | - PASSED_ENV_VARS="ERA_BELLMAN_CUDA_RELEASE,CUDA_ARCH,PROTOCOL_VERSION,RUST_FLAGS" \ + PASSED_ENV_VARS="ERA_BELLMAN_CUDA_RELEASE,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 \ No newline at end of file + ci_run cat /tmp/sccache_log.txt || true diff --git a/.github/workflows/new-build-prover-template.yml b/.github/workflows/new-build-prover-template.yml index cb254f602fc5..3a721e4425a8 100644 --- a/.github/workflows/new-build-prover-template.yml +++ b/.github/workflows/new-build-prover-template.yml @@ -30,8 +30,12 @@ on: CUDA_ARCH: description: "CUDA Arch to build" type: string - default: "89" + default: "75;80;89" required: false + # Details: https://arnon.dk/matching-sm-architectures-arch-and-gencode-for-various-nvidia-cards/ + # L4: 89 + # T4: 75 + # A100: 80 outputs: protocol_version: description: "Protocol version of the binary" @@ -210,7 +214,6 @@ jobs: --tag asia-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ inputs.image_tag_suffix }} \ us-docker.pkg.dev/matterlabs-infra/matterlabs-docker/${{ matrix.component }}:2.0-${{ 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 diff --git a/.github/workflows/new-build-witness-generator-template.yml b/.github/workflows/new-build-witness-generator-template.yml index bbd6aee23ed1..a96d217da832 100644 --- a/.github/workflows/new-build-witness-generator-template.yml +++ b/.github/workflows/new-build-witness-generator-template.yml @@ -21,11 +21,6 @@ on: type: string default: non-push 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 @@ -39,7 +34,7 @@ on: jobs: get-protocol-version: name: Get protocol version - runs-on: [ matterlabs-ci-runner-high-performance ] + runs-on: [matterlabs-ci-runner-high-performance] outputs: protocol_version: ${{ steps.protocolversion.outputs.protocol_version }} steps: @@ -85,7 +80,7 @@ jobs: needs: get-protocol-version env: PROTOCOL_VERSION: ${{ needs.get-protocol-version.outputs.protocol_version }} - runs-on: [ matterlabs-ci-runner-c3d ] + runs-on: [matterlabs-ci-runner-c3d] strategy: matrix: components: @@ -126,7 +121,6 @@ jobs: context: . push: ${{ inputs.action == 'push' }} build-args: | - CUDA_ARCH=${{ inputs.CUDA_ARCH }} SCCACHE_GCS_BUCKET=matterlabs-infra-sccache-storage SCCACHE_GCS_SERVICE_ACCOUNT=gha-ci-runners@matterlabs-infra.iam.gserviceaccount.com SCCACHE_GCS_RW_MODE=READ_WRITE diff --git a/.github/workflows/release-test-stage.yml b/.github/workflows/release-test-stage.yml index eb75ab179b8e..2e6c7882aa98 100644 --- a/.github/workflows/release-test-stage.yml +++ b/.github/workflows/release-test-stage.yml @@ -102,7 +102,6 @@ jobs: with: image_tag_suffix: ${{ needs.setup.outputs.image_tag_suffix }} ERA_BELLMAN_CUDA_RELEASE: ${{ vars.ERA_BELLMAN_CUDA_RELEASE }} - CUDA_ARCH: "60;70;75;80;89" action: "push" secrets: DOCKERHUB_USER: ${{ secrets.DOCKERHUB_USER }} @@ -116,7 +115,6 @@ jobs: 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;80;89" WITNESS_GENERATOR_RUST_FLAGS: "-Ctarget_feature=+avx512bw,+avx512cd,+avx512dq,+avx512f,+avx512vl " action: "push" secrets: From a7dc0ed5007f6b2f789f4c61cb3d137843151860 Mon Sep 17 00:00:00 2001 From: Patrick Date: Tue, 3 Dec 2024 17:37:41 +0100 Subject: [PATCH 3/3] fix(tee): fix race condition in batch locking (#3342) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## What ❔ After [scaling][1] [`zksync-tee-prover`][2] to two instances/replicas on Azure for `azure-stage2`, `azure-testnet2`, and `azure-mainnet2`, we started experiencing [duplicated proving for some batches][3]. ![logs](https://github.com/user-attachments/assets/39170805-d363-47bf-b275-468694593669) While this is not an erroneous situation, it is wasteful from a resource perspective. This was due to a race condition in batch locking. This PR fixes the issue by adding atomic batch locking. [1]: https://github.com/matter-labs/gitops-kubernetes/pull/7033/files [2]: https://github.com/matter-labs/zksync-era/blob/aaca32b6ab411d5cdc1234c20af8b5c1092195d7/core/bin/zksync_tee_prover/src/main.rs [3]: https://grafana.matterlabs.dev/goto/M1I_Bq7HR?orgId=1 ## Why ❔ To fix the bug that only activates after running `zksync-tee-prover` on multiple instances. ## Checklist - [x] 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. - [x] Code has been formatted via `zkstack dev fmt` and `zkstack dev lint`. --- ...82b0fa233913582fe9091cc1e8954dfd0eb1b.json | 30 ++++++ ...9f73e353da3bc6af7ecb81102c4194df631aa.json | 26 +++++ ...8203a62629904bc4956249e690a8ad7a48983.json | 32 ------ core/lib/dal/src/tee_proof_generation_dal.rs | 102 +++++++++++------- 4 files changed, 120 insertions(+), 70 deletions(-) create mode 100644 core/lib/dal/.sqlx/query-4777de5d3f313f1eb8c3b6a4c1782b0fa233913582fe9091cc1e8954dfd0eb1b.json create mode 100644 core/lib/dal/.sqlx/query-8ead57cdda5909348f31f8c4d989f73e353da3bc6af7ecb81102c4194df631aa.json delete mode 100644 core/lib/dal/.sqlx/query-e46c99b23db91800b27c717100f8203a62629904bc4956249e690a8ad7a48983.json diff --git a/core/lib/dal/.sqlx/query-4777de5d3f313f1eb8c3b6a4c1782b0fa233913582fe9091cc1e8954dfd0eb1b.json b/core/lib/dal/.sqlx/query-4777de5d3f313f1eb8c3b6a4c1782b0fa233913582fe9091cc1e8954dfd0eb1b.json new file mode 100644 index 000000000000..37adf92582e7 --- /dev/null +++ b/core/lib/dal/.sqlx/query-4777de5d3f313f1eb8c3b6a4c1782b0fa233913582fe9091cc1e8954dfd0eb1b.json @@ -0,0 +1,30 @@ +{ + "db_name": "PostgreSQL", + "query": "\n INSERT INTO\n tee_proof_generation_details (\n l1_batch_number, tee_type, status, created_at, updated_at, prover_taken_at\n )\n VALUES\n (\n $1,\n $2,\n $3,\n NOW(),\n NOW(),\n NOW()\n )\n ON CONFLICT (l1_batch_number, tee_type) DO\n UPDATE\n SET\n status = $3,\n updated_at = NOW(),\n prover_taken_at = NOW()\n RETURNING\n l1_batch_number,\n created_at\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "l1_batch_number", + "type_info": "Int8" + }, + { + "ordinal": 1, + "name": "created_at", + "type_info": "Timestamp" + } + ], + "parameters": { + "Left": [ + "Int8", + "Text", + "Text" + ] + }, + "nullable": [ + false, + false + ] + }, + "hash": "4777de5d3f313f1eb8c3b6a4c1782b0fa233913582fe9091cc1e8954dfd0eb1b" +} diff --git a/core/lib/dal/.sqlx/query-8ead57cdda5909348f31f8c4d989f73e353da3bc6af7ecb81102c4194df631aa.json b/core/lib/dal/.sqlx/query-8ead57cdda5909348f31f8c4d989f73e353da3bc6af7ecb81102c4194df631aa.json new file mode 100644 index 000000000000..6266f93e6545 --- /dev/null +++ b/core/lib/dal/.sqlx/query-8ead57cdda5909348f31f8c4d989f73e353da3bc6af7ecb81102c4194df631aa.json @@ -0,0 +1,26 @@ +{ + "db_name": "PostgreSQL", + "query": "\n SELECT\n p.l1_batch_number\n FROM\n proof_generation_details p\n LEFT JOIN\n tee_proof_generation_details tee\n ON\n p.l1_batch_number = tee.l1_batch_number\n AND tee.tee_type = $1\n WHERE\n (\n p.l1_batch_number >= $5\n AND p.vm_run_data_blob_url IS NOT NULL\n AND p.proof_gen_data_blob_url IS NOT NULL\n )\n AND (\n tee.l1_batch_number IS NULL\n OR (\n (tee.status = $2 OR tee.status = $3)\n AND tee.prover_taken_at < NOW() - $4::INTERVAL\n )\n )\n LIMIT 1\n FOR UPDATE OF p\n SKIP LOCKED\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "l1_batch_number", + "type_info": "Int8" + } + ], + "parameters": { + "Left": [ + "Text", + "Text", + "Text", + "Interval", + "Int8" + ] + }, + "nullable": [ + false + ] + }, + "hash": "8ead57cdda5909348f31f8c4d989f73e353da3bc6af7ecb81102c4194df631aa" +} diff --git a/core/lib/dal/.sqlx/query-e46c99b23db91800b27c717100f8203a62629904bc4956249e690a8ad7a48983.json b/core/lib/dal/.sqlx/query-e46c99b23db91800b27c717100f8203a62629904bc4956249e690a8ad7a48983.json deleted file mode 100644 index 7ca2c9e7e9fa..000000000000 --- a/core/lib/dal/.sqlx/query-e46c99b23db91800b27c717100f8203a62629904bc4956249e690a8ad7a48983.json +++ /dev/null @@ -1,32 +0,0 @@ -{ - "db_name": "PostgreSQL", - "query": "\n WITH upsert AS (\n SELECT\n p.l1_batch_number\n FROM\n proof_generation_details p\n LEFT JOIN\n tee_proof_generation_details tee\n ON\n p.l1_batch_number = tee.l1_batch_number\n AND tee.tee_type = $1\n WHERE\n (\n p.l1_batch_number >= $5\n AND p.vm_run_data_blob_url IS NOT NULL\n AND p.proof_gen_data_blob_url IS NOT NULL\n )\n AND (\n tee.l1_batch_number IS NULL\n OR (\n (tee.status = $2 OR tee.status = $3)\n AND tee.prover_taken_at < NOW() - $4::INTERVAL\n )\n )\n FETCH FIRST ROW ONLY\n )\n \n INSERT INTO\n tee_proof_generation_details (\n l1_batch_number, tee_type, status, created_at, updated_at, prover_taken_at\n )\n SELECT\n l1_batch_number,\n $1,\n $2,\n NOW(),\n NOW(),\n NOW()\n FROM\n upsert\n ON CONFLICT (l1_batch_number, tee_type) DO\n UPDATE\n SET\n status = $2,\n updated_at = NOW(),\n prover_taken_at = NOW()\n RETURNING\n l1_batch_number,\n created_at\n ", - "describe": { - "columns": [ - { - "ordinal": 0, - "name": "l1_batch_number", - "type_info": "Int8" - }, - { - "ordinal": 1, - "name": "created_at", - "type_info": "Timestamp" - } - ], - "parameters": { - "Left": [ - "Text", - "Text", - "Text", - "Interval", - "Int8" - ] - }, - "nullable": [ - false, - false - ] - }, - "hash": "e46c99b23db91800b27c717100f8203a62629904bc4956249e690a8ad7a48983" -} diff --git a/core/lib/dal/src/tee_proof_generation_dal.rs b/core/lib/dal/src/tee_proof_generation_dal.rs index e6b2df974b26..61a9e23ffea5 100644 --- a/core/lib/dal/src/tee_proof_generation_dal.rs +++ b/core/lib/dal/src/tee_proof_generation_dal.rs @@ -64,72 +64,98 @@ impl TeeProofGenerationDal<'_, '_> { ) -> DalResult> { let processing_timeout = pg_interval_from_duration(processing_timeout); let min_batch_number = i64::from(min_batch_number.0); + let mut transaction = self.storage.start_transaction().await?; + + // Lock rows in the proof_generation_details table to prevent race conditions. The + // tee_proof_generation_details table does not have corresponding entries yet if this is the + // first time the query is invoked for a batch. Locking rows in proof_generation_details + // ensures that two different TEE prover instances will not try to prove the same batch. + let batch_number = sqlx::query!( + r#" + SELECT + p.l1_batch_number + FROM + proof_generation_details p + LEFT JOIN + tee_proof_generation_details tee + ON + p.l1_batch_number = tee.l1_batch_number + AND tee.tee_type = $1 + WHERE + ( + p.l1_batch_number >= $5 + AND p.vm_run_data_blob_url IS NOT NULL + AND p.proof_gen_data_blob_url IS NOT NULL + ) + AND ( + tee.l1_batch_number IS NULL + OR ( + (tee.status = $2 OR tee.status = $3) + AND tee.prover_taken_at < NOW() - $4::INTERVAL + ) + ) + LIMIT 1 + FOR UPDATE OF p + SKIP LOCKED + "#, + tee_type.to_string(), + TeeProofGenerationJobStatus::PickedByProver.to_string(), + TeeProofGenerationJobStatus::Failed.to_string(), + processing_timeout, + min_batch_number + ) + .instrument("lock_batch_for_proving#get_batch_no") + .with_arg("tee_type", &tee_type) + .with_arg("processing_timeout", &processing_timeout) + .with_arg("min_batch_number", &min_batch_number) + .fetch_optional(&mut transaction) + .await?; + + let batch_number = match batch_number { + Some(batch) => batch.l1_batch_number, + None => { + return Ok(None); + } + }; + let locked_batch = sqlx::query_as!( StorageLockedBatch, r#" - WITH upsert AS ( - SELECT - p.l1_batch_number - FROM - proof_generation_details p - LEFT JOIN - tee_proof_generation_details tee - ON - p.l1_batch_number = tee.l1_batch_number - AND tee.tee_type = $1 - WHERE - ( - p.l1_batch_number >= $5 - AND p.vm_run_data_blob_url IS NOT NULL - AND p.proof_gen_data_blob_url IS NOT NULL - ) - AND ( - tee.l1_batch_number IS NULL - OR ( - (tee.status = $2 OR tee.status = $3) - AND tee.prover_taken_at < NOW() - $4::INTERVAL - ) - ) - FETCH FIRST ROW ONLY - ) - INSERT INTO tee_proof_generation_details ( l1_batch_number, tee_type, status, created_at, updated_at, prover_taken_at ) - SELECT - l1_batch_number, + VALUES + ( $1, $2, + $3, NOW(), NOW(), NOW() - FROM - upsert + ) ON CONFLICT (l1_batch_number, tee_type) DO UPDATE SET - status = $2, + status = $3, updated_at = NOW(), prover_taken_at = NOW() RETURNING l1_batch_number, created_at "#, + batch_number, tee_type.to_string(), TeeProofGenerationJobStatus::PickedByProver.to_string(), - TeeProofGenerationJobStatus::Failed.to_string(), - processing_timeout, - min_batch_number ) - .instrument("lock_batch_for_proving") + .instrument("lock_batch_for_proving#insert") + .with_arg("batch_number", &batch_number) .with_arg("tee_type", &tee_type) - .with_arg("processing_timeout", &processing_timeout) - .with_arg("l1_batch_number", &min_batch_number) - .fetch_optional(self.storage) + .fetch_optional(&mut transaction) .await? .map(Into::into); + transaction.commit().await?; Ok(locked_batch) }