From 5789b54bb6be2b129197bd47a6dc2f02c12e6e9c Mon Sep 17 00:00:00 2001
From: Georgi <georgi.goranov@zama.ai>
Date: Fri, 6 Dec 2024 13:01:26 +0200
Subject: [PATCH] fix(schedule): prevent panicking in loop_schedule routine

---
 fhevm-engine/scheduler/src/dfg/scheduler.rs | 33 +++++++++++++--------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/fhevm-engine/scheduler/src/dfg/scheduler.rs b/fhevm-engine/scheduler/src/dfg/scheduler.rs
index 79330481..b918cc5a 100644
--- 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?;