dbg: add syncs after each step (#425)

worldcoin · Sep 20, 2024 · d2fc51f · d2fc51f
1 parent 47fb009
commit d2fc51f
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 1 deletion.
diff --git a/deploy/stage/common-values-iris-mpc.yaml b/deploy/stage/common-values-iris-mpc.yaml
@@ -1,4 +1,4 @@
-image: "ghcr.io/worldcoin/iris-mpc:v0.7.4"
+image: "ghcr.io/worldcoin/iris-mpc:v0.7.5"
 
 environment: stage
 replicaCount: 1

diff --git a/iris-mpc-gpu/src/server/actor.rs b/iris-mpc-gpu/src/server/actor.rs
@@ -929,6 +929,8 @@ impl ServerActor {
             tracing::info!(party_id = self.party_id, "batch_dot end");
         });
 
+        self.device_manager.await_streams(batch_streams); // DEBUG
+
         record_stream_time!(
             &self.device_manager,
             batch_streams,
@@ -945,6 +947,8 @@ impl ServerActor {
             }
         );
 
+        self.device_manager.await_streams(batch_streams); // DEBUG
+
         let db_sizes_batch =
             vec![self.max_batch_size * ROTATIONS; self.device_manager.device_count()];
         let code_dots_batch = self.batch_codes_engine.result_chunk_shares(&db_sizes_batch);
@@ -966,6 +970,8 @@ impl ServerActor {
             }
         );
 
+        self.device_manager.await_streams(batch_streams); // DEBUG
+
         tracing::info!(party_id = self.party_id, "phase2_batch start");
 
         let res = self.phase2_batch.take_result_buffer();
@@ -985,6 +991,8 @@ impl ServerActor {
         );
         self.phase2_batch.return_result_buffer(res);
 
+        self.device_manager.await_streams(batch_streams); // DEBUG
+
         tracing::info!(party_id = self.party_id, "Finished batch deduplication");
         // ---- END BATCH DEDUP ----
 
@@ -1044,6 +1052,8 @@ impl ServerActor {
             self.device_manager
                 .await_event(request_streams, &current_dot_event);
 
+            self.device_manager.await_streams(request_streams); // DEBUG
+
             // ---- START PHASE 1 ----
             tracing::info!(party_id = self.party_id, "Start PHASE 1");
 
@@ -1060,6 +1070,8 @@ impl ServerActor {
                 );
             });
 
+            self.device_manager.await_streams(request_streams); // DEBUG
+
             // wait for the exchange result buffers to be ready
             tracing::info!(
                 party_id = self.party_id,
@@ -1069,6 +1081,8 @@ impl ServerActor {
             self.device_manager
                 .await_event(request_streams, &current_exchange_event);
 
+            self.device_manager.await_streams(request_streams); // DEBUG
+
             record_stream_time!(&self.device_manager, batch_streams, events, "db_reduce", {
                 compact_device_sums.compute_dot_reducer_against_db(
                     &mut self.codes_engine,
@@ -1081,6 +1095,8 @@ impl ServerActor {
                 );
             });
 
+            self.device_manager.await_streams(request_streams); // DEBUG
+
             tracing::info!(
                 party_id = self.party_id,
                 chunk = db_chunk_idx,
@@ -1097,6 +1113,8 @@ impl ServerActor {
             });
             tracing::info!(party_id = self.party_id, "End PHASE 1");
 
+            self.device_manager.await_streams(request_streams); // DEBUG
+
             // ---- END PHASE 1 ----
 
             tracing::info!(
@@ -1135,6 +1153,9 @@ impl ServerActor {
                         );
                     }
                 );
+
+                self.device_manager.await_streams(request_streams); // DEBUG
+
                 // we can now record the exchange event since the phase 2 is no longer using the
                 // code_dots/mask_dots which are just reinterpretations of the exchange result
                 // buffers
@@ -1161,6 +1182,8 @@ impl ServerActor {
                     request_streams,
                 );
                 self.phase2.return_result_buffer(res);
+
+                self.device_manager.await_streams(request_streams); // DEBUG
             }
             tracing::info!(
                 party_id = self.party_id,
@@ -1171,6 +1194,8 @@ impl ServerActor {
                 .record_event(request_streams, &next_phase2_event);
             tracing::info!(party_id = self.party_id, "End PHASE 2");
 
+            self.device_manager.await_streams(request_streams); // DEBUG
+
             // ---- END PHASE 2 ----
 
             // Update events for synchronization