Merge pull request #17802 from MinaProtocol/lyh/simplify-wire-data-snark-worker

Update SNARK worker & minimize data pased on wire - compatible

Author: glyh

changes/17724.md

@@ -0,0 +1,19 @@
+## Abstract
+Updated SNARK worker protocol. Now more parallelism can be utilized when constructing proofs for a zkApp command. Redeployment of SNARK worker and SNARK coordinator(i.e. any daemon that's attached to a SNARK worker) would be needed for this to take effect. 
+
+## Metrics
+The metrics related to SNARK workers has been updated as well.
+
+Following metrics are maintained by the worker:
+- `snark_worker_sub_zkapp_command_segment_time`
+- `snark_worker_sub_zkapp_command_merge_time`
+- `snark_worker_merge_time`
+- `snark_worker_nonzkapp_transition_time`
+
+addtionally, the coordinator maintains all metrics above, plus an extra metric:
+- `snark_worker_zkapp_transition_time`
+
+While having same name, these shared metrics have different meanings across coordinator and worker: Any time tracked by a worker represents the time this specific worker spent on that genre, while any time tracked by a coordinator represents the sum of time any worker attached to that coordinator spent on that genre.
+
+## Foot Note
+This is part of the Slot Reduction MIP from Mesa Hardfork.

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

@@ -380,56 +380,41 @@ let setup_local_server ?(client_trustlist = []) ?rest_server_port
   in
   let snark_worker_impls =
     [ implement Snark_worker.Rpcs_versioned.Get_work.Latest.rpc (fun () () ->
-          Deferred.return
-            (let open Option.Let_syntax in
-            let%bind key =
-              Option.merge
-                (Mina_lib.snark_worker_key mina)
-                (Mina_lib.snark_coordinator_key mina)
-                ~f:Fn.const
-            in
-            let%map work = Mina_lib.request_work mina in
-            let work =
-              Snark_work_lib.Work.Spec.map work
-                ~f:
-                  (Snark_work_lib.Work.Single.Spec.map
-                     ~f_proof:Ledger_proof.Cached.read_proof_from_disk
-                     ~f_witness:Transaction_witness.read_all_proofs_from_disk )
-            in
-            [%log trace]
-              ~metadata:
-                [ ( "work_spec"
-                  , Snark_work_lib.Selector.Spec.Stable.Latest.to_yojson work )
-                ]
-              "responding to a Get_work request with some new work" ;
-            Mina_metrics.(Counter.inc_one Snark_work.snark_work_assigned_rpc) ;
-            (work, key)) )
-    ; implement Snark_worker.Rpcs_versioned.Submit_work.Latest.rpc
-        (fun () (work : Snark_work_lib.Selector.Result.Stable.Latest.t) ->
-          [%log trace] "received completed work from a snark worker"
-            ~metadata:
-              [ ( "work_spec"
-                , Snark_work_lib.Selector.Spec.Stable.Latest.to_yojson work.spec
-                )
-              ] ;
+          match Mina_lib.request_work mina with
+          | None ->
+              Deferred.return None
+          | Some (Ok spec) ->
+              [%log debug] "responding to a Get_work request with some new work"
+                ~metadata:
+                  [ ( "work_id"
+                    , Snark_work_lib.(
+                        Spec.Partitioned.Poly.get_id spec |> Id.Any.to_yojson)
+                    )
+                  ] ;
 
+              Mina_metrics.(Counter.inc_one Snark_work.snark_work_assigned_rpc) ;
+              Deferred.return (Some spec)
+          | Some (Error (`Failed_to_generate_inputs (zkapp_cmd, e))) ->
+              let open Mina_base.Zkapp_command in
+              [%log error]
+                "Mina_lib.request_work failed to generate inputs for a zkapp \
+                 command"
+                ~metadata:
+                  [ ("error", `String (Error.to_string_hum e))
+                  ; ( "zkapp_cmd"
+                    , Stable.Latest.to_yojson
+                      @@ read_all_proofs_from_disk zkapp_cmd )
+                  ] ;
+              Deferred.return None )
+    ; implement Snark_worker.Rpcs_versioned.Submit_work.Latest.rpc
+        (fun () (result : Snark_work_lib.Result.Partitioned.Stable.Latest.t) ->
+          [%log debug] "received completed work from a snark worker"
+            ~metadata:[ ("work_id", Snark_work_lib.Id.Any.to_yojson result.id) ] ;
           Mina_metrics.(
             Counter.inc_one Snark_work.completed_snark_work_received_rpc) ;
-          One_or_two.zip_exn work.spec.instances work.metrics
-          |> One_or_two.iter ~f:(fun (single_spec, (elapsed, _tag)) ->
-                 Snark_work_lib.Metrics.(
-                   emit_single_metrics ~logger ~single_spec
-                     ~data:{ data = elapsed; proof = () }
-                     ~on_zkapp_command:emit_zkapp_metrics_legacy ()) ) ;
-          Deferred.return @@ Mina_lib.add_work mina work )
+          Deferred.return @@ Mina_lib.add_work mina result )
     ; implement Snark_worker.Rpcs_versioned.Failed_to_generate_snark.Latest.rpc
-        (fun
-          ()
-          ((error, _work_spec, _prover_public_key) :
-            Error.t
-            * Snark_work_lib.Selector.Spec.Stable.Latest.t
-            * Signature_lib.Public_key.Compressed.t )
-        ->
+        (fun () (error, _) ->
           [%str_log error]
             (Snark_worker.Events.Generating_snark_work_failed
                { error = Error_json.error_to_yojson error } ) ;

src/lib/mina_graphql/mina_graphql.ml

@@ -2351,7 +2351,7 @@ module Queries = struct
       ~args:Arg.[]
       ~typ:(non_null @@ list @@ non_null Types.pending_work)
       ~resolve:(fun { ctx = mina; _ } () ->
-        let snark_job_state = Mina_lib.snark_job_state mina in
+        let snark_job_state = Mina_lib.work_selector mina in
         let snark_pool = Mina_lib.snark_pool mina in
         let fee_opt =
           Mina_lib.(
@@ -2381,7 +2381,7 @@ module Queries = struct
           ]
       ~typ:(non_null @@ list @@ non_null Types.pending_work_spec)
       ~resolve:(fun { ctx = mina; _ } () start_idx end_idx ->
-        let snark_job_state = Mina_lib.snark_job_state mina in
+        let snark_job_state = Mina_lib.work_selector mina in
         let snark_pool = Mina_lib.snark_pool mina in
         let all_work = Work_selector.all_work ~snark_pool snark_job_state in
         let work_size = all_work |> List.length |> Unsigned.UInt32.of_int in

src/lib/mina_lib/mina_lib.ml

@@ -107,7 +107,8 @@ type t =
   ; pipes : pipes
   ; wallets : Secrets.Wallets.t
   ; coinbase_receiver : Consensus.Coinbase_receiver.t ref
-  ; snark_job_state : Work_selector.State.t
+  ; work_selector : Work_selector.State.t
+  ; work_partitioner : Work_partitioner.t
   ; mutable next_producer_timing :
       Daemon_rpcs.Types.Status.Next_producer_timing.t option
   ; subscriptions : Mina_subscriptions.t
@@ -796,7 +797,7 @@ let get_inferred_nonce_from_transaction_pool_and_ledger t
       let%map account = get_account t account_id in
       account.Account.nonce
 
-let snark_job_state t = t.snark_job_state
+let work_selector t = t.work_selector
 
 let add_block_subscriber t public_key =
   Mina_subscriptions.add_block_subscriber t.subscriptions public_key
@@ -869,46 +870,104 @@ let best_chain ?max_length t =
 
 let request_work t =
   let (module Work_selection_method) = t.config.work_selection_method in
+  let%bind.Option prover =
+    Option.first_some (snark_worker_key t) (snark_coordinator_key t)
+  in
   let fee = snark_work_fee t in
-  let instances_opt =
-    Work_selection_method.work ~logger:t.config.logger ~fee
-      ~snark_pool:(snark_pool t) (snark_job_state t)
+  let sok_message = Sok_message.create ~fee ~prover in
+  [%log' debug t.config.logger] "Received work request"
+    ~metadata:[ ("sok_message", Sok_message.to_yojson sok_message) ] ;
+  let work_from_selector =
+    lazy
+      (Work_selection_method.work ~snark_pool:(snark_pool t) ~fee
+         ~logger:t.config.logger t.work_selector )
   in
-  Option.map instances_opt ~f:(fun instances ->
-      { Snark_work_lib.Work.Spec.instances; fee } )
+  Work_partitioner.request_partitioned_work ~work_from_selector ~sok_message
+    ~partitioner:t.work_partitioner
 
 let work_selection_method t = t.config.work_selection_method
 
-let add_work t (work : Snark_work_lib.Selector.Result.Stable.Latest.t) =
+let add_complete_work ~logger ~fee ~prover
+    ~(results :
+       ( Snark_work_lib.Spec.Single.t
+       , Ledger_proof.Cached.t )
+       Snark_work_lib.Result.Single.Poly.t
+       One_or_two.t ) t =
   let update_metrics () =
     let snark_pool = snark_pool t in
     let fee_opt =
       Option.map (snark_worker_key t) ~f:(fun _ -> snark_work_fee t)
     in
     let pending_work =
-      Work_selector.pending_work_statements ~snark_pool ~fee_opt
-        t.snark_job_state
+      Work_selector.pending_work_statements ~snark_pool ~fee_opt t.work_selector
       |> List.length
     in
     Mina_metrics.(
       Gauge.set Snark_work.pending_snark_work (Int.to_float pending_work))
   in
-  let spec =
-    One_or_two.map work.spec.instances
-      ~f:Snark_work_lib.Work.Single.Spec.statement
+  let proofs = One_or_two.map ~f:(fun { proof; _ } -> proof) results in
+  let fee_with_prover = Fee_with_prover.{ fee; prover } in
+  let stmts =
+    One_or_two.map
+      ~f:(fun { spec; _ } -> Snark_work_lib.Spec.Single.Poly.statement spec)
+      results
   in
-  let cb _ =
-    (* remove it from seen jobs after attempting to adding it to the pool to avoid this work being reassigned
+  [%log debug] "Partitioner combined work"
+    ~metadata:
+      [ ("work_ids", Transaction_snark_work.Statement.compact_json stmts)
+      ; ("fee_with_prover", Fee_with_prover.to_yojson fee_with_prover)
+      ] ;
+  Or_error.try_with update_metrics
+  |> Result.iter_error ~f:(fun err ->
+         [%log warn] "Failed to update metrics on adding work"
+           ~metadata:[ ("error", `String (Error.to_string_hum err)) ] ) ;
+  let cb result =
+    (* remove it from seen jobs after attempting to adding it to the pool to
+           avoid this work being reassigned.
      * If the diff is accepted then remove it from the seen jobs.
-     * If not then the work should have already been in the pool with a lower fee or the statement isn't referenced anymore or any other error. In any case remove it from the seen jobs so that it can be picked up if needed *)
-    Work_selector.remove t.snark_job_state spec
+     * If not then the work should have already been in the pool with a 
+           lower fee or the statement isn't referenced anymore or any other 
+           error. In any case remove it from the seen jobs so that it can be 
+           picked up if needed *)
+    Work_selector.remove t.work_selector stmts ;
+    Result.iter_error result ~f:(fun err ->
+        (* Possible reasons of failure: receiving pipe's capacity exceeded,
+            fee that isn't the lowest, failure in verification or application to the pool *)
+        [%log warn] "Failed to add completed work to the pool"
+          ~metadata:
+            [ ("work_ids", Transaction_snark_work.Statement.compact_json stmts)
+            ; ("error", `String (Error.to_string_hum err))
+            ] )
   in
-  ignore (Or_error.try_with (fun () -> update_metrics ()) : unit Or_error.t) ;
   Network_pool.Snark_pool.(
     Local_sink.push t.pipes.snark_local_sink
-      (Resource_pool.Diff.of_result work, cb))
+      ( Add_solved_work
+          ( stmts
+          , Network_pool.Priced_proof.
+              { proof =
+                  proofs
+                  |> One_or_two.map ~f:Ledger_proof.Cached.read_proof_from_disk
+              ; fee = fee_with_prover
+              } )
+      , cb ))
   |> Deferred.don't_wait_for
 
+let add_work t (work : Snark_work_lib.Result.Partitioned.Stable.Latest.t) =
+  let logger = t.config.logger in
+  match
+    Work_partitioner.submit_partitioned_work ~result:work
+      ~partitioner:t.work_partitioner
+  with
+  | SpecUnmatched ->
+      `SpecUnmatched
+  | Removed ->
+      `Removed
+  | Processed None ->
+      `Ok
+  | Processed (Some { results; fee; prover }) ->
+      add_complete_work ~logger ~fee ~prover ~results t ;
+      `Ok
+
 let add_work_graphql t diff =
   let results_ivar = Ivar.create () in
   Network_pool.Snark_pool.Local_sink.push t.pipes.snark_local_sink
@@ -2076,12 +2135,20 @@ let create ~commit_id ?wallets (config : Config.t) =
               ; constraint_constants
               }
           in
-          let snark_jobs_state =
+          let work_selector =
             Work_selector.State.init
               ~reassignment_wait:config.work_reassignment_wait
               ~frontier_broadcast_pipe:frontier_broadcast_pipe_r
               ~logger:config.logger
           in
+          (* NOTE: [reassignment_wait] is interpreted as milliseconds *)
+          let work_partitioner =
+            Work_partitioner.create
+              ~signature_kind:Mina_signature_kind.t_DEPRECATED
+              ~reassignment_timeout:
+                (Time.Span.of_ms (Float.of_int config.work_reassignment_wait))
+              ~logger:config.logger ~proof_cache_db
+          in
           let sinks = (block_sink, tx_remote_sink, snark_remote_sink) in
           let%bind net =
             O1trace.thread "mina_networking" (fun () ->
@@ -2091,7 +2158,7 @@ let create ~commit_id ?wallets (config : Config.t) =
                   ~get_transition_frontier:(fun () ->
                     Broadcast_pipe.Reader.peek frontier_broadcast_pipe_r )
                   ~get_snark_pool:(fun () -> Some snark_pool)
-                  ~snark_job_state:(fun () -> Some snark_jobs_state)
+                  ~snark_job_state:(fun () -> Some work_selector)
                   ~get_node_status )
           in
           let user_command_input_reader, user_command_input_writer =
@@ -2515,7 +2582,8 @@ let create ~commit_id ?wallets (config : Config.t) =
               }
           ; wallets
           ; coinbase_receiver = ref config.coinbase_receiver
-          ; snark_job_state = snark_jobs_state
+          ; work_selector
+          ; work_partitioner
           ; subscriptions
           ; sync_status
           ; precomputed_block_writer

src/lib/mina_lib/mina_lib.mli

@@ -95,6 +95,10 @@ val current_epoch_delegators :
 val last_epoch_delegators :
   t -> pk:Public_key.Compressed.t -> Mina_base.Account.t list option
 
+(** [replace_snark_worker_key t key_opt] Replace all SNARK worker's key
+    associated with current coordinator.
+    - If the new key is [None], SNARK worker will be turned off if it's running;
+    - If the new key is [Some k], SNARK worker will be turn on if it's not running. *)
 val replace_snark_worker_key :
   t -> Public_key.Compressed.t option -> unit Deferred.t
 
@@ -113,11 +117,19 @@ val snark_work_fee : t -> Currency.Fee.t
 
 val set_snark_work_fee : t -> Currency.Fee.t -> unit
 
-val request_work : t -> Work_selector.work Snark_work_lib.Work.Spec.t option
+val request_work :
+     t
+  -> ( Snark_work_lib.Spec.Partitioned.Stable.Latest.t
+     , Work_partitioner.Snark_worker_shared.Failed_to_generate_inputs.t )
+     Result.t
+     option
 
 val work_selection_method : t -> (module Work_selector.Selection_method_intf)
 
-val add_work : t -> Snark_work_lib.Selector.Result.Stable.Latest.t -> unit
+val add_work :
+     t
+  -> Snark_work_lib.Result.Partitioned.Stable.Latest.t
+  -> [> `Ok | `Removed | `SpecUnmatched ]
 
 val add_work_graphql :
      t
@@ -127,7 +139,7 @@ val add_work_graphql :
      * Network_pool.Snark_pool.Resource_pool.Diff.rejected )
      Deferred.Or_error.t
 
-val snark_job_state : t -> Work_selector.State.t
+val work_selector : t -> Work_selector.State.t
 
 val get_current_nonce :
      t

src/lib/snark_work_lib/combined_result.mli

@@ -1,5 +1,6 @@
 type t =
-  { spec_with_proof : (Single_spec.t * Ledger_proof.t) One_or_two.t
+  { results :
+      (Single_spec.t, Ledger_proof.Cached.t) Single_result.Poly.t One_or_two.t
   ; fee : Currency.Fee.t
   ; prover : Signature_lib.Public_key.Compressed.t
   }

src/lib/snark_work_lib/id.ml

@@ -53,3 +53,17 @@ module Sub_zkapp = struct
   let to_single ({ which_one; pairing_id; _ } : t) : Single.t =
     { which_one; pairing_id }
 end
+
+module Any = struct
+  [%%versioned
+  module Stable = struct
+    module V1 = struct
+      type t =
+        | Single of Single.Stable.V1.t
+        | Sub_zkapp of Sub_zkapp.Stable.V1.t
+      [@@deriving compare, hash, sexp, yojson, equal]
+
+      let to_latest = Fn.id
+    end
+  end]
+end

src/lib/snark_work_lib/id.mli

@@ -51,3 +51,17 @@ module Sub_zkapp : sig
 
   val to_single : t -> Single.t
 end
+
+module Any : sig
+  [%%versioned:
+  module Stable : sig
+    module V1 : sig
+      type t =
+        | Single of Single.Stable.V1.t
+        | Sub_zkapp of Sub_zkapp.Stable.V1.t
+      [@@deriving compare, hash, sexp, yojson, equal]
+
+      val to_latest : t -> t
+    end
+  end]
+end