Merge pull request #12874 from openmina/feature/prover-internal-tracing

Internal tracing for prover and verifier (internal tracing PR 2 of 3)

src/app/archive/lib/test.ml

@@ -20,7 +20,8 @@ let%test_module "Archive node unit tests" =
       Async.Thread_safe.block_on_async_exn (fun () ->
           Verifier.create ~logger ~proof_level ~constraint_constants
             ~conf_dir:None
-            ~pids:(Child_processes.Termination.create_pid_table ()) )
+            ~pids:(Child_processes.Termination.create_pid_table ())
+            () )
 
     module Genesis_ledger = (val Genesis_ledger.for_unit_tests)
 

src/app/cli/src/cli_entrypoint/mina_cli_entrypoint.ml

@@ -1157,8 +1157,11 @@ let setup_daemon logger =
               ~default:Cli_lib.Default.stop_time stop_time
           in
           if enable_tracing then Mina_tracing.start conf_dir |> don't_wait_for ;
-          if enable_internal_tracing then
-            Internal_tracing.toggle ~logger `Enabled ;
+          let%bind () =
+            if enable_internal_tracing then
+              Internal_tracing.toggle ~logger `Enabled
+            else Deferred.unit
+          in
           let seed_peer_list_url =
             Option.value_map seed_peer_list_url ~f:Option.some
               ~default:
@@ -1615,7 +1618,7 @@ let internal_commands logger =
                      ~proof_level:Genesis_constants.Proof_level.compiled
                      ~constraint_constants:
                        Genesis_constants.Constraint_constants.compiled
-                     ~pids:(Pid.Table.create ()) ~conf_dir
+                     ~pids:(Pid.Table.create ()) ~conf_dir ()
                  in
                  Prover.prove_from_input_sexp prover sexp >>| ignore
              | `Eof ->
@@ -1713,7 +1716,7 @@ let internal_commands logger =
               ~proof_level:Genesis_constants.Proof_level.compiled
               ~constraint_constants:
                 Genesis_constants.Constraint_constants.compiled
-              ~pids:(Pid.Table.create ()) ~conf_dir:(Some conf_dir)
+              ~pids:(Pid.Table.create ()) ~conf_dir:(Some conf_dir) ()
           in
           let%bind result =
             match input with

src/app/cli/src/init/mina_run.ml

@@ -382,11 +382,9 @@ let setup_local_server ?(client_trustlist = []) ?rest_server_port
     ; implement Daemon_rpcs.Stop_tracing.rpc (fun () () ->
           Mina_tracing.stop () ; Deferred.unit )
     ; implement Daemon_rpcs.Start_internal_tracing.rpc (fun () () ->
-          Internal_tracing.toggle ~logger `Enabled ;
-          Deferred.unit )
+          Internal_tracing.toggle ~logger `Enabled )
     ; implement Daemon_rpcs.Stop_internal_tracing.rpc (fun () () ->
-          Internal_tracing.toggle ~logger `Disabled ;
-          Deferred.unit )
+          Internal_tracing.toggle ~logger `Disabled )
     ; implement Daemon_rpcs.Visualization.Frontier.rpc (fun () filename ->
           return (Mina_lib.visualize_frontier ~filename mina) )
     ; implement Daemon_rpcs.Visualization.Registered_masks.rpc

src/app/cli/src/init/transaction_snark_profiler.ml

@@ -16,7 +16,8 @@ let run ~user_command_profiler ~zkapp_profiler num_transactions ~max_num_updates
       Async.Thread_safe.block_on_async_exn (fun () ->
           Verifier.create ~logger ~proof_level ~constraint_constants
             ~conf_dir:None
-            ~pids:(Child_processes.Termination.create_pid_table ()) )
+            ~pids:(Child_processes.Termination.create_pid_table ())
+            () )
     in
     let rec go n =
       if n <= 0 then ()

src/lib/bootstrap_controller/bootstrap_controller.ml

@@ -693,7 +693,8 @@ let%test_module "Bootstrap_controller tests" =
       Async.Thread_safe.block_on_async_exn (fun () ->
           Verifier.create ~logger ~proof_level ~constraint_constants
             ~conf_dir:None
-            ~pids:(Child_processes.Termination.create_pid_table ()) )
+            ~pids:(Child_processes.Termination.create_pid_table ())
+            () )
 
     module Genesis_ledger = (val precomputed_values.genesis_ledger)
 

src/lib/crypto/kimchi_backend/common/dune

@@ -39,6 +39,8 @@
   allocation_functor
   snarky.intf
   promise
+  logger
+  internal_tracing.context_logger
   ppx_version.runtime))
 
 (rule

src/lib/crypto/kimchi_backend/common/plonk_dlog_proof.ml

@@ -396,14 +396,22 @@ module Make (Inputs : Inputs_intf) = struct
 
   let batch_verify' (conv : 'a -> Fq.t array)
       (ts : (Verifier_index.t * t * 'a * message option) list) =
+    let logger = Internal_tracing_context_logger.get () in
+    [%log internal] "Batch_verify_backend_convert_inputs" ;
     let vks_and_v =
       Array.of_list_map ts ~f:(fun (vk, t, xs, m) ->
           let p = to_backend' (Option.value ~default:[] m) (conv xs) t in
           (vk, p) )
     in
-    Backend.batch_verify
-      (Array.map ~f:fst vks_and_v)
-      (Array.map ~f:snd vks_and_v)
+    [%log internal] "Batch_verify_backend_convert_inputs_done" ;
+    [%log internal] "Batch_verify_backend" ;
+    let%map.Promise result =
+      Backend.batch_verify
+        (Array.map ~f:fst vks_and_v)
+        (Array.map ~f:snd vks_and_v)
+    in
+    [%log internal] "Batch_verify_backend_done" ;
+    result
 
   let batch_verify = batch_verify' (fun xs -> List.to_array xs)
 

src/lib/internal_tracing/context_call/dune

@@ -0,0 +1,16 @@
+(library
+ (name internal_tracing_context_call)
+ (public_name internal_tracing.context_call)
+ (library_flags -linkall)
+ (inline_tests)
+ (libraries
+  ;; opam libraries
+  base.base_internalhash_types
+  core_kernel
+  sexplib0
+  async_kernel)
+ (preprocess
+  (pps ppx_jane ppx_mina ppx_version ppx_deriving_yojson))
+ (instrumentation
+  (backend bisect_ppx))
+ (synopsis "Internal tracing context call ID helper"))

src/lib/internal_tracing/context_call/internal_tracing_context_call.ml

@@ -0,0 +1,13 @@
+open Core_kernel
+
+let last_call_id = ref 0
+
+let key = Univ_map.Key.create ~name:"call_id" sexp_of_int
+
+let with_call_id f =
+  incr last_call_id ;
+  Async_kernel.Async_kernel_scheduler.with_local key (Some !last_call_id) ~f
+
+let get_opt () = Async_kernel.Async_kernel_scheduler.find_local key
+
+let get () = Option.value (get_opt ()) ~default:0

src/lib/internal_tracing/context_call/internal_tracing_context_call.mli

@@ -0,0 +1,12 @@
+(** Helper module for setting the execution context call ID.
+
+    Useful for being able to detect context switches when
+    there are concurrent verifier/prover calls. *)
+
+(** [with_call_id f] runs [f] in a context in which
+    an unique identifier has been associated to the current call. *)
+val with_call_id : (unit -> 'a) -> 'a
+
+(** [get ()] returns current call ID bound to the current context (if any),
+    or 0 *)
+val get : unit -> int

src/lib/internal_tracing/context_logger/dune

@@ -0,0 +1,18 @@
+(library
+ (name internal_tracing_context_logger)
+ (public_name internal_tracing.context_logger)
+ (library_flags -linkall)
+ (inline_tests)
+ (libraries
+  ;; opam libraries
+  base.base_internalhash_types
+  core_kernel
+  sexplib0
+  async_kernel
+  ;; local libraries
+  logger)
+ (preprocess
+  (pps ppx_jane ppx_mina ppx_version ppx_deriving_yojson))
+ (instrumentation
+  (backend bisect_ppx))
+ (synopsis "Internal tracing context logger helper"))

src/lib/internal_tracing/context_logger/internal_tracing_context_logger.ml

@@ -0,0 +1,10 @@
+open Core_kernel
+
+let key = Univ_map.Key.create ~name:"logger" sexp_of_opaque
+
+let with_logger logger f =
+  Async_kernel.Async_kernel_scheduler.with_local key logger ~f
+
+let get_opt () = Async_kernel.Async_kernel_scheduler.find_local key
+
+let get () = Option.value (get_opt ()) ~default:(Logger.null ())

src/lib/internal_tracing/context_logger/internal_tracing_context_logger.mli

@@ -0,0 +1,12 @@
+(** Helper module for setting the execution context logger.
+
+    Useful for very deep call-stack where adding a logger parameter
+    to all the functions in the call-chain is too cumbersome. *)
+
+(** [with_logger logger f] runs [f] in a context in which
+    a call to [get] will return [logger]. *)
+val with_logger : Logger.t option -> (unit -> 'a) -> 'a
+
+(** [get ()] returns the logger bound to the current context (if any),
+    or the null logger otherwise. *)
+val get : unit -> Logger.t

src/lib/internal_tracing/dune

@@ -15,7 +15,10 @@
   logger
   mina_base
   mina_version
-  mina_numbers)
+  mina_numbers
+  internal_tracing.context_call
+  internal_tracing.context_logger
+  )
  (preprocess
   (pps ppx_jane ppx_mina ppx_version ppx_deriving_yojson))
  (instrumentation

src/lib/internal_tracing/internal_tracing.ml

@@ -17,10 +17,14 @@ let no_block = Block_id.Global_slot Mina_numbers.Global_slot.zero
 
 let last_block_id = ref no_block
 
+let last_call_id = ref 0
+
 let enabled = ref false
 
 let is_enabled () = !enabled
 
+let toggle_callbacks = ref []
+
 let disable_tracing ~logger () =
   [%log internal] "@internal_tracing_disabled" ;
   enabled := false
@@ -36,11 +40,24 @@ let enable_tracing ~logger () =
       ; ("branch", `String Mina_version.branch)
       ]
 
