[WIP]Add method for kernel manager to retrieve statistics

jupyter_client/kernelmetrics.py

@@ -0,0 +1,81 @@
+from abc import ABCMeta, abstractmethod
+from six import add_metaclass
+import re
+
+METRIC_NAME_PREFIX = "KernelMetric"
+
+
+class KernelMetricStore(ABCMeta):
+    """
+    KernelMetricStore Is acting as a store of all the metric classes created using KernelMetric base class.
+    TYPES is a dictionary that maps the names of the metrics to the metric classes themselves.
+
+    Additionally it defines defines a standard for how kernel metric classes should be named.
+    Example:
+    If the name of the class that exposes a metric is "KernelMetricMemoryUsage",
+    the dictionary will remove the must have prefix "KernelMetric",
+    split the rest by camel case letters with a "_" delimiter, and then turn them all to lower case.
+    the name that maps to the class of the metric will be "memory_usage".
+
+    *ALL KERNEL METRIC CLASSES MUST INHERIT FROM THE BASE CLASS `KernelMetric`
+    *ALL KERNEL METRIC CLASSES MUST HAVE THE PREFIX IN THEIR NAME
+
+    When we define a metric class to inherit from the `KernelMetric`, this class will automatically
+    check that its name has the need prefix, and will add it to the TYPES dict.
+    """
+    TYPES = {}
+
+    def __new__(mcs, name, bases, dct):
+        mcs._check_name_prefix(mcs, name)
+        new_metric_class = super(KernelMetricStore, mcs).__new__(mcs, name, bases, dct)
+        if new_metric_class.__name__ == "KernelMetric":
+            return new_metric_class
+        mcs._add_new_metric(mcs, new_metric_class)
+        return new_metric_class
+
+    def _add_new_metric(self, new_metric):
+        metric_name = new_metric.__name__[len(METRIC_NAME_PREFIX):]
+        metric_name = re.sub('(?!^)([A-Z][a-z]+)', r' \1', metric_name).split()
+        metric_name = "_".join(metric_name).lower()
+        self.TYPES[metric_name] = new_metric
+
+    def _check_name_prefix(self, name):
+        assert name.startswith(METRIC_NAME_PREFIX)
+
+
+@add_metaclass(KernelMetricStore)
+class KernelMetric(object):
+    """
+    An abstract base class for all of the metric classes.
+    """
+    def __init__(self, kernel_manager):
+        """
+        :param kernel_manager: desired kernel to pull the metric from
+        """
+        self.kernel_manager = kernel_manager
+
+    @abstractmethod
+    def poll(self):
+        """
+        Polls the necessary metric from the kernel and returns it
+        :return: Metric data
+        """
+        pass
+
+
+class KernelMetricMemoryUsage(KernelMetric):
+    """
+    Returns memory usage of the kernel in MB.
+    `poll` function should be overwritten if there is a
+    need for a different kind of implementation.
+    """
+    def poll(self):
+        from psutil import Process
+        kernel_pid = self.kernel_manager.kernel.pid
+        process = Process(kernel_pid)
+        mem = process.memory_info().vms
+        return self._bytes_to_mb(mem)
+
+    def _bytes_to_mb(self, num):
+        return int(num / 1024 ** 2)
+

jupyter_client/manager.py

@@ -27,6 +27,7 @@
 from .managerabc import (
     KernelManagerABC
 )
+from .kernelmetrics import KernelMetricStore
 
 
 class KernelManager(ConnectionFileMixin):
@@ -454,6 +455,37 @@ def is_alive(self):
             # we don't have a kernel
             return False
 
+    def kernel_metrics(self, *desired_metrics):
+        """
+        Returns a dict containing the results of the gathered
+        statistics.
+
+        Parameters
+        ----------
+
+        *desired_metrics: A list of strings representing the kinds of
+                metrics they are interested in polling.
+
+        Usage
+        -----
+        >>> ks = km.kernel_metrics("memory_usage", "latency")
+        The function call above will return a dict with two key-value
+        pairs (assuming there are function to gather those stats).
+        The keys are exactly the same as the inputs.
+        >>> print ks["memory_usage"]
+        Line above will print the memory usage of the kernel.
+
+        """
+        results = {}
+        for metric_name in desired_metrics:
+            try:
+                cls = KernelMetricStore.TYPES[metric_name]
+            except KeyError:
+                continue
+            else:
+                metric = cls(self)
+                results[metric_name] = metric.poll()
+        return results
 
 KernelManagerABC.register(KernelManager)
 

jupyter_client/tests/test_kernelmanager.py

@@ -131,6 +131,18 @@ def test_start_new_kernel(self):
         self.assertTrue(km.is_alive())
         self.assertTrue(kc.is_alive())
 
+    def test_kernel_stats(self):
+        c = Config()
+        km = KernelManager(config=c)
+        km.start_kernel(stdout=PIPE, stderr=PIPE)
+        # Call the kernel_stats function requesting the
+        # memory usage of the kernel.
+        ks = km.kernel_metrics("memory_usage")
+        km.shutdown_kernel(now=True)
+        # An arbitrary check for a value of a certain key in the result
+        # of `kernel_stats`, in this case its the memory usage one.
+        assert type(ks["memory_usage"]) is int
+
 @pytest.mark.parallel
 class TestParallel: