Merge pull request #171 from jtpio/enable-pss

Use `pss` if available

Author: krassowski

README.md

@@ -105,8 +105,9 @@ causes "lag"/pauses in the UI. To workaround this you can disable Prometheus met
 
 ## Resources Displayed
 
-Currently the server extension only reports memory usage (just RSS) and CPU usage. Other metrics will be
-added in the future as needed.
+Currently the server extension only reports memory usage and CPU usage. Other metrics will be added in the future as needed.
+
+Memory usage will show the PSS whenever possible (Linux only feature), and default to RSS otherwise.
 
 The notebook extension currently doesn't show CPU usage, only memory usage.
 

jupyter_resource_usage/api.py

@@ -42,24 +42,30 @@ async def get(self):
 
         # Get memory information
         rss = 0
+        pss = None
         for p in all_processes:
             try:
-                rss += p.memory_info().rss
+                info = p.memory_full_info()
+                if hasattr(info, "pss"):
+                    pss = (pss or 0) + info.pss
+                rss += info.rss
             except (psutil.NoSuchProcess, psutil.AccessDenied) as e:
                 pass
 
         if callable(config.mem_limit):
-            mem_limit = config.mem_limit(rss=rss)
+            mem_limit = config.mem_limit(rss=rss, pss=pss)
         else:  # mem_limit is an Int
             mem_limit = config.mem_limit
 
-        limits = {"memory": {"rss": mem_limit}}
+        limits = {"memory": {"rss": mem_limit, "pss": mem_limit}}
         if config.mem_limit and config.mem_warning_threshold != 0:
             limits["memory"]["warn"] = (mem_limit - rss) < (
                 mem_limit * config.mem_warning_threshold
             )
 
         metrics = {"rss": rss, "limits": limits}
+        if pss is not None:
+            metrics["pss"] = pss
 
         # Optionally get CPU information
         if config.track_cpu_percent:

jupyter_resource_usage/static/main.js

@@ -53,14 +53,16 @@ define([
         $.getJSON({
             url: utils.get_body_data('baseUrl') + 'api/metrics/v1',
             success: function (data) {
-                totalMemoryUsage = humanFileSize(data['rss']);
+                value = data['pss'] || data['rss'];
+                totalMemoryUsage = humanFileSize(value);
 
                 var limits = data['limits'];
                 var display = totalMemoryUsage;
 
                 if (limits['memory']) {
-                    if (limits['memory']['rss']) {
-                        maxMemoryUsage = humanFileSize(limits['memory']['rss']);
+                    limit = limits['memory']['pss'] ?? limits['memory']['rss'];
+                    if (limit) {
+                        maxMemoryUsage = humanFileSize(limit);
                         display += " / " + maxMemoryUsage
                     }
                     if (limits['memory']['warn']) {

packages/labextension/src/memoryUsage.tsx

@@ -180,10 +180,9 @@ export namespace MemoryUsage {
         this._units = 'B';
         this._warn = false;
       } else {
-        const numBytes = value.rss;
-        const memoryLimit = value.limits.memory
-          ? value.limits.memory.rss
-          : null;
+        const numBytes = value.pss ?? value.rss;
+        const memoryLimits = value.limits.memory;
+        const memoryLimit = memoryLimits?.pss ?? memoryLimits?.rss ?? null;
         const [currentMemory, units] = convertToLargestUnit(numBytes);
         const usageWarning = value.limits.memory
           ? value.limits.memory.warn
@@ -260,9 +259,11 @@ namespace Private {
    */
   export interface IMetricRequestResult {
     rss: number;
+    pss?: number;
     limits: {
       memory?: {
         rss: number;
+        pss?: number;
         warn: boolean;
       };
     };