optimize memory consumption for verifier subprocess #16201

dkijania · 2024-10-09T10:02:23Z

Implementation of https://www.notion.so/o1labs/RAM-savvy-Mina-node-114e79b1f91080968c47e815124a6279 regarding verifier/prover key duplication. Removed duplicated key generation in verifier process and reusing prover verification key (through GRPC) in verifier. Steps are pretty simple:

Start prover process
Let prover process generate the proving key and return the verification key back to main thread (via an RPC response)
Start a verifier process and initiailize Mina node with the verification key returned by the prover process

Testing:

Apart from full suite of integration tests i did tests locally (also in order to generate memory usage comparison before and after changed). For this purpose I used single node (seed without bp keys). Bootstrapped it and observed memory usage for mina and both sub-processes (verifier, prover). Before my changes prover and verifier were using proving keys in a lazy way. Currently prover works more eagerly, that's why mina takes up to 2 GB when bootstrapping, then GC frees memory leaving on only:

0.8 GB for prover
0.3 GB for main mina process
0.2 GB of verifier

To compare previously memory usage distribution was similar to:

0.2 GB for prover
0.3 GB for main mina process
0.2 GB of verifier

Before first usage of prover/verifier cryptographic material

Also what is worth mentioning that due to passing verification key from prover to verifier we gained ~ 15 seconds in bootstrap speed. In order to reach ~1 second we need to refactor :

Pickles.Side_loaded.srs_precomputation () ;

And bake those computation into rust binary. As Matthew mentioned , we can compute it at compile time and completely remove 10s form verifier boatload leaving only 1 sec.

Another test was to use local mina network in order to simulate small network method usage. I used network of:

1 seed
2 whales bp
2 fish bp
1 passive node
snark coordinator with 2 snark workers

Below some raw data of memory usage:
metrics_new.csv
metrics_compatible.csv

And some graphs:

Above graph show that verifiers processes are taking up to 0.9 GB after they are forced to be evaluated. Mina processes are taking around ~250 to 300 Mb of residential memory size. Only two provers (for mina whales) are taking up to 2.0 GB as they are asked to init calculation for produced blocks.

Laser focus on verifiers:

After optimization:

Here we can observe quite big spike of memory for seed (up to 1.7 GB) as it was first node to calculate proof data. Provers are acting quite the same as before change for active block producers. Inactive bp (seeds,fishes,passive nodes are taking up to 0.9 as they eagerly loaded proof data). This behavior could be tweak as well and as a result we could skip launching provers for non-block producers (seeds, archive node feeders) save ~ 1 GB of memory
The major change is that more main and verifier processes are taking around 250 MB.

More detailed look at verifiers processes:

And finally verifiers memory comparison:

georgeee · 2024-10-09T13:34:36Z

src/lib/transition_frontier/full_frontier/full_frontier.ml

@@ -960,12 +960,30 @@ module For_tests = struct
    let consensus_constants = precomputed_values.consensus_constants
  end

+  module T = Transaction_snark.Make (struct


Although it is located in the For_unit_tests module, I suspect it can be initialized even by the production code if it's somehow referenced by some function somewhere.

I'm not sure if making the module is an expensive operation, but I'd suggest:

Making it singleton for usage in unit tests (this would also reduce the code duplication)

Measure the time to initialize it, and if it's non-negligble, surround with some laziness

Previously initialization would happen from within a separate process, only when it's launched. Now it will happen if the module that contains this is somehow referenced by some code that is invoked in initialization, and it's possible that this is executed every time.

Once again, I'm not sure if making a functor is expensive itself or it is cheap, just generating a bunch of lazy values.

src/lib/verifier/prod.ml

georgeee · 2024-10-09T13:42:39Z

src/lib/verifier/prod.ml

-                 Lazy.force B.Proof.verification_key )
-
-             let get_blockchain_verification_key () = Lazy.force vk
+             let get_blockchain_verification_key () = Deferred.return verification_key


Maybe we could just remove this method?

oh yes, forgot to do it

georgeee · 2024-10-09T13:43:45Z

opam.export

