Load the second half of the database on demand for GPU version

.github/workflows/temp-branch-build-and-push.yaml

@@ -3,7 +3,7 @@ name: Branch - Build and push docker image
 on:
   push:
     branches:
-      - "chore/insert-sns-msg-id-upon-reset-update-msg"
+      - "dk/load_half_on_demand"
 
 concurrency:
   group: '${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}'

deploy/stage/common-values-iris-mpc.yaml

@@ -1,4 +1,4 @@
-image: "ghcr.io/worldcoin/iris-mpc@sha256:f1b79237c2e5eb7dd3c26977a49d21d09125d6dc7da69213b3d65c7693b82ca7" # v0.15.5
+image: "ghcr.io/worldcoin/iris-mpc:de51a749fbf8de65f71caf4b0c79e9e878b6874e" # v0.15.5
 
 environment: stage
 replicaCount: 1

deploy/stage/smpcv2-0-stage/values-iris-mpc.yaml

@@ -90,10 +90,10 @@ env:
     value: "true"
 
   - name: SMPC__INIT_DB_SIZE
-    value: "1000" # "1200000"
+    value: "1200000"
 
   - name: SMPC__MAX_DB_SIZE
-    value: "100000" # "1300000"
+    value: "1300000"
 
   - name: SMPC__MAX_BATCH_SIZE
     value: "64"
@@ -114,10 +114,10 @@ env:
     value: "true"
 
   - name: SMPC__LUC_ENABLED
-    value: "false"
+    value: "true"
 
   - name: SMPC__LUC_LOOKBACK_RECORDS
-    value: "50"
+    value: "1000"
 
   - name: SMPC__LUC_SERIAL_IDS_FROM_SMPC_REQUEST
     value: "true"
@@ -157,9 +157,12 @@ env:
   - name: SMPC__ENABLE_MODIFICATIONS_REPLAY
     value: "true"
 
-  - name : SMPC__ENABLE_DEBUG_TIMING
+  - name: SMPC__ENABLE_DEBUG_TIMING
     value: "true"
 
+  - name: SMPC__LOAD_ONLY_FULL_SCAN_SIDE_IN_MEMORY
+    value: "false"
+
 initContainer:
   enabled: true
   image: "amazon/aws-cli:2.17.62"

deploy/stage/smpcv2-1-stage/values-iris-mpc.yaml

@@ -90,10 +90,10 @@ env:
     value: "true"
 
   - name: SMPC__INIT_DB_SIZE
-    value: "1000" # "1200000"
+    value: "1200000"
 
   - name: SMPC__MAX_DB_SIZE
-    value: "100000" # "1300000"
+    value: "1300000"
 
   - name: SMPC__MAX_BATCH_SIZE
     value: "64"
@@ -114,10 +114,10 @@ env:
     value: "true"
 
   - name: SMPC__LUC_ENABLED
-    value: "false"
+    value: "true"
 
   - name: SMPC__LUC_LOOKBACK_RECORDS
-    value: "50"
+    value: "1000"
 
   - name: SMPC__LUC_SERIAL_IDS_FROM_SMPC_REQUEST
     value: "true"
@@ -160,6 +160,9 @@ env:
   - name : SMPC__ENABLE_DEBUG_TIMING
     value: "true"
 
+  - name: SMPC__LOAD_ONLY_FULL_SCAN_SIDE_IN_MEMORY
+    value: "false"
+
 initContainer:
   enabled: true
   image: "amazon/aws-cli:2.17.62"

deploy/stage/smpcv2-2-stage/values-iris-mpc.yaml

@@ -90,10 +90,10 @@ env:
     value: "true"
 
   - name: SMPC__INIT_DB_SIZE
-    value: "1000" # "1200000"
+    value: "1200000"
 
   - name: SMPC__MAX_DB_SIZE
-    value: "100000" # "1300000"
+    value: "1300000"
 
   - name: SMPC__MAX_BATCH_SIZE
     value: "64"
@@ -114,10 +114,10 @@ env:
     value: "true"
 
   - name: SMPC__LUC_ENABLED
-    value: "false"
+    value: "true"
 
   - name: SMPC__LUC_LOOKBACK_RECORDS
-    value: "50"
+    value: "1000"
 
   - name: SMPC__LUC_SERIAL_IDS_FROM_SMPC_REQUEST
     value: "true"
@@ -160,6 +160,9 @@ env:
   - name : SMPC__ENABLE_DEBUG_TIMING
     value: "true"
 
+  - name: SMPC__LOAD_ONLY_FULL_SCAN_SIDE_IN_MEMORY
+    value: "false"
+
 initContainer:
   enabled: true
   image: "amazon/aws-cli:2.17.62"

iris-mpc-common/Cargo.toml

@@ -26,6 +26,7 @@ aws-sdk-sns = { workspace = true, optional = true }
 aws-sdk-sqs = { workspace = true, optional = true }
 aws-sdk-s3 = { workspace = true, optional = true }
 aws-sdk-secretsmanager = { workspace = true, optional = true }
