[Ready for Review] add websockets instead of rest api for better resp…

…onse times (#188)

* add websockets instead of rest api

* fix pytests

* rm unused import

* update dependencies

* allow-root for jupyterlab server within the ci container

* fix testing in ci

* remove logs

* add SETTINGS text to config button

ci/test_python.sh

@@ -18,7 +18,7 @@ set +u
 conda activate test
 set -u
 
-# rapids-logger "Downloading artifacts from previous jobs"
+rapids-logger "Downloading artifacts from previous jobs"
 PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
 
 rapids-print-env

conda/environments/all_arch-any.yaml

@@ -12,6 +12,8 @@ dependencies:
 - psutil
 - pynvml
 - pytest
+- pytest-asyncio
 - pytest-jupyter[server]>=0.6.0
 - python>=3.8
+- websockets
 name: all_arch-any

dependencies.yaml

@@ -80,3 +80,5 @@ dependencies:
         packages:
           - pytest
           - pytest-jupyter[server]>=0.6.0
+          - pytest-asyncio
+          - websockets

jupyterlab_nvdashboard/apps/cpu.py

@@ -1,13 +1,11 @@
 import json
 import psutil
 import time
-import tornado
-from jupyter_server.base.handlers import APIHandler
+from jupyterlab_nvdashboard.apps.utils import CustomWebSocketHandler
 
 
-class CPUResourceHandler(APIHandler):
-    @tornado.web.authenticated
-    def get(self):
+class CPUResourceWebSocketHandler(CustomWebSocketHandler):
+    def send_data(self):
         now = time.time()
         stats = {
             "time": now * 1000,
@@ -18,5 +16,4 @@ def get(self):
             "network_read": psutil.net_io_counters().bytes_recv,
             "network_write": psutil.net_io_counters().bytes_sent,
         }
-        self.set_header("Content-Type", "application/json")
-        self.write(json.dumps(stats))
+        self.write_message(json.dumps(stats))

jupyterlab_nvdashboard/apps/gpu.py

@@ -1,8 +1,7 @@
 import json
+from jupyterlab_nvdashboard.apps.utils import CustomWebSocketHandler
 import pynvml
 import time
-import tornado
-from jupyter_server.base.handlers import APIHandler
 
 try:
     pynvml.nvmlInit()
@@ -41,19 +40,17 @@
         pci_gen = None
 
 
-class GPUUtilizationHandler(APIHandler):
-    @tornado.web.authenticated
-    def get(self):
+class GPUUtilizationWebSocketHandler(CustomWebSocketHandler):
+    def send_data(self):
         gpu_utilization = [
             pynvml.nvmlDeviceGetUtilizationRates(gpu_handles[i]).gpu
             for i in range(ngpus)
         ]
-        self.finish(json.dumps({"gpu_utilization": gpu_utilization}))
+        self.write_message(json.dumps({"gpu_utilization": gpu_utilization}))
 
 
-class GPUUsageHandler(APIHandler):
-    @tornado.web.authenticated
-    def get(self):
+class GPUUsageWebSocketHandler(CustomWebSocketHandler):
+    def send_data(self):
         memory_usage = [
             pynvml.nvmlDeviceGetMemoryInfo(handle).used
             for handle in gpu_handles
@@ -64,16 +61,15 @@ def get(self):
             for handle in gpu_handles
         ]
 
-        self.finish(
+        self.write_message(
             json.dumps(
                 {"memory_usage": memory_usage, "total_memory": total_memory}
             )
         )
 
 
-class GPUResourceHandler(APIHandler):
-    @tornado.web.authenticated
-    def get(self):
+class GPUResourceWebSocketHandler(CustomWebSocketHandler):
+    def send_data(self):
         now = time.time()
         stats = {
             "time": now * 1000,
@@ -118,15 +114,13 @@ def get(self):
         stats["gpu_memory_total"] = round(
             (stats["gpu_memory_total"] / gpu_mem_sum) * 100, 2
         )
-        self.set_header("Content-Type", "application/json")
-        self.write(json.dumps(stats))
+        self.write_message(json.dumps(stats))
 
 
-class NVLinkThroughputHandler(APIHandler):
+class NVLinkThroughputWebSocketHandler(CustomWebSocketHandler):
     prev_throughput = None
 
-    @tornado.web.authenticated
-    def get(self):
+    def send_data(self):
         throughput = [
             pynvml.nvmlDeviceGetFieldValues(
                 handle,
@@ -162,9 +156,8 @@ def get(self):
         # Store the current throughput for the next request
         self.prev_throughput = throughput
 
-        self.set_header("Content-Type", "application/json")
         # Send the change in throughput as part of the response
-        self.write(
+        self.write_message(
             json.dumps(
                 {
                     "nvlink_rx": [
@@ -191,9 +184,8 @@ def get(self):
         )
 
 
-class PCIStatsHandler(APIHandler):
-    @tornado.web.authenticated
-    def get(self):
+class PCIStatsWebSocketHandler(CustomWebSocketHandler):
+    def send_data(self):
         # Use device-0 to get "upper bound"
         pci_width = pynvml.nvmlDeviceGetMaxPcieLinkWidth(gpu_handles[0])
         pci_bw = {
@@ -231,5 +223,4 @@ def get(self):
             "max_rxtx_tp": max_rxtx_tp,
         }
 
-        self.set_header("Content-Type", "application/json")
-        self.write(json.dumps(stats))
+        self.write_message(json.dumps(stats))

jupyterlab_nvdashboard/apps/utils.py

@@ -0,0 +1,31 @@
+from tornado.websocket import WebSocketHandler
+import tornado
+import json
+
+
+class CustomWebSocketHandler(WebSocketHandler):
+    def open(self):
+        self.write_message(json.dumps({"status": "connected"}))
+        self.set_nodelay(True)
+        # Start a periodic callback to send data every 50ms
+        self.callback = tornado.ioloop.PeriodicCallback(self.send_data, 1000)
+        self.callback.start()
+
+    def on_message(self, message):
+        message_data = json.loads(message)
+        # Update the periodic callback frequency
+        new_frequency = message_data["updateFrequency"]
+        if hasattr(self, "callback"):
+            self.callback.stop()
+            self.callback = tornado.ioloop.PeriodicCallback(
+                self.send_data, new_frequency
+            )
+            if not message_data["isPaused"]:
+                self.callback.start()
+
+    def on_close(self):
+        if hasattr(self, "callback") and self.callback.is_running():
+            self.callback.stop()
+
+    def send_data(self):
+        pass

jupyterlab_nvdashboard/handlers.py

@@ -26,13 +26,13 @@ def setup_handlers(web_app):
             base_url, URL_PATH, "nvlink_throughput"
         )
         handlers += [
-            (route_pattern_gpu_util, apps.gpu.GPUUtilizationHandler),
-            (route_pattern_gpu_usage, apps.gpu.GPUUsageHandler),
-            (route_pattern_gpu_resource, apps.gpu.GPUResourceHandler),
-            (route_pattern_pci_stats, apps.gpu.PCIStatsHandler),
+            (route_pattern_gpu_util, apps.gpu.GPUUtilizationWebSocketHandler),
+            (route_pattern_gpu_usage, apps.gpu.GPUUsageWebSocketHandler),
+            (route_pattern_gpu_resource, apps.gpu.GPUResourceWebSocketHandler),
+            (route_pattern_pci_stats, apps.gpu.PCIStatsWebSocketHandler),
             (
                 route_pattern_nvlink_throughput,
-                apps.gpu.NVLinkThroughputHandler,
+                apps.gpu.NVLinkThroughputWebSocketHandler,
             ),
         ]
 
@@ -41,7 +41,7 @@ def setup_handlers(web_app):
     )
 
     handlers += [
-        (route_pattern_cpu_resource, apps.cpu.CPUResourceHandler),
+        (route_pattern_cpu_resource, apps.cpu.CPUResourceWebSocketHandler),
     ]
 
     web_app.add_handlers(host_pattern, handlers)

jupyterlab_nvdashboard/tests/conftest.py

@@ -0,0 +1,2 @@
+def pytest_configure(config):
+    config.addinivalue_line("markers", "asyncio: mark test as asyncio")

jupyterlab_nvdashboard/tests/test_cpu_handlers.py

@@ -0,0 +1,37 @@
+import json
+import pytest
+from unittest.mock import MagicMock, patch
+
+from jupyterlab_nvdashboard.apps.cpu import CPUResourceWebSocketHandler
+
+
+@pytest.fixture
+def mock_handler(monkeypatch):
+    mock = MagicMock()
+    monkeypatch.setattr(
+        "jupyterlab_nvdashboard.apps.cpu.CustomWebSocketHandler.write_message",
+        mock,
+    )
+    return mock
+
+
+@pytest.fixture
+def handler_args():
+    with patch("tornado.web.Application") as mock_application, patch(
+        "tornado.httputil.HTTPServerRequest"
+    ) as mock_request:
+        yield mock_application, mock_request
+
+
+def test_cpu_resource_handler(mock_handler, handler_args):
+    handler = CPUResourceWebSocketHandler(*handler_args)
+    handler.send_data()
+    args, _ = mock_handler.call_args
+    data = json.loads(args[0])
+    assert "time" in data
+    assert "cpu_utilization" in data
+    assert "memory_usage" in data
+    assert "disk_read" in data
+    assert "disk_write" in data
+    assert "network_read" in data
+    assert "network_write" in data

jupyterlab_nvdashboard/tests/test_gpu_handlers.py

@@ -0,0 +1,80 @@
+import json
+import pytest
+from unittest.mock import MagicMock, patch
+
+from jupyterlab_nvdashboard.apps.gpu import (
+    GPUUtilizationWebSocketHandler,
+    GPUUsageWebSocketHandler,
+    GPUResourceWebSocketHandler,
+    NVLinkThroughputWebSocketHandler,
+    PCIStatsWebSocketHandler,
+)
+
+
+@pytest.fixture
+def mock_handler(monkeypatch):
+    mock = MagicMock()
+    monkeypatch.setattr(
+        "jupyterlab_nvdashboard.apps.gpu.CustomWebSocketHandler.write_message",
+        mock,
+    )
+    return mock
+
+
+@pytest.fixture
+def handler_args():
+    with patch("tornado.web.Application") as mock_application, patch(
+        "tornado.httputil.HTTPServerRequest"
+    ) as mock_request:
+        yield mock_application, mock_request
+
+
+def test_gpu_utilization_handler(mock_handler, handler_args):
+    handler = GPUUtilizationWebSocketHandler(*handler_args)
+    handler.send_data()
+    args, _ = mock_handler.call_args
+    data = json.loads(args[0])
+    assert "gpu_utilization" in data
+
+
+def test_gpu_usage_handler(mock_handler, handler_args):
+    handler = GPUUsageWebSocketHandler(*handler_args)
+    handler.send_data()
+    args, _ = mock_handler.call_args
+    data = json.loads(args[0])
+    assert "memory_usage" in data
+    assert "total_memory" in data
+
+
+def test_gpu_resource_handler(mock_handler, handler_args):
+    handler = GPUResourceWebSocketHandler(*handler_args)
+    handler.send_data()
+    args, _ = mock_handler.call_args
+    data = json.loads(args[0])
+    assert "time" in data
+    assert "gpu_utilization_total" in data
+    assert "gpu_memory_total" in data
+    assert "rx_total" in data
+    assert "tx_total" in data
+    assert "gpu_memory_individual" in data
+    assert "gpu_utilization_individual" in data
+
+
+def test_nvlink_throughput_handler(mock_handler, handler_args):
+    handler = NVLinkThroughputWebSocketHandler(*handler_args)
+    handler.send_data()
+    args, _ = mock_handler.call_args
+    data = json.loads(args[0])
+    assert "nvlink_rx" in data
+    assert "nvlink_tx" in data
+    assert "max_rxtx_bw" in data
+
+
+def test_pci_stats_handler(mock_handler, handler_args):
+    handler = PCIStatsWebSocketHandler(*handler_args)
+    handler.send_data()
+    args, _ = mock_handler.call_args
+    data = json.loads(args[0])
+    assert "pci_tx" in data
+    assert "pci_rx" in data
+    assert "max_rxtx_tp" in data