From c27946052f503f10203cdf1062b1537bc242a2ac Mon Sep 17 00:00:00 2001
From: Maria Kuklina <maria@fluence.one>
Date: Thu, 29 Aug 2024 14:08:51 +0200
Subject: [PATCH] allocate memory based on physical cores

---
 crates/core-distributor/src/types.rs | 8 ++++++++
 crates/vm-utils/src/vm_utils.rs      | 5 ++++-
 crates/workers/src/workers.rs        | 8 +++++---
 3 files changed, 17 insertions(+), 4 deletions(-)

diff --git a/crates/core-distributor/src/types.rs b/crates/core-distributor/src/types.rs
index 94ee26759c..f745f6ddf6 100644
--- a/crates/core-distributor/src/types.rs
+++ b/crates/core-distributor/src/types.rs
@@ -107,6 +107,14 @@ impl Assignment {
             .collect::<Vec<_>>()
     }
 
+    pub fn physical_core_count(&self) -> usize {
+        self.cuid_cores
+            .values()
+            .map(|cores| cores.physical_core_id)
+            .collect::<BTreeSet<PhysicalCoreId>>()
+            .len()
+    }
+
     pub fn pin_current_thread_with(&self, thread_pinner: &dyn ThreadPinner) {
         thread_pinner.pin_current_thread_to_cpuset(&self.logical_core_ids());
     }
diff --git a/crates/vm-utils/src/vm_utils.rs b/crates/vm-utils/src/vm_utils.rs
index c6c7670d52..6299d11279 100644
--- a/crates/vm-utils/src/vm_utils.rs
+++ b/crates/vm-utils/src/vm_utils.rs
@@ -17,6 +17,7 @@ pub struct CreateVMDomainParams {
     name: String,
     image: PathBuf,
     cpus: NonEmpty<LogicalCoreId>,
+    cores_num: usize,
     bridge_name: String,
     allow_gpu: bool,
 }
@@ -26,6 +27,7 @@ impl CreateVMDomainParams {
         name: String,
         image: PathBuf,
         cpus: NonEmpty<LogicalCoreId>,
+        cores_num: usize,
         bridge_name: String,
         allow_gpu: bool,
     ) -> Self {
@@ -33,6 +35,7 @@ impl CreateVMDomainParams {
             name,
             image,
             cpus,
+            cores_num,
             bridge_name,
             allow_gpu,
         }
@@ -346,7 +349,7 @@ fn prepare_xml(
         }
         mapping.push_str(format!("<vcpupin vcpu='{index}' cpuset='{logical_id}'/>").as_str());
     }
-    let memory_in_kb = params.cpus.len() * 4 * 1024 * 1024; // 4Gbs per core
+    let memory_in_kb = params.cores_num * 4 * 1024 * 1024; // 4Gbs per core
     let gpu_pci_configuration = gpu_pci_location
         .iter()
         .map(prepare_pci_config)
diff --git a/crates/workers/src/workers.rs b/crates/workers/src/workers.rs
index cf5f9b05cc..367a066602 100644
--- a/crates/workers/src/workers.rs
+++ b/crates/workers/src/workers.rs
@@ -673,14 +673,16 @@ impl Workers {
         image: &Path,
         vm_config: &VmConfig,
     ) -> Result<String, WorkersError> {
-        let assignment = {
+        let (cores_count, assignment) = {
             let guard = self.assignments.read();
             let assignment = guard
                 .get(&worker_id)
                 .ok_or_else(|| WorkersError::WorkerNotFound(worker_id))?;
 
-            NonEmpty::from_vec(assignment.logical_core_ids())
-                .ok_or_else(|| WorkersError::WrongAssignment)?
+            let logical_cores = NonEmpty::from_vec(assignment.logical_core_ids())
+                .ok_or_else(|| WorkersError::WrongAssignment)?;
+            let physical_cores_count = assignment.physical_core_count();
+            (physical_cores_count, logical_cores)
         };
 
         let file_name = &image