+async-trait.workspace = true
 dotenvy.workspace = true
 clap.workspace = true
 rand.workspace = true
@@ -63,6 +64,7 @@ ring = "0.17.8"
 data-encoding = "2.6.0"
 bincode.workspace = true
 serde-big-array.workspace = true
+quick_cache = "0.6.13"
 metrics = "0.22.1"
 metrics-exporter-statsd = "0.7"
 

iris-mpc-common/src/config/mod.rs

@@ -242,6 +242,8 @@ pub struct Config {
     #[serde(default = "default_full_scan_side")]
     pub full_scan_side: Eye,
 
+    #[serde(default)]
+    pub load_only_full_scan_side_in_memory: bool,
     // used to fix max batch size to 1 for correctness testing purposes
     #[serde(default = "default_override_max_batch_size")]
     pub override_max_batch_size: bool,
@@ -629,7 +631,8 @@ impl From<Config> for CommonConfig {
             cpu_disable_persistence,
             hawk_server_resets_enabled,
             full_scan_side,
-            override_max_batch_size: _, // for testing purposes only
+            load_only_full_scan_side_in_memory: _, // could be different for each server
+            override_max_batch_size: _,            // for testing purposes only
         } = value;
 
         Self {

iris-mpc-common/src/galois_engine.rs

@@ -450,6 +450,7 @@ pub mod degree4 {
         })
     }
 
+    #[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
     pub struct FullGaloisRingIrisCodeShare {
         pub code: GaloisRingIrisCodeShare,
         pub mask: GaloisRingTrimmedMaskCodeShare,

iris-mpc-common/src/helpers/inmemory_store.rs

@@ -1,4 +1,6 @@
-use crate::vector_id::VectorId;
+use std::sync::Arc;
+
+use crate::{job::Eye, vector_id::VectorId};
 
 /// A helper trait encapsulating the functionality to add iris codes to some
 /// form of in-memory store.
@@ -117,3 +119,87 @@ pub trait InMemoryStore {
     /// unpredictable.
     fn fake_db(&mut self, size: usize);
 }
+
+/// A helper trait encapsulating the functionality to load iris codes on demand from some source (DB, File-backed, etc.).
+#[async_trait::async_trait]
+pub trait OnDemandLoader {
+    /// Loads records from the source.
+    /// The returned Vec has the form:
+    /// `(index, side_code, side_mask)`.
+    #[allow(clippy::type_complexity)]
+    async fn load_records(
+        &self,
+        side: Eye,
+        indices: &[usize],
+    ) -> eyre::Result<Vec<(usize, Vec<u16>, Vec<u16>)>>;
+}
+
+pub struct CachingOnDemandLoader {
+    inner: Arc<dyn OnDemandLoader + Send + Sync + 'static>,
+    cache: quick_cache::sync::Cache<usize, (Vec<u16>, Vec<u16>)>,
+    fixed_side: Eye,
+}
+
+impl CachingOnDemandLoader {
+    pub fn new(
+        inner: Arc<dyn OnDemandLoader + Send + Sync + 'static>,
+        num_cached_iris_codes: usize,
+        fixed_side: Eye,
+    ) -> Self {
+        let cache = quick_cache::sync::Cache::new(num_cached_iris_codes);
+        Self {
+            inner,
+            cache,
+            fixed_side,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl OnDemandLoader for CachingOnDemandLoader {
+    async fn load_records(
+        &self,
+        side: Eye,
+        indices: &[usize],
+    ) -> eyre::Result<Vec<(usize, Vec<u16>, Vec<u16>)>> {
+        if side != self.fixed_side {
+            return Err(eyre::eyre!(
+                "Invalid side requested, expected: {:?}",
+                self.fixed_side
+            ));
+        }
+
+        // grab cached iris codes
+        let cached: Vec<_> = indices.iter().map(|x| (*x, self.cache.get(x))).collect();
+        // grab missing ones
+        let missing_indices = cached
+            .iter()
+            .filter_map(|(i, c)| if c.is_none() { Some(*i) } else { None })
+            .collect::<Vec<_>>();
+
+        // for missing ones, load from the inner source
+        let loaded = if missing_indices.is_empty() {
+            vec![]
+        } else {
+            self.inner.load_records(side, &missing_indices).await?
+        };
+
+        // put collected iris codes into the cache
+        for (idx, code, mask) in loaded.iter() {
+            self.cache.insert(*idx, (code.clone(), mask.clone()));
+        }
+        let result = cached
+            .into_iter()
+            .filter_map(|x| {
+                if let Some((code, mask)) = x.1 {
+                    Some((x.0, code, mask))
+                } else {
+                    None
+                }
+            })
+            .chain(loaded.into_iter())
+            .collect::<Vec<_>>();
+
+        Ok(result)
+    }
+}