From 5d7d6e600db0637773837ab7be1a41ef0c067da9 Mon Sep 17 00:00:00 2001
From: Georgi <georgi.goranov@zama.ai>
Date: Wed, 4 Dec 2024 15:31:29 +0200
Subject: [PATCH 1/3] Set server keys for each thread in Rayon pool

---
 fhevm-engine/scheduler/src/dfg/scheduler.rs | 40 +++++++++++++--------
 1 file changed, 26 insertions(+), 14 deletions(-)

diff --git a/fhevm-engine/scheduler/src/dfg/scheduler.rs b/fhevm-engine/scheduler/src/dfg/scheduler.rs
index 79330481..5baddb78 100644
--- a/fhevm-engine/scheduler/src/dfg/scheduler.rs
+++ b/fhevm-engine/scheduler/src/dfg/scheduler.rs
@@ -79,6 +79,9 @@ impl<'a> Scheduler<'a> {
 
     pub async fn schedule(&mut self, server_key: tfhe::ServerKey) -> Result<()> {
         let schedule_type = std::env::var("FHEVM_DF_SCHEDULE");
+
+        println!("Scheduling with strategy {:?}", schedule_type);
+
         match schedule_type {
             Ok(val) => match val.as_str() {
                 "MAX_PARALLELISM" => {
@@ -89,11 +92,11 @@ impl<'a> Scheduler<'a> {
                     self.schedule_coarse_grain(PartitionStrategy::MaxLocality, server_key)
                         .await
                 }
-                "LOOP" => self.schedule_component_loop(server_key).await,
+                "LOOP" => self.schedule_component_loop(server_key),
                 "FINE_GRAIN" => self.schedule_fine_grain(server_key).await,
                 unhandled => panic!("Scheduling strategy {:?} does not exist", unhandled),
             },
-            _ => self.schedule_component_loop(server_key).await,
+            _ => self.schedule_component_loop(server_key),
         }
     }
 
@@ -248,7 +251,7 @@ impl<'a> Scheduler<'a> {
         Ok(())
     }
 
-    async fn schedule_component_loop(&mut self, server_key: tfhe::ServerKey) -> Result<()> {
+    fn schedule_component_loop(&mut self, server_key: tfhe::ServerKey) -> Result<()> {
         let mut execution_graph: Dag<ExecNode, ()> = Dag::default();
         let _ = partition_components(self.graph, &mut execution_graph);
         let mut comps = vec![];
@@ -268,21 +271,28 @@ impl<'a> Scheduler<'a> {
             }
         }
 
+        let (src, dest) = channel();
+        let rayon_threads = self.rayon_threads;
+
         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();
+        
+        // Execute computations in parallel using Rayon threadpool
+        tokio::task::block_in_place( move || {
+            comps.par_iter().for_each_with(src, |src, (args, index)| {
+                tfhe::set_server_key(server_key.clone());
+                src.send(execute_partition(
+                    args.to_vec(),
+                    *index,
+                    true,
+                    rayon_threads,
+                    server_key.clone(),
+                ))
+                .unwrap();
+            });
         });
+
         let results: Vec<_> = dest.iter().collect();
         for result in results {
             let mut output = result?;
@@ -292,6 +302,7 @@ impl<'a> Scheduler<'a> {
                 self.graph[node_index].result = Some(o.1);
             }
         }
+
         Ok(())
     }
 }
@@ -444,6 +455,7 @@ fn execute_partition(
                 }
             }
         }
+
         if use_global_threadpool {
             let (node_index, result) = run_computation(opcode, Ok(cts), nidx.index())?;
             res.insert(node_index, result);

From 242a99384d8ca8815e5b198fd5894606dd2c2e9a Mon Sep 17 00:00:00 2001
From: 0xawaz <amina.waddiz@gmail.com>
Date: Wed, 4 Dec 2024 15:18:00 +0100
Subject: [PATCH 2/3] ci: add current branch

---
 .github/workflows/fhevm-coprocessor.yml | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/.github/workflows/fhevm-coprocessor.yml b/.github/workflows/fhevm-coprocessor.yml
index 58775b1e..73f59493 100644
--- a/.github/workflows/fhevm-coprocessor.yml
+++ b/.github/workflows/fhevm-coprocessor.yml
@@ -4,10 +4,12 @@ on:
   push:
     branches:
       - main
+      - temp-fix-coprocessor-crash
     paths:
       - .github/workflows/fhevm-coprocessor.yml
       - .github/workflows/common-docker.yml
       - fhevm-engine/coprocessor/**
+      - fhevm-engine/scheduler/**
       - fhevm-engine/.sqlx
       - proto/**
   release:

From b7295da8f5ea49e045bd14ba5895c423a555513f Mon Sep 17 00:00:00 2001
From: 0xawaz <amina.waddiz@gmail.com>
Date: Wed, 4 Dec 2024 15:39:53 +0100
Subject: [PATCH 3/3] ci: enable push for current branch

---
 .github/workflows/fhevm-coprocessor.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/fhevm-coprocessor.yml b/.github/workflows/fhevm-coprocessor.yml
index 73f59493..e7d7eac5 100644
--- a/.github/workflows/fhevm-coprocessor.yml
+++ b/.github/workflows/fhevm-coprocessor.yml
@@ -29,7 +29,7 @@ jobs:
       packages: "write"
     with:
       working-directory: "."
-      push_image: ${{ github.event_name == 'release' || github.ref_name == 'main' }}
+      push_image: true
       image-name: "fhevm-coprocessor"
       generate-dev-image: false
       docker-file: "fhevm-engine/coprocessor/Dockerfile"