@@ -21,6 +21,7 @@ roots: [
  "js_of_ocaml-toplevel.4.0.0"
  "lmdb.1.0"
  "menhir.20210419"
+  "memtrace_viewer.0.16.0"


Was this dependency used somehow?

only locally for memory trace. i will remove it

dkijania · 2024-10-10T18:34:23Z

src/lib/verifier/for_test.ml

I'm still unsure if this is purely for tests purposes. Maybe we should also reuse part of this method in prover ? then there will be a single place where we create BlockhainSnark and TransactionSnark modules. Prover bootstrap code is anyhow executed in mina main thread.

dkijania · 2024-10-10T18:38:17Z

src/lib/prover/prover.ml

@@ -56,6 +56,9 @@ module Worker_state = struct
    val toggle_internal_tracing : bool -> unit

    val set_itn_logger_data : daemon_port:int -> unit
+
+    val get_blockchain_verification_key :


I wonder also what to do in Check | No_check modes. Where we return dummy keys. Maybe this method should return Pickles.Verification_key.t Deferred.t but then in Mina_grahpql modules we are exposing Verifiaction key from verifier so we would need to either break the interface and also there wrap repsone in optional or handle returning diummy as well

Let's update graphql module to use key that is returned from prover. It's okay if it's dummy in case of non-full modes.

It's non-ideal that we create a dummy prover that returns dummy keys. I'd prefer it to be refactored to not initialize prover unless it's full mode. But let's save this idea for one of next PRs, not this one.

…n_bootstrap

dkijania · 2024-10-14T18:19:37Z

!ci-build-me

dkijania · 2024-10-14T18:19:43Z

!ci-nightly-me

dkijania · 2024-10-14T18:21:31Z

scripts/mina-local-network/monitor_memory.py

I'm still unsure if this is purely for tests purposes. Maybe we should also reuse part of this method in prover ? then there will be a single place where we create BlockhainSnark and TransactionSnark modules. Prover bootstrap code is anyhow executed in mina main thread.

dkijania · 2024-10-14T19:43:53Z

!ci-build-me

dkijania · 2024-10-14T20:41:50Z

!ci-nightly-me

mrmr1993 · 2024-10-14T20:49:00Z

src/lib/verifier/prod.ml

    ; commit_id : string
+    ; verification_key : Pickles.Verification_key.Stable.Latest.t


Please rename to blockchain_verification_key here and elsewhere.

mrmr1993 · 2024-10-14T20:49:42Z

src/lib/verifier/prod.ml

-               match Or_error.try_with (fun () -> T.verify ts) with
+               match
+                 Or_error.try_with (fun () ->
+                     Transaction_snark.verify ts ~key:verification_key )


This is incorrect. When we verify a transaction snark proof, we need to use the transaction snark verification key, not the blockchain snark verification key. You'll need to pipe that through.

…n_bootstrap

dkijania · 2024-10-15T17:56:08Z

!ci-build-me

dkijania · 2024-10-15T18:38:35Z

!ci-build-me

dkijania · 2024-10-16T07:59:08Z

!ci-nightly-me

…tion_on_bootstrap' into dkijania/optimize_memory_consumption_on_bootstrap

georgeee · 2024-10-21T09:23:04Z

scripts/mina-local-network/monitor_memory.py

@@ -0,0 +1,165 @@
+import collections
+import os
+import re


Let's introduce this script as a separate PR, it's not directly related to this PR

src/lib/verifier/for_test.mli

georgeee · 2024-10-21T09:34:23Z

src/lib/verifier/for_test.mli

@@ -0,0 +1,12 @@
+val get_verification_keys :


Let's remove get_verification_keys from for_test.mli, it's never used

georgeee · 2024-10-21T09:35:42Z

src/lib/verifier/prod.ml

@@ -191,7 +189,7 @@ module Worker_state = struct
               result

             let get_blockchain_verification_key () =
-               Lazy.force B.Proof.verification_key
+               Deferred.return blockchain_verification_key


Let's remove get_blockchain_verification_key from verifier and its interface, it's never used

georgeee · 2024-10-21T09:36:23Z

src/lib/verifier/dummy.ml

-let get_blockchain_verification_key { verification_key; _ } =
-  Deferred.Or_error.try_with ~here:[%here] (fun () ->
-      Lazy.force verification_key )
+let get_blockchain_verification_key { blockchain_verification_key; _ } =


Let's remove this

georgeee · 2024-10-21T09:42:02Z

src/lib/prover/prover.ml

@@ -56,6 +56,9 @@ module Worker_state = struct
    val toggle_internal_tracing : bool -> unit

    val set_itn_logger_data : daemon_port:int -> unit
+
+    val get_blockchain_verification_key :


Let's update graphql module to use key that is returned from prover. It's okay if it's dummy in case of non-full modes.

It's non-ideal that we create a dummy prover that returns dummy keys. I'd prefer it to be refactored to not initialize prover unless it's full mode. But let's save this idea for one of next PRs, not this one.

…n_bootstrap

Co-authored-by: George Agapov <[email protected]>

…tion_on_bootstrap' into dkijania/optimize_memory_consumption_on_bootstrap

dkijania · 2024-10-21T14:28:48Z

!ci-build-me

dkijania · 2024-10-21T14:50:36Z

!ci-build-me

dkijania · 2024-10-21T17:12:19Z

!ci-build-me

dkijania · 2024-10-21T20:51:22Z

!ci-nightly-me

coveralls · 2024-10-22T07:08:18Z

coverage: 60.999% (+0.01%) from 60.989%
when pulling cabd22b on dkijania/optimize_memory_consumption_on_bootstrap
into 4f78a48 on compatible.

dkijania · 2024-10-22T14:12:40Z

!ci-build-me

dkijania · 2024-10-22T14:13:28Z

Nightly run: https://buildkite.com/o-1-labs-2/mina-end-to-end-nightlies/builds/2784

georgeee

Approving with the last suggestion to simplify code

georgeee · 2024-10-25T09:15:43Z

src/lib/verifier/for_test.mli

+  -> proof_level:Genesis_constants.Proof_level.t
+  -> ( [ `Blockchain of Pickles.Verification_key.t ]
+     * [ `Transaction of Pickles.Verification_key.t ] )
+     Async.Deferred.t


Sorry for nit-picking last moment, but only now I noticed Async.Deferred.t at the end which is completely unnecessary here (we just return at the end).

Also, let's add a helper function for the following code:

let ( `Blockchain blockchain_verification_key , `Transaction transaction_verification_key ) = Verifier.For_test.get_verification_keys_eagerly ~constraint_constants ~proof_level in Verifier.create ~logger ~proof_level ~conf_dir:None ~pids:(Child_processes.Termination.create_pid_table ()) ~commit_id:"not specified for unit tests" () ~blockchain_verification_key ~transaction_verification_key )