+let register_toggle_callback (f : bool -> unit Async_kernel.Deferred.t) =
+  toggle_callbacks := !toggle_callbacks @ [ f ]
+
+let call_toggle_callbacks enable =
+  Async_kernel.Deferred.List.iter
+    ~f:(fun cb -> cb enable)
+    ~how:`Parallel !toggle_callbacks
+
 let toggle ~logger ?(force = false) = function
   | `Enabled ->
-      if force || not (is_enabled ()) then enable_tracing ~logger ()
+      if force || not (is_enabled ()) then (
+        enable_tracing ~logger () ; call_toggle_callbacks true )
+      else Async_kernel.Deferred.unit
   | `Disabled ->
-      disable_tracing ~logger ()
+      if force || is_enabled () then (
+        disable_tracing ~logger () ;
+        call_toggle_callbacks false )
+      else Async_kernel.Deferred.unit
 
 let current_block_id_key =
   Univ_map.Key.create ~name:"current_block_id" Block_id.sexp_of_t
@@ -67,11 +84,22 @@ let handling_current_block_id_change' block_id =
     events that relate to a block to detect and record context switches
     made by the Async scheduler. *)
 let handling_current_block_id_change events =
-  match get_current_block_id () with
-  | None ->
-      handling_current_block_id_change' no_block @ events
-  | Some block_id ->
-      handling_current_block_id_change' block_id @ events
+  let block_id = Option.value ~default:no_block @@ get_current_block_id () in
+  handling_current_block_id_change' block_id @ events
+
+let handling_current_call_id_change' call_id =
+  if Int.equal !last_call_id call_id then []
+  else (
+    last_call_id := call_id ;
+    [ `Assoc [ ("current_call_id", `Int call_id) ] ] )
+
+(** If the current context has a different call ID from last time,
+    record a call ID context change. It is important to call this when generating
+    events that relate to different concurrent calls to detect and record
+    context switches made by the Async scheduler. *)
+let handling_current_call_id_change events =
+  let call_id = Internal_tracing_context_call.get () in
+  handling_current_call_id_change' call_id @ events
 
 module For_logger = struct
   module Json_lines_rotate_transport = struct
@@ -131,8 +159,9 @@ module For_logger = struct
         let json_lines : Yojson.Safe.t list =
           match message with
           | "@metadata" ->
-              handling_current_block_id_change
-                [ Event.checkpoint_metadata metadata ]
+              handling_current_call_id_change
+              @@ handling_current_block_id_change
+                   [ Event.checkpoint_metadata metadata ]
           | "@block_metadata" ->
               handling_current_block_id_change [ Event.block_metadata metadata ]
           | "@internal_tracing_enabled" ->
@@ -152,8 +181,9 @@ module For_logger = struct
                 if String.Map.is_empty metadata then []
                 else [ Event.checkpoint_metadata metadata ]
               in
-              handling_current_block_id_change
-                (checkpoint_event_json :: metadata_event_json)
+              handling_current_call_id_change
+              @@ handling_current_block_id_change
+                   (checkpoint_event_json :: metadata_event_json)
         in
         Some
           ( String.concat ~sep:"\n"
@@ -171,3 +201,6 @@ module For_logger = struct
 
   let processor = Logger.Processor.create (module Processor) ()
 end
+
+module Context_logger = Internal_tracing_context_logger
+module Context_call = Internal_tracing_context_call

src/lib/internal_tracing/internal_tracing.mli

@@ -57,6 +57,8 @@
 
      - ["@current_block"]: used to notify of a execution context change that
        brings a different block into context.
+     - ["@current_call_id"]: used to notify of a execution context change that
+       brings a different concurrent verifier or prover call into context.
      - ["@internal_tracing_enabled"]: issued whenever internal tracing is enabled.
      - ["@internal_tracing_disabled"]: issued whenever internal tracing is disabled.
      - ["@mina_node_metadata"]: associates global metadata about the current node
@@ -90,18 +92,36 @@
 (** [is_enabled ()] returns [true] if internal tracing is enabled, and [false] otherwise. *)
 val is_enabled : unit -> bool
 
+(** [register_toggle_callback callback] will register [callback] to be called whenever
+    internal tracing is toggled.
+
+    This is useful to synchronize internal tracing done by subprocesses like the verifier
+    and prover.
+
+    [callback] will be called with [true] if internal tracing must be enabled, and with
+    [false] if it must be disabled. It must return a [Deferred.t] that will be resolved
+    once the call completes. *)
+val register_toggle_callback : (bool -> unit Async_kernel.Deferred.t) -> unit
+
 (** [toggle `Enabled] will enable tracing.
     [toggle `Disabled] will disable tracing.
 
     If [force] is [false] (the default), and if tracing is already active,
-    then trying to enable tracing is a noop. *)
-val toggle : logger:Logger.t -> ?force:bool -> [ `Enabled | `Disabled ] -> unit
+    then trying to enable tracing is a noop.
+
+    The returned promise will be resolved when all the calls to the registered toggle
+    callbacks have been resolved. *)
+val toggle :
+     logger:Logger.t
+  -> ?force:bool
+  -> [ `Enabled | `Disabled ]
+  -> unit Async_kernel.Deferred.t
 
 (** [with_state_hash state_hash f] runs [f] in a context in which checkpoints
     and metadata will be associated to a block with state hash equal to [state_hash].
 
     Any context in which checkpoints or metadata are recorded must be wrapped
-    in either [with_state_hash] or [with_slot] for the checkpoints to be
+    in either {!val:with_state_hash} or {!val:with_slot} for the checkpoints to be
     properly associated to the block being processed/produced. *)
 val with_state_hash : Mina_base.State_hash.t -> (unit -> 'a) -> 'a
 
@@ -110,7 +130,7 @@ val with_state_hash : Mina_base.State_hash.t -> (unit -> 'a) -> 'a
     equal to [global_slot].
 
     Any context in which checkpoints or metadata are recorded must be wrapped
-    in either [with_state_hash] or [with_slot] for the checkpoints to be
+    in either {!val:with_state_hash} or {!val:with_slot} for the checkpoints to be
     properly associated to the block being processed/produced. *)
 val with_slot : Mina_numbers.Global_slot.t -> (unit -> 'a) -> 'a
 
@@ -128,3 +148,7 @@ module For_logger : sig
   (** Processor for the "internal" log level used to record checkpoints *)
   val processor : Logger.Processor.t
 end
+
+module Context_logger : module type of Internal_tracing_context_logger
+
+module Context_call : module type of Internal_tracing_context_call

src/lib/ledger_catchup/normal_catchup.ml

@@ -878,7 +878,8 @@ let%test_module "Ledger_catchup tests" =
       Async.Thread_safe.block_on_async_exn (fun () ->
           Verifier.create ~logger ~proof_level ~constraint_constants
             ~conf_dir:None
-            ~pids:(Child_processes.Termination.create_pid_table ()) )
+            ~pids:(Child_processes.Termination.create_pid_table ())
+            () )
 
     module Context = struct
       let logger = logger

src/lib/ledger_catchup/super_catchup.ml

@@ -1408,7 +1408,8 @@ let%test_module "Ledger_catchup tests" =
       Async.Thread_safe.block_on_async_exn (fun () ->
           Verifier.create ~logger ~proof_level ~constraint_constants
             ~conf_dir:None
-            ~pids:(Child_processes.Termination.create_pid_table ()) )
+            ~pids:(Child_processes.Termination.create_pid_table ())
+            () )
 
     module Context = struct
       let logger = logger