Add biases and track executing CPU

1 files changed, 34 insertions, 9 deletions

diff --git a/src/energy.rs b/src/energy.rs
index d3845a8..1b52c13 100644
--- a/src/energy.rs
+++ b/src/energy.rs
@@ -18,8 +18,9 @@ use dashmap::DashMap;
 pub use budget::BudgetPolicy;
 pub use trackers::{KernelDriver, PerfEstimator};
 
+//TODO: add new command to update process_info or make process_info shared
 pub enum Request {
-    NewTask(Pid),
+    NewTask(Pid, i32),
     RemoveTask(Pid),
     LogEnergyTask(Pid),
 }
@@ -30,6 +31,7 @@ pub struct ProcessInfo {
     tree_energy: f64,
     last_update: std::time::Instant,
     parent: Pid,
+    cpu: i32,
 }
 
 pub struct EnergyLog {
@@ -40,6 +42,7 @@ pub struct EnergyLog {
 pub struct EnergyService {
     estimator: Box<dyn Estimator>,
     budget_policy: Option<Box<dyn BudgetPolicy>>,
+    // TODO: remove
     active_processes: BTreeSet<Pid>,
     process_info: HashMap<Pid, ProcessInfo>,
     shared_budgets: Arc<DashMap<Pid, u64>>,
@@ -52,6 +55,8 @@ pub struct EnergyService {
     logged_energy: Arc<RwLock<HashMap<Pid, EnergyLog>>>,
     rapl_offset: f64,
     old_rapl: f64,
+    bias: f64,
+    offset: f64,
 }
 
 impl EnergyService {
@@ -81,6 +86,8 @@ impl EnergyService {
             logged_energy,
             rapl_offset: rapl::read_package_energy().unwrap(),
             old_rapl: 0.,
+            bias: 1.,
+            offset: 0.,
         }
     }
 
@@ -110,7 +117,10 @@ impl EnergyService {
 
     fn handle_request(&mut self, request: Request) {
         match request {
-            Request::NewTask(pid) => {
+            Request::NewTask(pid, cpu) => {
+                if self.process_info.contains_key(&pid) {
+                    return;
+                }
                 self.estimator.start_trace(pid as u64);
                 self.active_processes.insert(pid);
                 let parent = (|| {
@@ -125,12 +135,13 @@ impl EnergyService {
                         tree_energy: 0.,
                         last_update: std::time::Instant::now(),
                         parent,
+                        cpu,
                     },
                 );
                 // Initialize with default budget
                 self.shared_budgets.insert(pid, u64::MAX);
                 if !self.process_info.contains_key(&parent) && parent != 0 {
-                    self.handle_request(Request::NewTask(parent));
+                    self.handle_request(Request::NewTask(parent, cpu));
                 }
             }
             Request::RemoveTask(pid) => {
@@ -175,30 +186,44 @@ impl EnergyService {
         let old_energy = self.process_info.get(&1).unwrap().tree_energy;
         for &pid in &self.active_processes {
             if let Some(info) = self.process_info.get_mut(&pid) {
-                if let Some(energy) = self.estimator.read_consumption(pid as u64) {
-                    info.energy += energy;
+                if let Some(energy) = self.estimator.read_consumption(pid as u64, info.cpu) {
+                    info.energy += energy * self.bias;
                     let mut parent = info.parent;
                     while let Some(info) = self.process_info.get_mut(&parent) {
-                        info.tree_energy += energy;
+                        info.tree_energy += energy * self.bias;
                         info.last_update = std::time::Instant::now();
                         parent = info.parent;
                     }
                 }
             }
         }
-        if let Some(init) = self.process_info.get(&1) {
+        if let Some(init) = self.process_info.get_mut(&1) {
             let rapl = rapl::read_package_energy().unwrap() - self.rapl_offset;
             let rapl_diff = rapl - self.old_rapl;
+            let est_diff = (init.tree_energy - old_energy);
+            if est_diff < 0.1 {
+                self.offset = (self.offset + (rapl_diff - est_diff)) * 0.5;
+            }
             self.old_rapl = rapl;
+            init.tree_energy = (init.tree_energy + self.offset);
+            let offset_bias = (rapl / init.tree_energy).clamp(0.1, 2.);
+            let diff_bias = (rapl_diff / est_diff).clamp(0.1, 2.);
+            let current_bias = (offset_bias + diff_bias) * 0.5;
+            //dbg!(current_bias, self.bias);
+            self.bias = (self.bias * ((1. / 3.) * current_bias + (2. / 3.))).clamp(0.1, 20.);
             //let v = logged_energy.get(&1).unwrap();
             println!(
                 "Energy estimation: {:.1} rapl: {:.1}, est diff: {:.1} rapl diff: {:.1}",
                 init.tree_energy,
                 rapl,
-                (init.tree_energy - old_energy),
+                est_diff,
                 rapl_diff,
                 //(rapl_diff - (init.tree_energy - old_energy)) / rapl_diff * 100.,
             );
+            // if init.tree_energy > 1000. || rapl > 1000. {
+            //     self.rapl_offset += rapl;
+            //     init.tree_energy = 0.;
+            // }
         }
     }
 
@@ -277,7 +302,7 @@ pub fn start_energy_service(
     );
 
     service.run();
-    socket::start_logging_socket_service("/tmp/pm-sched", logged_energy, Duration::from_secs(1))?;
+    //socket::start_logging_socket_service("/tmp/pm-sched", logged_energy, Duration::from_secs(1))?;
 
     Ok(request_sender)
 }