Skip to content

Commit

Permalink
Add a websocket endpoint and a UI to get live usage data
Browse files Browse the repository at this point in the history
  • Loading branch information
dormant-user committed Sep 4, 2024
1 parent b8b334d commit 1bcc935
Show file tree
Hide file tree
Showing 9 changed files with 295 additions and 2 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -17,3 +17,4 @@ logging.ini
*.db

!samples/*
discard/
6 changes: 5 additions & 1 deletion docs/genindex.html
Original file line number Diff line number Diff line change
Expand Up @@ -429,6 +429,8 @@ <h2 id="S">S</h2>
<li><a href="index.html#pyninja.models.ServiceStatus.status_code">status_code (pyninja.models.ServiceStatus attribute)</a>
</li>
<li><a href="index.html#pyninja.service.stopped">stopped() (in module pyninja.service)</a>
</li>
<li><a href="index.html#pyninja.squire.system_resources">system_resources() (in module pyninja.squire)</a>
</li>
</ul></td>
</tr></table>
Expand Down Expand Up @@ -458,10 +460,12 @@ <h2 id="U">U</h2>
<h2 id="W">W</h2>
<table style="width: 100%" class="indextable genindextable"><tr>
<td style="width: 33%; vertical-align: top;"><ul>
<li><a href="index.html#pyninja.models.ServiceManager.windows">windows (pyninja.models.ServiceManager attribute)</a>
<li><a href="index.html#pyninja.routers.websocket_endpoint">websocket_endpoint() (in module pyninja.routers)</a>
</li>
</ul></td>
<td style="width: 33%; vertical-align: top;"><ul>
<li><a href="index.html#pyninja.models.ServiceManager.windows">windows (pyninja.models.ServiceManager attribute)</a>
</li>
<li><a href="index.html#pyninja.models.EnvConfig.workers">workers (pyninja.models.EnvConfig attribute)</a>
</li>
</ul></td>
Expand Down
25 changes: 25 additions & 0 deletions docs/index.html
Original file line number Diff line number Diff line change
Expand Up @@ -372,6 +372,17 @@ <h1>Welcome to PyNinja’s documentation!<a class="headerlink" href="#welcome-to
</dl>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="pyninja.routers.websocket_endpoint">
<em class="property"><span class="k"><span class="pre">async</span></span><span class="w"> </span></em><span class="sig-prename descclassname"><span class="pre">pyninja.routers.</span></span><span class="sig-name descname"><span class="pre">websocket_endpoint</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">websocket</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">WebSocket</span></span></em><span class="sig-paren">)</span><a class="headerlink" href="#pyninja.routers.websocket_endpoint" title="Permalink to this definition"></a></dt>
<dd><p>Websocket endpoint to fetch live system resource usage.</p>
<dl class="field-list simple">
<dt class="field-odd">Parameters<span class="colon">:</span></dt>
<dd class="field-odd"><p><strong>websocket</strong> – Reference to the websocket object.</p>
</dd>
</dl>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="pyninja.routers.get_all_routes">
<span class="sig-prename descclassname"><span class="pre">pyninja.routers.</span></span><span class="sig-name descname"><span class="pre">get_all_routes</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">List</span><span class="p"><span class="pre">[</span></span><span class="pre">APIRoute</span><span class="p"><span class="pre">]</span></span></span></span><a class="headerlink" href="#pyninja.routers.get_all_routes" title="Permalink to this definition"></a></dt>
Expand Down Expand Up @@ -1025,6 +1036,20 @@ <h1>Models<a class="headerlink" href="#models" title="Permalink to this heading"
</dl>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="pyninja.squire.system_resources">
<span class="sig-prename descclassname"><span class="pre">pyninja.squire.</span></span><span class="sig-name descname"><span class="pre">system_resources</span></span><span class="sig-paren">(</span><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">Dict</span><span class="p"><span class="pre">[</span></span><span class="pre">str</span><span class="p"><span class="pre">,</span></span><span class="w"> </span><span class="pre">dict</span><span class="p"><span class="pre">]</span></span></span></span><a class="headerlink" href="#pyninja.squire.system_resources" title="Permalink to this definition"></a></dt>
<dd><p>Get system resources like CPU, virtual memory and swap memory information.</p>
<dl class="field-list simple">
<dt class="field-odd">Returns<span class="colon">:</span></dt>
<dd class="field-odd"><p>Returns a nested dictionary.</p>
</dd>
<dt class="field-even">Return type<span class="colon">:</span></dt>
<dd class="field-even"><p>Dict[str, dict]</p>
</dd>
</dl>
</dd></dl>

<dl class="py function">
<dt class="sig sig-object py" id="pyninja.squire.format_nos">
<span class="sig-prename descclassname"><span class="pre">pyninja.squire.</span></span><span class="sig-name descname"><span class="pre">format_nos</span></span><span class="sig-paren">(</span><em class="sig-param"><span class="n"><span class="pre">input_</span></span><span class="p"><span class="pre">:</span></span><span class="w"> </span><span class="n"><span class="pre">float</span></span></em><span class="sig-paren">)</span> <span class="sig-return"><span class="sig-return-icon">&#x2192;</span> <span class="sig-return-typehint"><span class="pre">int</span><span class="w"> </span><span class="p"><span class="pre">|</span></span><span class="w"> </span><span class="pre">float</span></span></span><a class="headerlink" href="#pyninja.squire.format_nos" title="Permalink to this definition"></a></dt>
Expand Down
Binary file modified docs/objects.inv
Binary file not shown.
2 changes: 1 addition & 1 deletion docs/searchindex.js

Large diffs are not rendered by default.

223 changes: 223 additions & 0 deletions index.html
Original file line number Diff line number Diff line change
@@ -0,0 +1,223 @@
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="UTF-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>System Monitor</title>
<script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
<style id="main-css" disabled>
body {
font-family: Arial, sans-serif;
}
.container {
display: flex;
justify-content: space-between;
margin-top: 50px;
}
.box {
border: 1px solid #ccc;
padding: 20px;
width: 30%;
text-align: center;
}
.progress {
width: 100%;
background-color: #f3f3f3;
border-radius: 5px;
overflow: hidden;
transition: background-color 0.5s ease;
}
.progress-bar {
height: 25px;
transition: width 0.5s ease, background-color 0.5s ease;
width: 0%;
}
.progress-bar-green {
background-color: #4caf50;
}
.progress-bar-yellow {
background-color: #ffeb3b;
}
.progress-bar-red {
background-color: #f44336;
}
.chart-container {
position: relative;
height: 200px;
width: 80%;
margin: 0 auto;
max-width: 100%;
}
canvas {
width: 100% !important;
height: inherit !important;
max-height: 100% !important;
}
</style>
</head>
<body>
<h1>System Monitor</h1>
<div class="container">
<div class="box">
<h3>CPU Usage</h3>
<div class="cpu-box" id="cpuUsageContainer">
<!-- CPU Usage will be dynamically added here -->
</div>
</div>
<div class="box">
<h3>Memory and Swap Usage</h3>
<div class="progress">
<div id="memoryUsage" class="progress-bar"></div>
</div>
<p id="memoryUsageText">Memory: 0%</p>
<div class="progress">
<div id="swapUsage" class="progress-bar"></div>
</div>
<p id="swapUsageText">Swap: 0%</p>
</div>
<div class="box">
<h4>Memory Usage</h4>
<div class="chart-container">
<canvas id="memoryChart"></canvas>
</div>
<h4>Swap Usage</h4>
<div class="chart-container">
<canvas id="swapChart"></canvas>
</div>
</div>
</div>

<script>
window.addEventListener('load', () => {
document.getElementById('main-css').disabled = false;
});

const ws = new WebSocket("ws://localhost:8080/ws/system");

let memoryChartInstance = null;
let swapChartInstance = null;

ws.onmessage = function(event) {
const data = JSON.parse(event.data);

// Update CPU usage
const cpuUsage = data.cpu_usage;
const cpuContainer = document.getElementById('cpuUsageContainer');
cpuContainer.innerHTML = ''; // Clear previous content
cpuUsage.forEach((usage, index) => {
const cpuDiv = document.createElement('div');
cpuDiv.innerHTML = `
<strong>CPU ${index}:</strong> ${usage}%
<div class="progress">
<div id="cpu${index}" class="progress-bar"></div>
</div>
`;
cpuContainer.appendChild(cpuDiv);
updateProgressBar(`cpu${index}`, usage);
});

// Update Memory and Swap usage
const memoryInfo = data.memory_info;
const memoryUsage = (memoryInfo.used / memoryInfo.total) * 100;
document.getElementById('memoryUsage').style.width = memoryUsage.toFixed(2) + '%';
document.getElementById('memoryUsageText').innerText = `Memory: ${memoryUsage.toFixed(2)}%`;
updateProgressBar('memoryUsage', memoryUsage);

const swapInfo = data.swap_info;
const swapUsage = (swapInfo.used / swapInfo.total) * 100;
document.getElementById('swapUsage').style.width = swapUsage.toFixed(2) + '%';
document.getElementById('swapUsageText').innerText = `Swap: ${swapUsage.toFixed(2)}%`;
updateProgressBar('swapUsage', swapUsage);

// Create/update pie charts
if (memoryChartInstance) {
memoryChartInstance.data.datasets[0].data = [memoryInfo.used, memoryInfo.total - memoryInfo.used];
memoryChartInstance.update();
} else {
const memoryChart = document.getElementById('memoryChart').getContext('2d');
memoryChartInstance = new Chart(memoryChart, {
type: 'pie',
data: {
labels: ['Used', 'Free'],
datasets: [{
label: 'Memory Usage',
data: [memoryInfo.used, memoryInfo.total - memoryInfo.used],
backgroundColor: ['#FF6384', '#36A2EB']
}]
},
options: {
responsive: true,
plugins: {
tooltip: {
callbacks: {
label: function(tooltipItem) {
const value = tooltipItem.raw;
const formattedValue = formatBytes(value);
return `${tooltipItem.label}: ${formattedValue}`;
}
}
}
}
}
});
}

if (swapChartInstance) {
swapChartInstance.data.datasets[0].data = [swapInfo.used, swapInfo.total - swapInfo.used];
swapChartInstance.update();
} else {
const swapChart = document.getElementById('swapChart').getContext('2d');
swapChartInstance = new Chart(swapChart, {
type: 'pie',
data: {
labels: ['Used', 'Free'],
datasets: [{
label: 'Swap Usage',
data: [swapInfo.used, swapInfo.total - swapInfo.used],
backgroundColor: ['#FFCE56', '#E7E9ED']
}]
},
options: {
responsive: true,
plugins: {
tooltip: {
callbacks: {
label: function(tooltipItem) {
const value = tooltipItem.raw;
const formattedValue = formatBytes(value);
return `${tooltipItem.label}: ${formattedValue}`;
}
}
}
}
}
});
}
};

function updateProgressBar(id, percentage) {
const bar = document.getElementById(id);
bar.style.width = percentage + '%';

// Remove old color classes
bar.classList.remove('progress-bar-green', 'progress-bar-yellow', 'progress-bar-red');

// Add new color class based on percentage
if (percentage <= 40) {
bar.classList.add('progress-bar-green');
} else if (percentage <= 80) {
bar.classList.add('progress-bar-yellow');
} else {
bar.classList.add('progress-bar-red');
}
}

function formatBytes(bytes) {
if (bytes < 1024) return bytes + ' bytes';
else if (bytes < 1048576) return (bytes / 1024).toFixed(2) + ' KB';
else if (bytes < 1073741824) return (bytes / 1048576).toFixed(2) + ' MB';
else return (bytes / 1073741824).toFixed(2) + ' GB';
}
</script>
</body>
</html>
1 change: 1 addition & 0 deletions pyninja/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -37,6 +37,7 @@ def start(**kwargs) -> None:
description="Lightweight OS-agnostic service monitoring API",
version=pyninja.version,
)
app.add_websocket_route(path="/ws/system", route=routers.websocket_endpoint)
kwargs = dict(
host=models.env.ninja_host,
port=models.env.ninja_port,
Expand Down
24 changes: 24 additions & 0 deletions pyninja/routers.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,7 @@
import asyncio
import logging
import subprocess
import time
from http import HTTPStatus
from typing import List, Optional

Expand All @@ -8,6 +10,7 @@
from fastapi.responses import RedirectResponse
from fastapi.routing import APIRoute
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
from fastapi.websockets import WebSocket, WebSocketDisconnect
from pydantic import PositiveFloat, PositiveInt

from pyninja import (
Expand Down Expand Up @@ -329,6 +332,27 @@ async def health():
raise exceptions.APIResponse(status_code=HTTPStatus.OK, detail=HTTPStatus.OK.phrase)


async def websocket_endpoint(websocket: WebSocket):
"""Websocket endpoint to fetch live system resource usage.
Args:
websocket: Reference to the websocket object.
"""
await websocket.accept()
refresh_time = time.time()
data = squire.system_resources()
while True:
if time.time() - refresh_time > 3:
refresh_time = time.time()
LOGGER.debug("Fetching new charts")
data = squire.system_resources()
try:
await websocket.send_json(data)
await asyncio.sleep(1)
except WebSocketDisconnect:
break


def get_all_routes() -> List[APIRoute]:
"""Get all the routes to be added for the API server.
Expand Down
15 changes: 15 additions & 0 deletions pyninja/squire.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
import subprocess
from typing import Dict, List

import psutil
import requests
import yaml
from pydantic import PositiveFloat, PositiveInt
Expand Down Expand Up @@ -67,6 +68,20 @@ def private_ip_address() -> str | None:
return ip_address_


def system_resources() -> Dict[str, dict]:
"""Get system resources like CPU, virtual memory and swap memory information.
Returns:
Dict[str, dict]:
Returns a nested dictionary.
"""
return dict(
cpu_usage=psutil.cpu_percent(interval=2, percpu=True),
memory_info=psutil.virtual_memory()._asdict(),
swap_info=psutil.swap_memory()._asdict(),
)


def format_nos(input_: float) -> int | float:
"""Removes ``.0`` float values.
Expand Down

0 comments on commit 1bcc935

Please sign in to comment.