Move estimator trait

Use frequency service in perf estimator

7 files changed, 38 insertions, 18 deletions

diff --git a/src/energy.rs b/src/energy.rs
index 6490efd..d3845a8 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -1,7 +1,9 @@
 mod budget;
+pub mod estimator;
 pub mod rapl;
 mod trackers;
 
+use crate::energy::estimator::Estimator;
 use std::collections::{BTreeSet, HashMap};
 use std::ops::RangeInclusive;
 use std::sync::{mpsc, Arc, RwLock};
@@ -14,7 +16,7 @@ use crate::Pid;
 use dashmap::DashMap;
 
 pub use budget::BudgetPolicy;
-pub use trackers::{Estimator, KernelDriver, PerfEstimator};
+pub use trackers::{KernelDriver, PerfEstimator};
 
 pub enum Request {
     NewTask(Pid),
@@ -248,7 +250,7 @@ pub fn start_energy_service(
 
     // Create the appropriate estimator based on configuration
     let estimator: Box<dyn Estimator> = if use_mocking {
-        Box::new(PerfEstimator::default())
+        Box::new(PerfEstimator::new(shared_cpu_current_frequencies.clone()))
     } else {
         Box::new(KernelDriver::default())
     };
diff --git a/src/energy/estimator.rs b/src/energy/estimator.rs
index f7f1e5d..85cb0eb 100644
--- a/src/energy/estimator.rs
+++ b/src/energy/estimator.rs
@@ -1,5 +1,5 @@
 pub trait Estimator: Send + 'static {
     fn start_trace(&mut self, pid: u64);
     fn stop_trace(&mut self, pid: u64);
-    fn read_consumption(&mut self, pid: u64) -> u64;
+    fn read_consumption(&mut self, pid: u64, cpu: i32) -> Option<f64>;
 }
diff --git a/src/energy/trackers.rs b/src/energy/trackers.rs
index 1665177..e510e8d 100644
--- a/src/energy/trackers.rs
+++ b/src/energy/trackers.rs
@@ -7,8 +7,8 @@ pub use kernel::*;
 pub use mock::*;
 pub use perf::*;
 
-pub trait Estimator: Send + 'static {
-    fn start_trace(&mut self, pid: u64);
-    fn stop_trace(&mut self, pid: u64);
-    fn read_consumption(&mut self, pid: u64) -> Option<f64>;
-}
+// pub trait Estimator: Send + 'static {
+//     fn start_trace(&mut self, pid: u64);
+//     fn stop_trace(&mut self, pid: u64);
+//     fn read_consumption(&mut self, pid: u64) -> Option<f64>;
+// }
diff --git a/src/energy/trackers/kernel.rs b/src/energy/trackers/kernel.rs
index 2a19e4c..89a0d89 100644
--- a/src/energy/trackers/kernel.rs
+++ b/src/energy/trackers/kernel.rs
@@ -38,7 +38,7 @@ impl Estimator for KernelDriver {
         let _ = STOP_TRACE.ioctl(&mut self.file, &pid);
     }
 
-    fn read_consumption(&mut self, pid: u64) -> Option<f64> {
+    fn read_consumption(&mut self, pid: u64, _cpu: i32) -> Option<f64> {
         let mut arg = pid;
         if READ_POWER.ioctl(&mut self.file, &mut arg).is_ok() {
             Some(arg as f64)
diff --git a/src/energy/trackers/mock.rs b/src/energy/trackers/mock.rs
index e3ef377..597cab0 100644
--- a/src/energy/trackers/mock.rs
+++ b/src/energy/trackers/mock.rs
@@ -1,5 +1,4 @@
-use super::Estimator;
-
+use crate::energy::estimator::Estimator;
 #[derive(Default)]
 pub struct MockEstimator;
 
@@ -8,7 +7,7 @@ impl Estimator for MockEstimator {
 
     fn stop_trace(&mut self, _pid: u64) {}
 
-    fn read_consumption(&mut self, _pid: u64) -> Option<f64> {
+    fn read_consumption(&mut self, _pid: u64, _cpu: i32) -> Option<f64> {
         Some(14.)
     }
 }
diff --git a/src/energy/trackers/perf.rs b/src/energy/trackers/perf.rs
index d340e1f..1c7943c 100644
--- a/src/energy/trackers/perf.rs
+++ b/src/energy/trackers/perf.rs
@@ -6,22 +6,38 @@ use perf_event::{
     Builder, Counter, Group,
 };
 
+use crate::freq::FrequencyKHZ;
+use std::sync::{Arc, RwLock};
+
 use crate::energy::{rapl, Estimator};
-type ArrayBackend = burn_ndarray::NdArray<f32>;
+use crate::model::ArrayBackend;
 
 pub struct PerfEstimator {
     registry: HashMap<u64, Counters>,
     model: crate::model::Net<ArrayBackend>,
     device: <ArrayBackend as burn::prelude::Backend>::Device,
+    shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>,
 }
 
-impl Default for PerfEstimator {
-    fn default() -> Self {
+// impl Default for PerfEstimator {
+//     fn default() -> Self {
+//         let model = crate::model::load_model();
+//         Self {
+//             registry: Default::default(),
+//             model,
+//             device: Default::default(),
+//         }
+//     }
+// }
+
+impl PerfEstimator {
+    pub fn new(shared_cpu_current_frequencies: Arc<RwLock<Vec<FrequencyKHZ>>>) -> Self {
         let model = crate::model::load_model();
         Self {
             registry: Default::default(),
             model,
             device: Default::default(),
+            shared_cpu_current_frequencies,
         }
     }
 }
@@ -105,7 +121,7 @@ impl Estimator for PerfEstimator {
         self.registry.remove(&pid);
     }
 
-    fn read_consumption(&mut self, pid: u64) -> Option<f64> {
+    fn read_consumption(&mut self, pid: u64, cpu: i32) -> Option<f64> {
         let Some(counters) = self.registry.get_mut(&pid) else {
             return None;
         };
@@ -130,8 +146,9 @@ impl Estimator for PerfEstimator {
         //dbg!(time_running_ns);
         //dbg!(counters.old_time.elapsed().as_millis());
 
-        // TODO:: Respect per core frequency for the running task
-        let mut values = vec![crate::benchmark::read_cpu_frequency(0).unwrap()];
+        let mut values =
+            vec![(self.shared_cpu_current_frequencies.read().unwrap()[cpu as usize] / 1000) as f64];
+        //let mut values = vec![crate::benchmark::read_cpu_frequency(0).unwrap()];
         for ty in counters.counters.iter().take(7) {
             let count: u64 = counts[&ty];
             values.push((count as f64) * correction_factor);
diff --git a/src/freq.rs b/src/freq.rs
index 7bba4f8..20a8619 100644
--- a/src/freq.rs
+++ b/src/freq.rs
@@ -208,6 +208,8 @@ impl SysFSFrequencyService {
 
             loop {
                 self.handle_requests().unwrap();
+                *self.cpu_current_frequencies.write().unwrap() =
+                    self.get_current_frequencies().unwrap();
                 thread::sleep(self.update_interval);
             }
         });