enhancement: Track disk utilization in $HOME

Similar to memory usage tracking, polls the disk used/available
in $HOME.

Author: holzman

jupyter_resource_usage/api.py

@@ -1,4 +1,5 @@
 import json
+import os
 from concurrent.futures import ThreadPoolExecutor
 from inspect import isawaitable
 
@@ -58,9 +59,15 @@ async def get(self):
             )
 
         metrics = {"rss": rss, "limits": limits}
+
         if pss is not None:
             metrics["pss"] = pss
 
+        if config.track_disk_usage:
+            disk_info = psutil.disk_usage(os.getenv("HOME"))
+            metrics["disk_used"] = disk_info.used
+            metrics["disk_total"] = disk_info.total
+
         # Optionally get CPU information
         if config.track_cpu_percent:
             cpu_count = psutil.cpu_count()

jupyter_resource_usage/config.py

@@ -37,6 +37,11 @@ class ResourceUseDisplay(Configurable):
     Holds server-side configuration for jupyter-resource-usage
     """
 
+    system_disk_metrics = List(
+        trait=PSUtilMetric(),
+        default_value=[{"name": "disk_usage", "attribute": "disk_used"}],
+    )
+
     process_memory_metrics = List(
         trait=PSUtilMetric(),
         default_value=[{"name": "memory_info", "attribute": "rss"}],
@@ -129,3 +134,10 @@ def _cpu_limit_default(self):
         Set to False in order to disable reporting of Prometheus style metrics.
         """,
     ).tag(config=True)
+
+    track_disk_usage = Bool(
+        default_value=False,
+        help="""
+        Set to True in order to enable reporting of disk usage statistics.
+        """,
+    ).tag(config=True)

jupyter_resource_usage/metrics.py

@@ -80,3 +80,6 @@ def cpu_metrics(self):
         return self.metrics(
             self.config.process_cpu_metrics, self.config.system_cpu_metrics
         )
+
+    def disk_metrics(self):
+        return self.metrics(self.config.system_disk_metrics)

packages/labextension/src/memoryView.tsx

