Merge branch 'bump/cudarc-0.13.3' into feat/LUC-with-bitmaps

worldcoin · Jan 24, 2025 · 6ba1406 · 6ba1406
2 parents b970cb6 + e4cb609
commit 6ba1406
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 5 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:
-      - "add-error-handling-skipped-requests"
+      - "add-graceful-shutdown-during-s3-sync"
 
 concurrency:
   group: '${{ github.workflow }} @ ${{ github.event.pull_request.head.label || github.head_ref || github.ref }}'

diff --git a/iris-mpc-gpu/src/bin/nccl.rs b/iris-mpc-gpu/src/bin/nccl.rs
@@ -85,9 +85,15 @@ async fn main() -> eyre::Result<()> {
         for i in 0..n_devices {
             devs[i].bind_to_thread().unwrap();
 
-            comms[i].broadcast(&slices[i], &mut slices1[i], 0).unwrap();
-            comms[i].broadcast(&slices[i], &mut slices2[i], 1).unwrap();
-            comms[i].broadcast(&slices[i], &mut slices3[i], 2).unwrap();
+            comms[i]
+                .broadcast(slices[i].as_ref(), &mut slices1[i], 0)
+                .unwrap();
+            comms[i]
+                .broadcast(slices[i].as_ref(), &mut slices2[i], 1)
+                .unwrap();
+            comms[i]
+                .broadcast(slices[i].as_ref(), &mut slices3[i], 2)
+                .unwrap();
         }
 
         for dev in devs.iter() {

diff --git a/iris-mpc/src/bin/server.rs b/iris-mpc/src/bin/server.rs
@@ -989,6 +989,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
     let load_chunks_parallelism = config.load_chunks_parallelism;
     let db_chunks_bucket_name = config.db_chunks_bucket_name.clone();
     let db_chunks_folder_name = config.db_chunks_folder_name.clone();
+    let download_shutdown_handler = Arc::clone(&shutdown_handler);
 
     let (tx, rx) = oneshot::channel();
     background_tasks.spawn_blocking(move || {
@@ -1068,6 +1069,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
                         "Initialize iris db: Loading from DB (parallelism: {})",
                         parallelism
                     );
+                    let download_shutdown_handler = Arc::clone(&download_shutdown_handler);
                     let s3_store = S3Store::new(s3_client_clone, db_chunks_bucket_name);
                     tokio::runtime::Handle::current().block_on(async {
                         let mut stream = match config.enable_s3_importer {
@@ -1199,6 +1201,10 @@ async fn server_main(config: Config) -> eyre::Result<()> {
                                     elapsed,
                                     record_counter as f64 / elapsed.as_secs_f64()
                                 );
+                                if download_shutdown_handler.is_shutting_down() {
+                                    tracing::warn!("Shutdown requested by shutdown_handler.");
+                                    return Err(eyre::eyre!("Shutdown requested"));
+                                }
                             }
 
                             // if the serial id hasn't been loaded before, count is as unique record
@@ -1282,7 +1288,7 @@ async fn server_main(config: Config) -> eyre::Result<()> {
     let sns_client_bg = sns_client.clone();
     let config_bg = config.clone();
     let store_bg = store.clone();
-    let shutdown_handler_bg = shutdown_handler.clone();
+    let shutdown_handler_bg = Arc::clone(&shutdown_handler);
     let _result_sender_abort = background_tasks.spawn(async move {
         while let Some(ServerJobResult {
             merged_results,