The code is repeated many times, and helper would make the code more readable.

Great! will do

georgeee · 2024-10-25T09:15:50Z

src/lib/verifier/for_test.ml

+  let open Async.Deferred.Let_syntax in
+  let%bind blockchain_vk = Lazy.force B.Proof.verification_key in
+  let%bind transaction_vk = Lazy.force T.verification_key in
+  return (`Blockchain blockchain_vk, `Transaction transaction_vk)


Suggested change

return (`Blockchain blockchain_vk, `Transaction transaction_vk)

(`Blockchain blockchain_vk, `Transaction transaction_vk)

Hm, compiler is not happy when i removed return.

Matthew's concern was covered in the follow-up commits

dkijania · 2024-10-30T16:05:40Z

!ci-build-me

dkijania · 2024-10-30T19:52:06Z

!ci-nightly-me

georgeee reviewed Oct 9, 2024

View reviewed changes

dkijania commented Oct 10, 2024

View reviewed changes

dkijania force-pushed the dkijania/optimize_memory_consumption_on_bootstrap branch from dd5ef45 to 081b26d Compare October 10, 2024 19:24

dkijania added 3 commits October 14, 2024 15:51

Remove duplicated proof loading from verifier process

9cfeeea

Add script to dump memory usage for local network

5cd4d64

clean up code

5a5884d

dkijania force-pushed the dkijania/optimize_memory_consumption_on_bootstrap branch from 081b26d to 5a5884d Compare October 14, 2024 13:56