@@ -20,11 +20,11 @@ const MemoryViewComponent = ({
   const [values, setValues] = useState<number[]>([]);
 
   const update = (): void => {
-    const { memoryLimit, currentMemory, units } = model;
-    const precision = ['B', 'KB', 'MB'].indexOf(units) > 0 ? 0 : 2;
+    const { memoryLimit, currentMemory, memoryUnits } = model;
+    const precision = ['B', 'KB', 'MB'].indexOf(memoryUnits) > 0 ? 0 : 2;
     const newText = `${currentMemory.toFixed(precision)} ${
       memoryLimit ? '/ ' + memoryLimit.toFixed(precision) : ''
-    } ${units}`;
+    } ${memoryUnits}`;
     const newValues = model.values.map((value) => value.memoryPercent);
     setText(newText);
     setValues(newValues);

packages/labextension/src/model.ts

@@ -1,5 +1,5 @@
-// Copyright (c) Jupyter Development Team.
-// Distributed under the terms of the Modified BSD License.
+// Copyright (c) Jupyter Development Team.  Distributed under the
+// terms of the Modified BSD License.
 
 import { VDomModel } from '@jupyterlab/apputils';
 
@@ -52,14 +52,22 @@ export namespace ResourceUsage {
         if (phase === 'rejected') {
           const oldMemoryAvailable = this._memoryAvailable;
           const oldCpuAvailable = this._cpuAvailable;
+          const oldDiskAvailable = this._diskAvailable;
+
           this._memoryAvailable = false;
           this._cpuAvailable = false;
+          this._diskAvailable = false;
+
           this._currentMemory = 0;
           this._memoryLimit = null;
           this._cpuLimit = null;
-          this._units = 'B';
+          this._memoryUnits = 'B';
 
-          if (oldMemoryAvailable || oldCpuAvailable) {
+          this._diskUnits = 'B';
+          this._diskUsed = 0;
+          this._diskTotal = null;
+
+          if (oldMemoryAvailable || oldCpuAvailable || oldDiskAvailable) {
             this.stateChanged.emit();
           }
           return;
@@ -79,7 +87,7 @@ export namespace ResourceUsage {
      * Whether the metrics server extension is available.
      */
     get metricsAvailable(): boolean {
-      return this._memoryAvailable || this._cpuAvailable;
+      return this._memoryAvailable || this._cpuAvailable || this._diskAvailable;
     }
 
     /**
@@ -120,8 +128,8 @@ export namespace ResourceUsage {
     /**
      * The units for memory usages and limits.
      */
-    get units(): MemoryUnit {
-      return this._units;
+    get memoryUnits(): MemoryUnit {
+      return this._memoryUnits;
     }
 
     /**
@@ -131,6 +139,22 @@ export namespace ResourceUsage {
       return this._currentCpuPercent;
     }
 
+    get diskAvailable(): boolean {
+      return this._diskAvailable;
+    }
+
+    get diskUsed(): number {
+      return this._diskUsed;
+    }
+
+    get diskTotal(): number | null {
+      return this._diskTotal;
+    }
+
+    get diskUnits(): MemoryUnit {
+      return this._diskUnits;
+    }
+
     /**
      * Get a list of the last metric values.
      */
@@ -164,26 +188,30 @@ export namespace ResourceUsage {
       if (value === null) {
         this._memoryAvailable = false;
         this._cpuAvailable = false;
+        this._diskAvailable = false;
         this._currentMemory = 0;
         this._memoryLimit = null;
-        this._units = 'B';
+        this._memoryUnits = 'B';
+        this._diskUnits = 'B';
+        this._diskUsed = 0;
+        this._diskTotal = null;
         this._warn = false;
         return;
       }
 
       const numBytes = value.pss ?? value.rss;
       const memoryLimits = value.limits.memory;
       const memoryLimit = memoryLimits?.pss ?? memoryLimits?.rss ?? null;
-      const [currentMemory, units] = convertToLargestUnit(numBytes);
+      const [currentMemory, memoryUnits] = convertToLargestUnit(numBytes);
       const usageWarning = value.limits.memory
         ? value.limits.memory.warn
         : false;
 
       this._memoryAvailable = numBytes !== undefined;
       this._currentMemory = currentMemory;
-      this._units = units;
+      this._memoryUnits = memoryUnits;
       this._memoryLimit = memoryLimit
-        ? memoryLimit / MEMORY_UNIT_LIMITS[units]
+        ? memoryLimit / MEMORY_UNIT_LIMITS[memoryUnits]
         : null;
       const memoryPercent = this.memoryLimit
         ? Math.min(this._currentMemory / this.memoryLimit, 1)
@@ -195,19 +223,34 @@ export namespace ResourceUsage {
       this._currentCpuPercent =
         value.cpu_percent !== undefined ? value.cpu_percent / 100 : 0;
 
+      const diskUsedBytes = value.disk_used;
+      const diskTotal = value.disk_total ?? null;
+      const [diskUsed, diskUnits] = convertToLargestUnit(diskUsedBytes);
+
+      this._diskAvailable = diskTotal ?? false;
+      this._diskUnits = diskUnits;
+      this._diskUsed = diskUsed;
+      this._diskTotal = diskTotal
+        ? diskTotal / MEMORY_UNIT_LIMITS[diskUnits]
+        : null;
+
       this._values.push({ memoryPercent, cpuPercent: this._currentCpuPercent });
       this._values.shift();
       this.stateChanged.emit(void 0);
     }
 
     private _memoryAvailable = false;
     private _cpuAvailable = false;
+    private _diskAvailable = false;
     private _currentMemory = 0;
     private _currentCpuPercent = 0;
     private _memoryLimit: number | null = null;
     private _cpuLimit: number | null = null;
+    private _diskUsed = 0;
+    private _diskTotal: number | null = null;
     private _poll: Poll<Private.IMetricRequestResult | null>;
-    private _units: MemoryUnit = 'B';
+    private _memoryUnits: MemoryUnit = 'B';
+    private _diskUnits: MemoryUnit = 'B';
     private _warn = false;
     private _values: Model.IMetricValue[] = [];
   }
@@ -268,6 +311,8 @@ namespace Private {
     pss?: number;
     cpu_percent?: number;
     cpu_count?: number;
+    disk_used?: number;
+    disk_total?: number;
     limits: {
       memory?: {
         rss: number;

packages/labextension/src/resourceUsage.tsx

@@ -37,16 +37,25 @@ export class ResourceUsageStatus extends VDomRenderer<ResourceUsage.Model> {
       text = this._trans.__(
         'Mem: %1 %2',
         this.model.currentMemory.toFixed(Private.DECIMAL_PLACES),
-        this.model.units
+        this.model.memoryUnits
       );
     } else {
       text = this._trans.__(
         'Mem: %1 / %2 %3',
         this.model.currentMemory.toFixed(Private.DECIMAL_PLACES),
         this.model.memoryLimit.toFixed(Private.DECIMAL_PLACES),
-        this.model.units
+        this.model.memoryUnits
       );
     }
+
+    if (this.model.diskAvailable && this.model.diskTotal != null) {
+      text = `${text} ${this._trans.__(
+        'Disk: %1 / %2 %3',
+        this.model.diskUsed.toFixed(Private.DECIMAL_PLACES),
+        this.model.diskTotal.toFixed(Private.DECIMAL_PLACES),
+        this.model.diskUnits
+      )}`;
+    }
     if (this.model.cpuAvailable) {
       text = `CPU: ${(this.model.currentCpuPercent * 100).toFixed(
         Private.DECIMAL_PLACES