feat(coprocessor): enable GPU backend in SnS worker (#1801)

* feat(coprocessor): enable GPU support in SnS worker

* chore(coprocessor): logs if gpu or cpu backend is enabled

* chore(coprocessor): enable test-harness/gpu feature

* chore(coprocessor): update sns_pk in setup_test_user

* chore(coprocessor): log if gpu or cpu backend is enabled, tfhe-worker

* chore(coprocessor): disable s3 asserts in SnS unit-testing on GPU backend

* chore(coprocessor): scope gpu testing to gpu-dependent crates

Author: goshawk-3

.github/workflows/coprocessor-gpu-tests.yml

@@ -185,12 +185,18 @@ jobs:
         env:
           HARDHAT_NETWORK: hardhat
 
-      - name: Run tests on GPU
+      - name: Run GPU tests for the worker services.
         run: |
-          DATABASE_URL=postgresql://postgres:postgres@localhost:5432/coprocessor cargo test --release --features=gpu -- --test-threads=1
+          export DATABASE_URL=postgresql://postgres:postgres@localhost:5432/coprocessor
+          cargo test \
+          -p tfhe-worker \
+          -p sns-worker \
+          --release \
+          --features=gpu \
+          -- \
+          --test-threads=1
         working-directory: coprocessor/fhevm-engine
 
-
   teardown-instance:
     name: coprocessor-gpu-tests/teardown
     if: ${{ always() && needs.setup-instance.result == 'success' }}

coprocessor/fhevm-engine/fhevm-engine-common/src/utils.rs

@@ -244,3 +244,27 @@ impl FromStr for DatabaseURL {
         Ok(Self(s.to_owned()))
     }
 }
+
+/// Logs whether the GPU backend is enabled or not.
+pub fn log_backend() -> bool {
+    log_backend_impl()
+}
+
+#[cfg(feature = "gpu")]
+fn log_backend_impl() -> bool {
+    use tfhe::core_crypto::gpu::{get_number_of_gpus, get_number_of_sms};
+    let num_gpus = get_number_of_gpus();
+    let streaming_multiprocessors = get_number_of_sms();
+    tracing::info!(
+        num_gpus,
+        streaming_multiprocessors,
+        "GPU feature is enabled"
+    );
+    true
+}
+
+#[cfg(not(feature = "gpu"))]
+fn log_backend_impl() -> bool {
+    tracing::info!("GPU feature is disabled, using CPU backend");
+    false
+}

coprocessor/fhevm-engine/sns-worker/Cargo.toml

@@ -50,6 +50,7 @@ name = "sns_worker"
 path = "src/bin/sns_worker.rs"
 
 [features]
+gpu = ["tfhe/gpu", "fhevm-engine-common/gpu", "test-harness/gpu"]
 test_decrypt_128 = []
 test_s3_use_handle_as_key = []
 

coprocessor/fhevm-engine/sns-worker/src/keyset.rs

@@ -23,9 +23,11 @@ pub(crate) async fn fetch_keyset(
     }
 
     info!(tenant_api_key, "Cache miss");
+
     let Some((client_key, server_key)) = fetch_keys(pool, tenant_api_key).await? else {
         return Ok(None);
     };
+
     let key_set: KeySet = KeySet {
         client_key,
         server_key,
@@ -45,21 +47,35 @@ pub(crate) async fn fetch_keyset(
 pub async fn fetch_keys(
     pool: &PgPool,
     tenant_api_key: &String,
-) -> anyhow::Result<Option<(Option<tfhe::ClientKey>, tfhe::ServerKey)>> {
+) -> anyhow::Result<Option<(Option<tfhe::ClientKey>, crate::ServerKey)>> {
     let blob = read_keys_from_large_object(
         pool,
         tenant_api_key,
         "sns_pk",
         SKS_KEY_WITH_NOISE_SQUASHING_SIZE,
     )
     .await?;
-    info!(bytes_len = blob.len(), "Retrieved sns_pk");
+    info!(
+        bytes_len = blob.len(),
+        "Fetched sns_pk/sks_ns bytes from LOB"
+    );
     if blob.is_empty() {
         return Ok(None);
     }
 
+    #[cfg(not(feature = "gpu"))]
     let server_key: tfhe::ServerKey = safe_deserialize_sns_key(&blob)?;
 
+    #[cfg(feature = "gpu")]
+    let server_key = {
+        let compressed_server_key: tfhe::CompressedServerKey = safe_deserialize_sns_key(&blob)?;
+        info!("Deserialized sns_pk/sks_ns to CompressedServerKey");
+
+        let server_key = compressed_server_key.decompress_to_gpu();
+        info!("Decompressed sns_pk/sks_ns to CudaServerKey");
+        server_key
+    };
+
     // Optionally retrieve the ClientKey for testing purposes
     let client_key = fetch_client_key(pool, tenant_api_key).await?;
     Ok(Some((client_key, server_key)))

coprocessor/fhevm-engine/sns-worker/src/lib.rs

@@ -27,7 +27,6 @@ use fhevm_engine_common::{
     utils::{to_hex, DatabaseURL},
 };
 use futures::join;
-use serde::{Deserialize, Serialize};
 use sqlx::{Postgres, Transaction};
 use thiserror::Error;
 use tokio::{
@@ -51,10 +50,16 @@ pub const UPLOAD_QUEUE_SIZE: usize = 20;
 pub const SAFE_SER_LIMIT: u64 = 1024 * 1024 * 66;
 pub type InternalEvents = Option<tokio::sync::mpsc::Sender<&'static str>>;
 
-#[derive(Serialize, Deserialize, Clone)]
+#[cfg(feature = "gpu")]
+type ServerKey = tfhe::CudaServerKey;
+#[cfg(not(feature = "gpu"))]
+type ServerKey = tfhe::ServerKey;
+
+#[derive(Clone)]
 pub struct KeySet {
-    pub server_key: tfhe::ServerKey,
+    /// Optional ClientKey for decrypting on testing
     pub client_key: Option<tfhe::ClientKey>,
+    pub server_key: ServerKey,
 }
 
 #[derive(Clone)]
@@ -117,7 +122,7 @@ pub struct Config {
     pub pg_auto_explain_with_min_duration: Option<Duration>,
 }
 
-#[derive(Debug, Clone, Copy, Default)]
+#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
 pub enum SchedulePolicy {
     Sequential,
     #[default]
@@ -501,7 +506,8 @@ pub async fn run_all(
         mpsc::channel::<UploadJob>(10 * config.s3.max_concurrent_uploads as usize);
 
     let rayon_threads = rayon::current_num_threads();
-    info!(config = %config, rayon_threads, "Starting SNS worker");
+    let gpu_enabled = fhevm_engine_common::utils::log_backend();
+    info!(gpu_enabled, rayon_threads, config = %config, "Starting SNS worker");
 
     if !config.service_name.is_empty() {
         if let Err(err) = telemetry::setup_otlp(&config.service_name) {

coprocessor/fhevm-engine/sns-worker/src/tests/mod.rs

@@ -231,6 +231,7 @@ async fn run_batch_computations(
 
 #[tokio::test]
 #[serial(db)]
+#[cfg(not(feature = "gpu"))]
 async fn test_lifo_mode() {
     init_tracing();
 
@@ -311,6 +312,7 @@ async fn test_lifo_mode() {
 
 #[tokio::test]
 #[serial(db)]
+#[cfg(not(feature = "gpu"))]
 async fn test_garbage_collect() {
     init_tracing();
 
@@ -433,7 +435,15 @@ async fn setup(enable_compression: bool) -> anyhow::Result<TestEnvironment> {
         .await?;
 
     // Set up S3 storage
-    let (s3_instance, s3_client) = setup_localstack(&conf).await?;
+    let (s3_instance, s3_client) = if cfg!(feature = "gpu") {
+        info!("GPU feature is enabled, avoid testing S3-related functionality");
+        (
+            None,
+            aws_sdk_s3::Client::new(&aws_config::load_defaults(BehaviorVersion::latest()).await),
+        )
+    } else {
+        setup_localstack(&conf).await?
+    };
 
     let token = db_instance.parent_token.child_token();
     let config: Config = conf.clone();
@@ -684,6 +694,7 @@ async fn assert_ciphertext128(
 }
 
 /// Asserts that ciphertext exists in S3
+#[cfg(not(feature = "gpu"))]
 async fn assert_ciphertext_uploaded(
     test_env: &TestEnvironment,
     bucket: &String,
@@ -700,7 +711,18 @@ async fn assert_ciphertext_uploaded(
     .await;
 }
 
+#[cfg(feature = "gpu")]
+async fn assert_ciphertext_uploaded(
+    _test_env: &TestEnvironment,
+    _bucket: &String,
+    _handle: &Vec<u8>,
+    _expected_ct_len: Option<i64>,
+) {
+    // No-op when GPU feature is enabled
+}
+
 /// Asserts that the number of ciphertext128 objects in S3 matches the expected count
+#[cfg(not(feature = "gpu"))]
 async fn assert_ciphertext_s3_object_count(
     test_env: &TestEnvironment,
     bucket: &String,
@@ -710,6 +732,15 @@ async fn assert_ciphertext_s3_object_count(
         .await;
 }
 
+#[cfg(feature = "gpu")]
+async fn assert_ciphertext_s3_object_count(
+    _te: &TestEnvironment,
+    _bucket: &String,
+    _expected_count: i64,
+) {
+    // No-op when GPU feature is enabled
+}
+
 fn build_test_config(url: DatabaseURL, enable_compression: bool) -> Config {
     let batch_limit = std::env::var("BATCH_LIMIT")
         .ok()

coprocessor/fhevm-engine/test-harness/src/db_utils.rs

@@ -138,6 +138,9 @@ pub async fn setup_test_user(
     pool: &sqlx::PgPool,
     with_sns_pk: bool,
 ) -> Result<(), Box<dyn std::error::Error>> {
+    let gpu_enabled = cfg!(feature = "gpu");
+    info!(gpu_enabled, "Setting up test user...");
+
     let (sks, cks, pks, pp, sns_pk) = if !cfg!(feature = "gpu") {
         (
             "../fhevm-keys/sks",
@@ -152,7 +155,7 @@ pub async fn setup_test_user(
             "../fhevm-keys/gpu-cks",
             "../fhevm-keys/gpu-pks",
             "../fhevm-keys/gpu-pp",
-            "../fhevm-keys/sns_pk",
+            "../fhevm-keys/gpu-csks",
         )
     };
     let sks = tokio::fs::read(sks).await.expect("can't read sks key");

coprocessor/fhevm-engine/tfhe-worker/src/lib.rs

@@ -81,7 +81,9 @@ pub async fn async_main(
     }
 
     if args.run_bg_worker {
-        info!(target: "async_main", "Initializing background worker");
+        let gpu_enabled = fhevm_engine_common::utils::log_backend();
+        info!(target: "async_main", gpu_enabled,  "Initializing background worker");
+
         set.spawn(tfhe_worker::run_tfhe_worker(
             args.clone(),
             health_check.clone(),