PID controller w/ mem & cpu

Author: Sushisource

core/Cargo.toml

@@ -48,6 +48,7 @@ opentelemetry_sdk = { version = "0.22", features = ["rt-tokio", "metrics"], opti
 opentelemetry-otlp = { version = "0.15", features = ["tokio", "metrics"], optional = true }
 opentelemetry-prometheus = { version = "0.15", optional = true }
 parking_lot = { version = "0.12", features = ["send_guard"] }
+pid = "4.0"
 pin-project = "1.0"
 prometheus = "0.13"
 prost = { workspace = true }

core/src/worker/slot_supplier/resource_based.rs

@@ -2,8 +2,10 @@ use crossbeam_utils::atomic::AtomicCell;
 use parking_lot::Mutex;
 use std::{
     marker::PhantomData,
-    ops::SubAssign,
-    sync::{atomic::AtomicU64, Arc},
+    sync::{
+        atomic::{AtomicU32, AtomicU64, Ordering},
+        Arc,
+    },
     time::{Duration, Instant},
 };
 use temporal_sdk_core_api::worker::{
@@ -14,40 +16,57 @@ use tokio::sync::watch;
 
 pub struct ResourceBasedSlots<MI> {
     target_mem_usage: f64,
-    assumed_maximum_marginal_contribution: f64,
-    mem_info_supplier: MI,
-    max: usize,
+    target_cpu_usage: f32,
+    sys_info_supplier: MI,
 }
 pub struct ResourceBasedSlotsForType<MI, SK> {
     inner: Arc<ResourceBasedSlots<MI>>,
     minimum: usize,
+    /// Maximum amount of slots of this type permitted
+    max: usize,
     /// Minimum time we will wait (after passing the minimum slots number) between handing out new
     /// slots
     ramp_throttle: Duration,
 
+    pids: Arc<Mutex<PidControllers>>,
     last_slot_issued_tx: watch::Sender<Instant>,
     last_slot_issued_rx: watch::Receiver<Instant>,
     _slot_kind: PhantomData<SK>,
 }
+struct PidControllers {
+    mem: pid::Pid<f64>,
+    cpu: pid::Pid<f32>,
+}
 
-impl ResourceBasedSlots<SysinfoMem> {
-    pub fn new(target_mem_usage: f64, marginal_contribution: f64, max: usize) -> Self {
+impl ResourceBasedSlots<RealSysInfo> {
+    pub fn new(target_mem_usage: f64, target_cpu_usage: f32) -> Self {
         Self {
             target_mem_usage,
-            assumed_maximum_marginal_contribution: marginal_contribution,
-            mem_info_supplier: SysinfoMem::new(),
-            max,
+            target_cpu_usage,
+            sys_info_supplier: RealSysInfo::new(),
         }
     }
 }
 
+impl PidControllers {
+    fn new(mem_target: f64, cpu_target: f32) -> Self {
+        let mut mem = pid::Pid::new(mem_target, 100.0);
+        mem.p(5.0, 100).i(0.0, 100).d(1.0, 100);
+        let mut cpu = pid::Pid::new(cpu_target, 100.0);
+        cpu.p(5.0, 100.).i(0.0, 100.).d(1.0, 100.);
+        Self { mem, cpu }
+    }
+}
+
 trait MemoryInfo {
     /// Return total available system memory in bytes
     fn total_mem(&self) -> u64;
     /// Return memory used by this process in bytes
     // TODO: probably needs to be just overall used... won't work w/ subprocesses for example
     fn process_used_mem(&self) -> u64;
 
+    fn used_cpu_percent(&self) -> f32;
+
     fn process_used_percent(&self) -> f64 {
         self.process_used_mem() as f64 / self.total_mem() as f64
     }
@@ -79,7 +98,9 @@ where
 
     fn try_reserve_slot(&self, ctx: &dyn SlotReservationContext) -> Option<SlotSupplierPermit> {
         if self.time_since_last_issued() > self.ramp_throttle
-            && self.inner.can_reserve(ctx.num_issued_slots())
+            && ctx.num_issued_slots() < self.max
+            && self.pid_decision()
+            && self.inner.can_reserve()
         {
             let _ = self.last_slot_issued_tx.send(Instant::now());
             Some(SlotSupplierPermit::NoData)
@@ -102,84 +123,105 @@ where
     MI: MemoryInfo + Send + Sync,
     SK: SlotKind + Send + Sync,
 {
+    fn new(
+        inner: Arc<ResourceBasedSlots<MI>>,
+        minimum: usize,
+        max: usize,
+        ramp_throttle: Duration,
+    ) -> Self {
+        let (tx, rx) = watch::channel(Instant::now());
+        Self {
+            minimum,
+            max,
+            ramp_throttle,
+            pids: Arc::new(Mutex::new(PidControllers::new(
+                inner.target_mem_usage,
+                inner.target_cpu_usage,
+            ))),
+            inner,
+            last_slot_issued_tx: tx,
+            last_slot_issued_rx: rx,
+            _slot_kind: PhantomData,
+        }
+    }
+
     fn time_since_last_issued(&self) -> Duration {
         Instant::now()
             .checked_duration_since(*self.last_slot_issued_rx.borrow())
             .unwrap_or_default()
     }
+
+    /// Returns true if the pid controllers think a new slot should be given out
+    fn pid_decision(&self) -> bool {
+        let mut pids = self.pids.lock();
+        let mem_output = pids
+            .mem
+            .next_control_output(self.inner.sys_info_supplier.process_used_percent())
+            .output;
+        let cpu_output = pids
+            .cpu
+            .next_control_output(self.inner.sys_info_supplier.used_cpu_percent())
+            .output;
+        mem_output > 0.25 && cpu_output > 0.25
+    }
 }
 
 impl<MI> WorkflowCacheSizer for ResourceBasedSlots<MI>
 where
-    MI: MemoryInfo + Sync,
+    MI: MemoryInfo + Sync + Send,
 {
-    fn can_allow_workflow(
-        &self,
-        slots_info: &WorkflowSlotsInfo,
-        _new_task: &WorkflowSlotInfo,
-    ) -> bool {
-        self.can_reserve(slots_info.used_slots.len())
+    fn can_allow_workflow(&self, _: &WorkflowSlotsInfo, _: &WorkflowSlotInfo) -> bool {
+        self.can_reserve()
     }
 }
 
-impl<MI: MemoryInfo + Sync> ResourceBasedSlots<MI> {
+impl<MI: MemoryInfo + Sync + Send> ResourceBasedSlots<MI> {
     // TODO: Can just be an into impl probably?
     pub fn as_kind<SK: SlotKind + Send + Sync>(
         self: &Arc<Self>,
         minimum: usize,
+        max: usize,
         ramp_throttle: Duration,
     ) -> Arc<ResourceBasedSlotsForType<MI, SK>> {
-        let (tx, rx) = watch::channel(Instant::now());
-        Arc::new(ResourceBasedSlotsForType {
-            inner: self.clone(),
+        Arc::new(ResourceBasedSlotsForType::new(
+            self.clone(),
             minimum,
+            max,
             ramp_throttle,
-            last_slot_issued_tx: tx,
-            last_slot_issued_rx: rx,
-            _slot_kind: PhantomData,
-        })
+        ))
     }
 
     pub fn into_kind<SK: SlotKind + Send + Sync>(self) -> ResourceBasedSlotsForType<MI, SK> {
-        let (tx, rx) = watch::channel(Instant::now());
-        ResourceBasedSlotsForType {
-            inner: Arc::new(self),
-            // TODO: Configure
-            minimum: 1,
-            ramp_throttle: Duration::from_millis(0),
-            last_slot_issued_tx: tx,
-            last_slot_issued_rx: rx,
-            _slot_kind: PhantomData,
-        }
+        // TODO: remove or parameterize
+        ResourceBasedSlotsForType::new(Arc::new(self), 1, 1000, Duration::from_millis(0))
     }
 
-    fn can_reserve(&self, num_used: usize) -> bool {
-        if num_used > self.max {
-            return false;
-        }
-        self.mem_info_supplier.process_used_percent() + self.assumed_maximum_marginal_contribution
-            <= self.target_mem_usage
+    fn can_reserve(&self) -> bool {
+        self.sys_info_supplier.process_used_percent() <= self.target_mem_usage
     }
 }
 
 #[derive(Debug)]
-pub struct SysinfoMem {
+pub struct RealSysInfo {
     sys: Mutex<sysinfo::System>,
     pid: sysinfo::Pid,
     cur_mem_usage: AtomicU64,
+    cur_cpu_usage: AtomicU32,
     last_refresh: AtomicCell<Instant>,
 }
-impl SysinfoMem {
+impl RealSysInfo {
     fn new() -> Self {
         let mut sys = sysinfo::System::new();
         let pid = sysinfo::get_current_pid().expect("get pid works");
         sys.refresh_processes();
         sys.refresh_memory();
+        sys.refresh_cpu();
         Self {
             sys: Default::default(),
             last_refresh: AtomicCell::new(Instant::now()),
             pid,
             cur_mem_usage: AtomicU64::new(0),
+            cur_cpu_usage: AtomicU32::new(0),
         }
     }
     fn refresh_if_needed(&self) {
@@ -189,23 +231,32 @@ impl SysinfoMem {
             let mut lock = self.sys.lock();
             lock.refresh_memory();
             lock.refresh_processes();
+            lock.refresh_cpu_usage();
             let proc = lock.process(self.pid).expect("exists");
             self.cur_mem_usage
-                .store(dbg!(proc.memory()), std::sync::atomic::Ordering::Release);
+                .store(dbg!(proc.memory()), Ordering::Release);
+            self.cur_cpu_usage.store(
+                dbg!(lock.global_cpu_info().cpu_usage()).to_bits(),
+                Ordering::Release,
+            );
             self.last_refresh.store(Instant::now())
         }
     }
 }
-impl MemoryInfo for SysinfoMem {
+impl MemoryInfo for RealSysInfo {
     fn total_mem(&self) -> u64 {
         self.refresh_if_needed();
         self.sys.lock().total_memory()
     }
 
     fn process_used_mem(&self) -> u64 {
         self.refresh_if_needed();
-        self.cur_mem_usage
-            .load(std::sync::atomic::Ordering::Acquire)
+        self.cur_mem_usage.load(Ordering::Acquire)
+    }
+
+    fn used_cpu_percent(&self) -> f32 {
+        self.refresh_if_needed();
+        f32::from_bits(self.cur_cpu_usage.load(Ordering::Acquire))
     }
 }
 
@@ -235,6 +286,10 @@ mod tests {
         fn process_used_mem(&self) -> u64 {
             self.used.load(Ordering::Acquire)
         }
+
+        fn used_cpu_percent(&self) -> f32 {
+            todo!()
+        }
     }
     struct FakeResCtx {}
     impl SlotReservationContext for FakeResCtx {
@@ -248,9 +303,8 @@ mod tests {
         let (fmis, used) = FakeMIS::new();
         let rbs = ResourceBasedSlots {
             target_mem_usage: 0.8,
-            assumed_maximum_marginal_contribution: 0.1,
-            mem_info_supplier: fmis,
-            max: 1000,
+            target_cpu_usage: 1.0,
+            sys_info_supplier: fmis,
         }
         .into_kind::<WorkflowSlotKind>();
         assert!(rbs.try_reserve_slot(&FakeResCtx {}).is_some());
@@ -264,9 +318,8 @@ mod tests {
         used.store(90_000, Ordering::Release);
         let rbs = ResourceBasedSlots {
             target_mem_usage: 0.8,
-            assumed_maximum_marginal_contribution: 0.1,
-            mem_info_supplier: fmis,
-            max: 1000,
+            target_cpu_usage: 1.0,
+            sys_info_supplier: fmis,
         }
         .into_kind::<WorkflowSlotKind>();
         let order = crossbeam_queue::ArrayQueue::new(2);

tests/heavy_tests.rs

@@ -10,7 +10,9 @@ use temporal_sdk_core_protos::{
     coresdk::{workflow_commands::ActivityCancellationType, AsJsonPayloadExt},
     temporal::api::enums::v1::WorkflowIdReusePolicy,
 };
-use temporal_sdk_core_test_utils::{workflows::la_problem_workflow, CoreWfStarter};
+use temporal_sdk_core_test_utils::{
+    canned_histories::cancel_scheduled_activity, workflows::la_problem_workflow, CoreWfStarter,
+};
 
 mod fuzzy_workflow;
 
@@ -96,11 +98,16 @@ async fn chunky_activities() {
     //     .worker_config
     //     .max_outstanding_activities(25)
     //     .max_outstanding_workflow_tasks(25);
-    let resource_slots = Arc::new(ResourceBasedSlots::new(0.5, 0.01, 1024));
+    // TODO: Fix /1 or /100 thing
+    let resource_slots = Arc::new(ResourceBasedSlots::new(0.5, 90.0));
     starter
         .worker_config
-        .workflow_task_slot_supplier(resource_slots.as_kind(25, Duration::from_millis(0)))
-        .activity_task_slot_supplier(resource_slots.as_kind(5, Duration::from_millis(100)));
+        .workflow_task_slot_supplier(resource_slots.as_kind(
+            25,
+            WORKFLOWS,
+            Duration::from_millis(0),
+        ))
+        .activity_task_slot_supplier(resource_slots.as_kind(5, 1000, Duration::from_millis(50)));
     let mut worker = starter.worker().await;
 
     let activity_id = "act-1";
@@ -177,8 +184,11 @@ async fn workflow_load() {
     let mut starter = CoreWfStarter::new("workflow_load");
     starter.worker_config.max_concurrent_wft_polls(10_usize);
     starter.worker_config.workflow_task_slot_supplier(
-        Arc::new(ResourceBasedSlots::new(0.7, 0.01, cache_size))
-            .as_kind(10, Duration::from_millis(0)),
+        Arc::new(ResourceBasedSlots::new(0.7, 0.9)).as_kind(
+            10,
+            cache_size,
+            Duration::from_millis(0),
+        ),
     );
     starter
         // .max_wft(100)