Feat: even odd import from s3 (#897)

* import bytes in even-odd orders from s3 * trigger image push * deploy to stage * fix code len assert * remove leftover change on loading from db * change even/odd order * bring back iris-mpc-py to workspace * remove unused method * simplify size * move invalid id check to pre-load * bump stage version
worldcoin · Jan 10, 2025 · f3e2ae0 · f3e2ae0
1 parent 8568198
commit f3e2ae0
Show file tree

Hide file tree

Showing 10 changed files with 297 additions and 117 deletions.
diff --git a/.github/workflows/temp-branch-build-and-push.yaml b/.github/workflows/temp-branch-build-and-push.yaml
@@ -3,7 +3,7 @@ name: Branch - Build and push docker image
 on:
   push:
     branches:
-      - "ps/potential-phantom-match"
+      - "feat/even-odd-import-from-s3"
 
 concurrency:
   group: '${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}'

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.13.10"
+image: "ghcr.io/worldcoin/iris-mpc:v0.13.11"
 
 environment: stage
 replicaCount: 1

diff --git a/deploy/stage/smpcv2-0-stage/values-iris-mpc.yaml b/deploy/stage/smpcv2-0-stage/values-iris-mpc.yaml
@@ -75,7 +75,7 @@ env:
     value: "iris-mpc-db-exporter-store-node-0-stage-eu-north-1"
 
   - name: SMPC__DB_CHUNKS_FOLDER_NAME
-    value: "binary_output_16k"
+    value: "even_odd_binary_output_16k"
 
   - name: SMPC__LOAD_CHUNKS_PARALLELISM
     value: "32"

diff --git a/deploy/stage/smpcv2-1-stage/values-iris-mpc.yaml b/deploy/stage/smpcv2-1-stage/values-iris-mpc.yaml
@@ -75,7 +75,7 @@ env:
     value: "iris-mpc-db-exporter-store-node-1-stage-eu-north-1"
 
   - name: SMPC__DB_CHUNKS_FOLDER_NAME
-    value: "binary_output_16k"
+    value: "even_odd_binary_output_16k"
 
   - name: SMPC__LOAD_CHUNKS_PARALLELISM
     value: "32"

diff --git a/deploy/stage/smpcv2-2-stage/values-iris-mpc.yaml b/deploy/stage/smpcv2-2-stage/values-iris-mpc.yaml
@@ -75,7 +75,7 @@ env:
     value: "iris-mpc-db-exporter-store-node-2-stage-eu-north-1"
 
   - name: SMPC__DB_CHUNKS_FOLDER_NAME
-    value: "binary_output_16k"
+    value: "even_odd_binary_output_16k"
 
   - name: SMPC__LOAD_CHUNKS_PARALLELISM
     value: "32"

diff --git a/iris-mpc-gpu/src/dot/share_db.rs b/iris-mpc-gpu/src/dot/share_db.rs
@@ -302,7 +302,7 @@ impl ShareDB {
         }
     }
 
