Add option for CPU limit to be a callable as well

Author: krinsman

nbresuse/__init__.py

@@ -93,7 +93,8 @@ def _mem_limit_default(self):
         """,
     ).tag(config=True)
 
-    cpu_limit = Float(
+    cpu_limit = Union(
+        trait_types=[Float(), Callable()],
         default_value=0,
         help="""
         CPU usage limit to display to the user.

nbresuse/prometheus.py

@@ -44,13 +44,15 @@ def __call__(self, *args, **kwargs):
     def apply_memory_limits(self, metrics: Optional[MemoryMetrics]) -> Optional[MemoryMetrics]:
         if metrics is not None:
             if callable(self.config.mem_limit):
-                metrics.max_memory = self.config.mem_limit(rss=metrics.max_memory)
+                metrics.max_memory = self.config.mem_limit(rss=metrics.current_memory)
             elif self.config.mem_limit > 0:  # mem_limit is an Int
                 metrics.max_memory = self.config.mem_limit
         return metrics
 
     def apply_cpu_limits(self, metrics: Optional[CPUMetrics]) -> Optional[CPUMetrics]:
         if metrics is not None:
-            if self.config.cpu_limit > 0:
+            if callable(self.config.cpu_limit):
+                metrics.cpu_max = self.config.cpu_limit(cpu_percent=metrics.cpu_usage)
+            elif self.config.cpu_limit > 0.0:  # cpu_limit is a Float
                 metrics.cpu_max = self.config.cpu_limit
         return metrics