fix(coprocessor): maximise dependence fetch from DB (#1950) (#1961)

* fix(coprocessor): maximise dependence fetch from DB  (#1950)

fix(coprocessor): maximise fetched dependence ciphertexts from DB

* chore(coprocessor): update time and bytes package versions

Author: antoniupop

coprocessor/fhevm-engine/Cargo.lock

@@ -312,6 +312,7 @@ pub struct DFComponentGraph {
     pub needed_map: HashMap<Handle, Vec<NodeIndex>>,
     pub produced: HashMap<Handle, Vec<(NodeIndex, Handle)>>,
     pub results: Vec<DFGTxResult>,
+    deferred_dependences: Vec<(NodeIndex, NodeIndex, Handle)>,
 }
 impl DFComponentGraph {
     pub fn build(&mut self, nodes: &mut Vec<ComponentNode>) -> Result<()> {
@@ -350,7 +351,16 @@ impl DFComponentGraph {
                         error!(target: "scheduler", { output_handle = ?hex::encode(i.clone()) },
 				   "Missing producer for handle");
                     } else {
-                        dependence_pairs.push((producer[0].0, consumer));
+                        // Cross-transaction dependence: defer until
+                        // after DB fetch. If the handle is found in
+                        // DB, we use the fetched value and skip the
+                        // dependence edge.
+                        self.deferred_dependences
+                            .push((producer[0].0, consumer, i.clone()));
+                        self.needed_map
+                            .entry(i.clone())
+                            .and_modify(|uses| uses.push(consumer))
+                            .or_insert(vec![consumer]);
                     }
                 } else {
                     self.needed_map
@@ -361,19 +371,52 @@ impl DFComponentGraph {
             }
         }
 
-        // We build a replica of the graph and map it to the
-        // underlying DiGraph so we can identify cycles.
-        let mut digraph = self.graph.map(|idx, _| idx, |_, _| ()).graph().clone();
-        // Add transaction dependence edges
+        // Same-transaction dependences are always acyclic (they
+        // derive from the transaction's internal DAG). Add them
+        // directly; cycle detection runs once in
+        // resolve_dependences() over the full edge set.
         for (producer, consumer) in dependence_pairs.iter() {
+            if self.graph.add_edge(*producer, *consumer, ()).is_err() {
+                let prod = self
+                    .graph
+                    .node_weight(*producer)
+                    .ok_or(SchedulerError::DataflowGraphError)?;
+                let cons = self
+                    .graph
+                    .node_weight(*consumer)
+                    .ok_or(SchedulerError::DataflowGraphError)?;
+                error!(target: "scheduler", { producer_id = ?hex::encode(prod.transaction_id.clone()), consumer_id = ?hex::encode(cons.transaction_id.clone()) },
+		       "Unexpected cycle in same-transaction dependence");
+                return Err(SchedulerError::CyclicDependence.into());
+            }
+        }
+        Ok(())
+    }
+
+    // Resolve deferred cross-transaction dependences after DB fetch.
+    // Dependences whose handle was successfully fetched are dropped
+    // (the consumer already has the data). Remaining dependences are
+    // added as graph edges after cycle detection.
+    pub fn resolve_dependences(&mut self, fetched_handles: &HashSet<Handle>) -> Result<()> {
+        let remaining: Vec<(NodeIndex, NodeIndex)> = self
+            .deferred_dependences
+            .drain(..)
+            .filter(|(_, _, handle)| !fetched_handles.contains(handle))
+            .map(|(prod, cons, _)| (prod, cons))
+            .collect();
+        if remaining.is_empty() {
+            return Ok(());
+        }
+        // Build a digraph replica including existing edges +
+        // remaining deferred edges and check for cycles
+        let mut digraph = self.graph.map(|idx, _| idx, |_, _| ()).graph().clone();
+        for (producer, consumer) in remaining.iter() {
             digraph.add_edge(*producer, *consumer, ());
         }
         let mut tarjan = daggy::petgraph::algo::TarjanScc::new();
         let mut sccs = Vec::new();
         tarjan.run(&digraph, |scc| {
             if scc.len() > 1 {
-                // All non-singleton SCCs in a directed graph are
-                // dependence cycles
                 sccs.push(scc.to_vec());
             }
         });
@@ -389,9 +432,6 @@ impl DFComponentGraph {
                         .graph
                         .node_weight_mut(*idx)
                         .ok_or(SchedulerError::DataflowGraphError)?;
-                    // Mark the node as uncomputable so we don't go
-                    // and mark as completed operations that are in
-                    // error.
                     tx.is_uncomputable = true;
                     error!(target: "scheduler", { transaction_id = ?hex::encode(tx.transaction_id.clone()) },
 		       "Transaction is part of a dependence cycle");
@@ -405,26 +445,20 @@ impl DFComponentGraph {
                 }
             }
             return Err(SchedulerError::CyclicDependence.into());
-        } else {
-            // If no dependence cycles were found, then we can
-            // complete the graph and proceed to execution
-            for (producer, consumer) in dependence_pairs.iter() {
-                // The error case here should not happen as we've
-                // already covered it by testing for SCCs in the graph
-                // first
-                if self.graph.add_edge(*producer, *consumer, ()).is_err() {
-                    let prod = self
-                        .graph
-                        .node_weight(*producer)
-                        .ok_or(SchedulerError::DataflowGraphError)?;
-                    let cons = self
-                        .graph
-                        .node_weight(*consumer)
-                        .ok_or(SchedulerError::DataflowGraphError)?;
-                    error!(target: "scheduler", { producer_id = ?hex::encode(prod.transaction_id.clone()), consumer_id = ?hex::encode(cons.transaction_id.clone()) },
+        }
+        for (producer, consumer) in remaining.iter() {
+            if self.graph.add_edge(*producer, *consumer, ()).is_err() {
+                let prod = self
+                    .graph
+                    .node_weight(*producer)
+                    .ok_or(SchedulerError::DataflowGraphError)?;
+                let cons = self
+                    .graph
+                    .node_weight(*consumer)
+                    .ok_or(SchedulerError::DataflowGraphError)?;
+                error!(target: "scheduler", { producer_id = ?hex::encode(prod.transaction_id.clone()), consumer_id = ?hex::encode(cons.transaction_id.clone()) },
 		       "Dependence cycle when adding dependence - initial cycle detection failed");
-                    return Err(SchedulerError::CyclicDependence.into());
-                }
+                return Err(SchedulerError::CyclicDependence.into());
             }
         }
         Ok(())

coprocessor/fhevm-engine/scheduler/src/dfg.rs

@@ -531,23 +531,26 @@ async fn build_transaction_graph_and_execute<'a>(
     let cts_to_query = tx_graph.needed_map.keys().cloned().collect::<Vec<_>>();
     let ciphertext_map =
         query_ciphertexts(&cts_to_query, *tenant_id, trx, tracer, loop_ctx).await?;
-    // Check if we retrieved all needed CTs - if not, we may not want to proceed to execution
-    if cts_to_query.len() != ciphertext_map.len() {
+    let fetched_handles: std::collections::HashSet<_> = ciphertext_map.keys().cloned().collect();
+    if cts_to_query.len() != fetched_handles.len() {
         if let Some(dcid_lock) = dcid_mngr.get_current_lock() {
-            warn!(target: "tfhe_worker", { missing_inputs = ?(cts_to_query.len() - ciphertext_map.len()), dcid = %hex::encode(dcid_lock.dependence_chain_id) },
+            warn!(target: "tfhe_worker", { missing_inputs = ?(cts_to_query.len() - fetched_handles.len()), dcid = %hex::encode(dcid_lock.dependence_chain_id) },
 	      "some inputs are missing to execute the dependence chain");
         }
-        // Do not stop execution, we will allow the scheduler to run
-        // and complete as many operations as can be computed given
-        // the inputs fetched
     }
-
     for (handle, (ct_type, mut ct)) in ciphertext_map.into_iter() {
         tx_graph.add_input(
             &handle,
             &DFGTxInput::Compressed(((ct_type, std::mem::take(&mut ct)), true)),
         )?;
     }
+    // Resolve deferred cross-transaction dependences: edges whose
+    // handle was fetched from DB are dropped (data already available),
+    // remaining edges are added after cycle detection.
+    if let Err(e) = tx_graph.resolve_dependences(&fetched_handles) {
+        warn!(target: "tfhe_worker", { error = %e }, "error resolving cross-transaction dependences");
+        return Ok(tx_graph);
+    }
     // Execute the DFG with the current tenant's keys
     let mut s_compute = tracer.start_with_context("compute_fhe_ops", loop_ctx);
     {