From 929a2fc718a16d1c19cbc44fb0966f0c1c3f5903 Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <funtowiczmo@gmail.com>
Date: Thu, 28 Nov 2024 14:53:46 +0100
Subject: [PATCH] feat(backend): add some test to the backend for core
 allocation

---
 backends/llamacpp/build.rs       |  2 +-
 backends/llamacpp/src/backend.rs | 76 ++++++++++++++++++++++++++++++++
 2 files changed, 77 insertions(+), 1 deletion(-)

diff --git a/backends/llamacpp/build.rs b/backends/llamacpp/build.rs
index 023ccfba..e22fa07c 100644
--- a/backends/llamacpp/build.rs
+++ b/backends/llamacpp/build.rs
@@ -94,7 +94,7 @@ fn main() {
         .unwrap_or(out_dir.join("dist"));
 
     // Build the backend
-    let deps_path = build_backend(is_debug, opt_level, out_dir.as_path(), &install_path);
+    let _ = build_backend(is_debug, opt_level, out_dir.as_path(), &install_path);
 
     // Build the FFI layer calling the backend above
     build_ffi_layer(is_debug, &install_path);
diff --git a/backends/llamacpp/src/backend.rs b/backends/llamacpp/src/backend.rs
index e662e207..d8f28ab9 100644
--- a/backends/llamacpp/src/backend.rs
+++ b/backends/llamacpp/src/backend.rs
@@ -28,6 +28,7 @@ use tracing::{debug, error, info};
 ///
 /// returns: usize Integer greater than 0 representing the number of CPU cores on the machine
 ///
+#[cfg(not(test))]
 fn get_num_cores() -> usize {
     match option_env!("TGI_USE_PHYSICAL_CORES")
         .unwrap_or("OFF")
@@ -45,6 +46,18 @@ fn get_num_cores() -> usize {
     }
 }
 
+#[cfg(test)]
+fn get_num_cores() -> usize {
+    match option_env!("TGI_USE_PHYSICAL_CORES")
+        .unwrap_or("OFF")
+        .to_uppercase()
+        .as_str()
+    {
+        "ON" => 16,
+        _ => 32,
+    }
+}
+
 /// Subdivide the set of CPU cores available on the system to equal, non-overlapping, subsets of CPU cores
 ///
 /// # Arguments
@@ -417,3 +430,66 @@ impl Backend for LlamaCppBackend {
         true
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::backend::{get_cores_allocation, get_num_cores};
+
+    fn test_get_num_cores() {
+        std::env::set_var("TGI_USE_PHYSICAL_CORES", "OFF");
+        assert_eq!(get_num_cores(), 32);
+
+        std::env::set_var("TGI_USE_PHYSICAL_CORES", "ON");
+        assert_eq!(get_num_cores(), 16);
+    }
+
+    fn test_get_cores_allocation_single_instance() {
+        std::env::set_var("TGI_USE_PHYSICAL_CORES", "OFF");
+        let smt_allocation = get_cores_allocation(0);
+        assert_eq!(smt_allocation.len(), 1);
+        assert_eq!(
+            smt_allocation[0].clone().collect::<Vec<_>>(),
+            (0..32).collect::<Vec<_>>()
+        );
+
+        std::env::set_var("TGI_USE_PHYSICAL_CORES", "ON");
+        let smt_allocation = get_cores_allocation(0);
+        assert_eq!(smt_allocation.len(), 1);
+        assert_eq!(
+            smt_allocation[0].clone().collect::<Vec<_>>(),
+            (0..16).collect::<Vec<_>>()
+        );
+    }
+
+    fn test_get_cores_allocation_multi_instances() {
+        for cores_per_instance in [1, 2, 4, 8, 16, 3, 7] {
+            std::env::set_var("TGI_USE_PHYSICAL_CORES", "OFF");
+
+            let num_instances = 32 / cores_per_instance;
+            let smt_allocation = get_cores_allocation(cores_per_instance);
+
+            for i in 0..num_instances {
+                let start = i * cores_per_instance;
+                let end = start + cores_per_instance;
+                assert_eq!(
+                    smt_allocation[i].clone().collect::<Vec<_>>(),
+                    (start..end).collect::<Vec<_>>()
+                );
+            }
+
+            std::env::set_var("TGI_USE_PHYSICAL_CORES", "ON");
+            let num_instances = 16 / cores_per_instance;
+            let smt_allocation = get_cores_allocation(cores_per_instance);
+            assert_eq!(smt_allocation.len(), num_instances);
+
+            for i in 0..num_instances {
+                let start = i * cores_per_instance;
+                let end = start + cores_per_instance;
+                assert_eq!(
+                    smt_allocation[i].clone().collect::<Vec<_>>(),
+                    (start..end).collect::<Vec<_>>()
+                );
+            }
+        }
+    }
+}