fix(schedule): prevent panicking in loop_schedule routine

zama-ai · Dec 6, 2024 · 5789b54 · 5789b54
1 parent b4c2d53
commit 5789b54
Showing 1 changed file with 20 additions and 13 deletions.
diff --git a/fhevm-engine/scheduler/src/dfg/scheduler.rs b/fhevm-engine/scheduler/src/dfg/scheduler.rs
@@ -268,21 +268,28 @@ impl<'a> Scheduler<'a> {
             }
         }
 
-        rayon::broadcast(|_| {
-            tfhe::set_server_key(server_key.clone());
-        });
+
         let (src, dest) = channel();
         let rayon_threads = self.rayon_threads;
-        comps.par_iter().for_each_with(src, |src, (args, index)| {
-            src.send(execute_partition(
-                args.to_vec(),
-                *index,
-                true,
-                rayon_threads,
-                server_key.clone(),
-            ))
-            .unwrap();
-        });
+
+        tokio::task::spawn_blocking(move || {
+            rayon::broadcast(|_| {
+                tfhe::set_server_key(server_key.clone());
+            });
+
+            tfhe::set_server_key(server_key.clone());
+            comps.par_iter().for_each_with(src, |src, (args, index)| {
+                src.send(execute_partition(
+                    args.to_vec(),
+                    *index,
+                    true,
+                    rayon_threads,
+                    server_key.clone(),
+                ))
+                .unwrap();
+            });
+        }).await?;
+
         let results: Vec<_> = dest.iter().collect();
         for result in results {
             let mut output = result?;