dkijania changed the title ~~[Do Not Merge] optimize memory consumption on bootstrap~~ optimize memory consumption on bootstrap for verifier Oct 14, 2024

Merge branch 'compatible' into dkijania/optimize_memory_consumption_o…

21aead5

…n_bootstrap

dkijania commented Oct 14, 2024

View reviewed changes

reformat

3ffbc4a

dkijania marked this pull request as ready for review October 14, 2024 20:41

dkijania requested review from a team as code owners October 14, 2024 20:41

dkijania changed the title ~~optimize memory consumption on bootstrap for verifier~~ optimize memory consumption on bootstrap for verifier subprocess Oct 14, 2024

dkijania changed the title ~~optimize memory consumption on bootstrap for verifier subprocess~~ optimize memory consumption for verifier subprocess Oct 14, 2024

mrmr1993 previously requested changes Oct 14, 2024

View reviewed changes

dkijania and others added 2 commits October 15, 2024 19:52

expse transaction_verification_key and use it for transaction snarks

73677a7

Merge branch 'compatible' into dkijania/optimize_memory_consumption_o…

2c0276b

…n_bootstrap

fix unit tests (verifier setup)

aa673fc

fix proof_level argument

b43535a

WIP

b1bc384

Merge remote-tracking branch 'origin/dkijania/optimize_memory_consump…

b027b7d

…tion_on_bootstrap' into dkijania/optimize_memory_consumption_on_bootstrap

georgeee requested changes Oct 21, 2024

View reviewed changes

Base automatically changed from georgeee/correct-snark-verify-functions to compatible October 21, 2024 11:56

dkijania and others added 5 commits October 21, 2024 15:19

Merge branch 'compatible' into dkijania/optimize_memory_consumption_o…

a749b66

…n_bootstrap

Apply suggestions from code review

0a30799

Co-authored-by: George Agapov <[email protected]>

remove get_verification_keys

6885222

Merge remote-tracking branch 'origin/dkijania/optimize_memory_consump…

804d1a4

…tion_on_bootstrap' into dkijania/optimize_memory_consumption_on_bootstrap

applied review comments

732619c

remove unused method expose prover in mina_lib

b567dc1

dkijania added 2 commits October 21, 2024 18:13

format

9122e38

fix compilation issues for mina_lib tests

a80029b

reformat

842b539

dkijania requested review from georgeee and mrmr1993 October 23, 2024 12:22

georgeee approved these changes Oct 25, 2024

View reviewed changes

move for_test module to sub module of verifier

cabd22b

dkijania self-assigned this Oct 30, 2024

georgeee approved these changes Nov 7, 2024

View reviewed changes

dkijania merged commit 4970213 into compatible Nov 7, 2024
65 checks passed

dkijania deleted the dkijania/optimize_memory_consumption_on_bootstrap branch November 7, 2024 21:33

		; commit_id : string
		; verification_key : Pickles.Verification_key.Stable.Latest.t

	return (`Blockchain blockchain_vk, `Transaction transaction_vk)
	(`Blockchain blockchain_vk, `Transaction transaction_vk)

optimize memory consumption for verifier subprocess #16201

optimize memory consumption for verifier subprocess #16201

Conversation

dkijania commented Oct 9, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dkijania commented Oct 14, 2024

dkijania commented Oct 14, 2024

Choose a reason for hiding this comment

dkijania commented Oct 14, 2024

dkijania commented Oct 14, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dkijania commented Oct 15, 2024

dkijania commented Oct 15, 2024

dkijania commented Oct 16, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dkijania commented Oct 21, 2024

dkijania commented Oct 21, 2024

dkijania commented Oct 21, 2024

dkijania commented Oct 21, 2024

coveralls commented Oct 22, 2024 • edited Loading

dkijania commented Oct 22, 2024

dkijania commented Oct 22, 2024

georgeee left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dkijania commented Oct 30, 2024

dkijania commented Oct 30, 2024

dkijania commented Oct 9, 2024 •

edited

Loading

coveralls commented Oct 22, 2024 •

edited

Loading