-    pub fn load_single_record(
+    pub fn load_single_record_from_db(
         index: usize,
         db: &CudaVec2DSlicerRawPointer,
         record: &[u16],
@@ -339,6 +339,35 @@ impl ShareDB {
         };
     }
 
+    pub fn load_single_record_from_s3(
+        index: usize,
+        db: &CudaVec2DSlicerRawPointer,
+        a0_host: &[u8],
+        a1_host: &[u8],
+        n_shards: usize,
+        code_length: usize,
+    ) {
+        assert_eq!(a0_host.len(), code_length);
+        assert_eq!(a1_host.len(), code_length);
+
+        let device_index = index % n_shards;
+        let device_db_index = index / n_shards;
+
+        unsafe {
+            std::ptr::copy(
+                a0_host.as_ptr() as *const _,
+                (db.limb_0[device_index] + (device_db_index * code_length) as u64) as *mut _,
+                code_length,
+            );
+
+            std::ptr::copy(
+                a1_host.as_ptr() as *const _,
+                (db.limb_1[device_index] + (device_db_index * code_length) as u64) as *mut _,
+                code_length,
+            );
+        };
+    }
+
     pub fn preprocess_db(&self, db: &mut SlicedProcessedDatabase, db_lens: &[usize]) {
         let code_len = self.code_length;
         for device_index in 0..self.device_manager.device_count() {
@@ -381,7 +410,7 @@ impl ShareDB {
             .par_chunks(self.code_length)
             .enumerate()
             .for_each(|(idx, chunk)| {
-                Self::load_single_record(idx, &db.code_gr, chunk, n_shards, code_length);
+                Self::load_single_record_from_db(idx, &db.code_gr, chunk, n_shards, code_length);
             });
 
         // Calculate the number of entries per shard

diff --git a/iris-mpc-gpu/src/server/actor.rs b/iris-mpc-gpu/src/server/actor.rs
@@ -461,36 +461,36 @@ impl ServerActor {
         self.current_db_sizes = db_lens1;
     }
 
-    pub fn load_single_record(
+    pub fn load_single_record_from_db(
         &mut self,
         index: usize,
         left_code: &[u16],
         left_mask: &[u16],
         right_code: &[u16],
         right_mask: &[u16],
     ) {
-        ShareDB::load_single_record(
+        ShareDB::load_single_record_from_db(
             index,
             &self.left_code_db_slices.code_gr,
             left_code,
             self.device_manager.device_count(),
             IRIS_CODE_LENGTH,
         );
-        ShareDB::load_single_record(
+        ShareDB::load_single_record_from_db(
             index,
             &self.left_mask_db_slices.code_gr,
             left_mask,
             self.device_manager.device_count(),
             MASK_CODE_LENGTH,
         );
-        ShareDB::load_single_record(
+        ShareDB::load_single_record_from_db(
             index,
             &self.right_code_db_slices.code_gr,
             right_code,
             self.device_manager.device_count(),
             IRIS_CODE_LENGTH,
         );
-        ShareDB::load_single_record(
+        ShareDB::load_single_record_from_db(
             index,
             &self.right_mask_db_slices.code_gr,
             right_mask,
@@ -499,6 +499,53 @@ impl ServerActor {
         );
     }
 
+    #[allow(clippy::too_many_arguments)]
+    pub fn load_single_record_from_s3(
+        &mut self,
+        index: usize,
+        left_code_odd: &[u8],
+        left_code_even: &[u8],
+        right_code_odd: &[u8],
+        right_code_even: &[u8],
+        left_mask_odd: &[u8],
+        left_mask_even: &[u8],
+        right_mask_odd: &[u8],
+        right_mask_even: &[u8],
+    ) {
+        ShareDB::load_single_record_from_s3(
+            index,
+            &self.left_code_db_slices.code_gr,
+            left_code_odd,
+            left_code_even,
+            self.device_manager.device_count(),
+            IRIS_CODE_LENGTH,
+        );
+        ShareDB::load_single_record_from_s3(
+            index,
+            &self.left_mask_db_slices.code_gr,
+            left_mask_odd,
+            left_mask_even,
+            self.device_manager.device_count(),
+            MASK_CODE_LENGTH,
+        );
+        ShareDB::load_single_record_from_s3(
+            index,
+            &self.right_code_db_slices.code_gr,
+            right_code_odd,
+            right_code_even,
+            self.device_manager.device_count(),
+            IRIS_CODE_LENGTH,
+        );
+        ShareDB::load_single_record_from_s3(
+            index,
+            &self.right_mask_db_slices.code_gr,
+            right_mask_odd,
+            right_mask_even,
+            self.device_manager.device_count(),
+            MASK_CODE_LENGTH,
+        );
+    }
+
     pub fn increment_db_size(&mut self, index: usize) {
         self.current_db_sizes[index % self.device_manager.device_count()] += 1;
     }