cpu-loader 0.1.2__tar.gz → 0.2.3__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {cpu_loader-0.1.2 → cpu_loader-0.2.3}/MANIFEST.in +1 -1
- {cpu_loader-0.1.2/src/cpu_loader.egg-info → cpu_loader-0.2.3}/PKG-INFO +1 -1
- {cpu_loader-0.1.2 → cpu_loader-0.2.3}/pyproject.toml +6 -3
- {cpu_loader-0.1.2 → cpu_loader-0.2.3}/setup.py +2 -2
- cpu_loader-0.2.3/src/cpu_loader/__init__.py +5 -0
- cpu_loader-0.2.3/src/cpu_loader/__main__.py +6 -0
- {cpu_loader-0.1.2/src → cpu_loader-0.2.3/src/cpu_loader}/cpu_loader.py +1 -1
- cpu_loader-0.2.3/src/cpu_loader/main.py +328 -0
- cpu_loader-0.1.2/src/main.py → cpu_loader-0.2.3/src/cpu_loader/templates/index.html +0 -321
- {cpu_loader-0.1.2 → cpu_loader-0.2.3/src/cpu_loader.egg-info}/PKG-INFO +1 -1
- cpu_loader-0.2.3/src/cpu_loader.egg-info/SOURCES.txt +19 -0
- cpu_loader-0.2.3/src/cpu_loader.egg-info/entry_points.txt +2 -0
- cpu_loader-0.2.3/src/cpu_loader.egg-info/top_level.txt +1 -0
- cpu_loader-0.1.2/src/cpu_loader.egg-info/SOURCES.txt +0 -16
- cpu_loader-0.1.2/src/cpu_loader.egg-info/entry_points.txt +0 -2
- cpu_loader-0.1.2/src/cpu_loader.egg-info/top_level.txt +0 -5
- {cpu_loader-0.1.2 → cpu_loader-0.2.3}/LICENSE +0 -0
- {cpu_loader-0.1.2 → cpu_loader-0.2.3}/README.md +0 -0
- {cpu_loader-0.1.2 → cpu_loader-0.2.3}/setup.cfg +0 -0
- {cpu_loader-0.1.2/src → cpu_loader-0.2.3/src/cpu_loader}/cpu_loader_core.c +0 -0
- {cpu_loader-0.1.2/src → cpu_loader-0.2.3/src/cpu_loader}/example.py +0 -0
- {cpu_loader-0.1.2/src → cpu_loader-0.2.3/src/cpu_loader}/mqtt_publisher.py +0 -0
- {cpu_loader-0.1.2 → cpu_loader-0.2.3}/src/cpu_loader.egg-info/dependency_links.txt +0 -0
- {cpu_loader-0.1.2 → cpu_loader-0.2.3}/src/cpu_loader.egg-info/requires.txt +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
[project]
|
|
2
2
|
name = "cpu-loader"
|
|
3
|
-
version = "0.
|
|
3
|
+
version = "0.2.3"
|
|
4
4
|
description = "A tool to generate CPU load on a system with runtime configuration through a WebUI and REST API"
|
|
5
5
|
readme = "README.md"
|
|
6
6
|
requires-python = ">=3.8"
|
|
@@ -37,7 +37,7 @@ dependencies = [
|
|
|
37
37
|
]
|
|
38
38
|
|
|
39
39
|
[project.scripts]
|
|
40
|
-
cpu-loader = "
|
|
40
|
+
cpu-loader = "cpu_loader.__main__:run"
|
|
41
41
|
|
|
42
42
|
[project.urls]
|
|
43
43
|
Homepage = "https://github.com/the78mole/cpu-loader"
|
|
@@ -51,4 +51,7 @@ build-backend = "setuptools.build_meta"
|
|
|
51
51
|
|
|
52
52
|
[tool.setuptools]
|
|
53
53
|
package-dir = {"" = "src"}
|
|
54
|
-
|
|
54
|
+
packages = ["cpu_loader"]
|
|
55
|
+
|
|
56
|
+
[tool.setuptools.package-data]
|
|
57
|
+
cpu_loader = ["templates/*.html", "*.c"]
|
|
@@ -9,8 +9,8 @@ if platform.system() != 'Windows':
|
|
|
9
9
|
extra_link_args = ['-pthread']
|
|
10
10
|
|
|
11
11
|
module = Extension(
|
|
12
|
-
'cpu_loader_core',
|
|
13
|
-
sources=['src/cpu_loader_core.c'],
|
|
12
|
+
'cpu_loader.cpu_loader_core',
|
|
13
|
+
sources=['src/cpu_loader/cpu_loader_core.c'],
|
|
14
14
|
extra_compile_args=extra_compile_args,
|
|
15
15
|
extra_link_args=extra_link_args,
|
|
16
16
|
)
|
|
@@ -7,7 +7,7 @@ import multiprocessing
|
|
|
7
7
|
from typing import Dict
|
|
8
8
|
|
|
9
9
|
try:
|
|
10
|
-
import cpu_loader_core
|
|
10
|
+
from cpu_loader import cpu_loader_core # type: ignore[attr-defined]
|
|
11
11
|
except ImportError:
|
|
12
12
|
raise ImportError(
|
|
13
13
|
"cpu_loader_core module not found. Please build the C extension with: "
|
|
@@ -0,0 +1,328 @@
|
|
|
1
|
+
"""
|
|
2
|
+
CPU Loader FastAPI Application
|
|
3
|
+
Provides REST API and WebUI for controlling CPU load.
|
|
4
|
+
"""
|
|
5
|
+
import argparse
|
|
6
|
+
import asyncio
|
|
7
|
+
import logging
|
|
8
|
+
from contextlib import asynccontextmanager
|
|
9
|
+
from typing import Dict, List, Optional, Set
|
|
10
|
+
|
|
11
|
+
import psutil
|
|
12
|
+
import uvicorn
|
|
13
|
+
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
|
|
14
|
+
from fastapi.responses import HTMLResponse
|
|
15
|
+
from pydantic import BaseModel, Field
|
|
16
|
+
|
|
17
|
+
from cpu_loader.cpu_loader import CPULoader
|
|
18
|
+
from cpu_loader.mqtt_publisher import MQTTPublisher
|
|
19
|
+
|
|
20
|
+
# Configure logging
|
|
21
|
+
logging.basicConfig(
|
|
22
|
+
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
|
23
|
+
)
|
|
24
|
+
logger = logging.getLogger(__name__)
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
# Pydantic models for request/response validation
|
|
28
|
+
class ThreadLoadRequest(BaseModel):
|
|
29
|
+
load_percent: float = Field(
|
|
30
|
+
..., ge=0, le=100, description="Load percentage (0-100)"
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class AllThreadsLoadRequest(BaseModel):
|
|
35
|
+
load_percent: float = Field(
|
|
36
|
+
..., ge=0, le=100, description="Load percentage (0-100)"
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
class ThreadCountRequest(BaseModel):
|
|
41
|
+
num_threads: int = Field(
|
|
42
|
+
..., gt=0, description="Number of threads (must be positive)"
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class ThreadsStatusResponse(BaseModel):
|
|
47
|
+
num_threads: int
|
|
48
|
+
loads: Dict[int, float]
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
class CPUMetricsResponse(BaseModel):
|
|
52
|
+
total_cpu_percent: float
|
|
53
|
+
per_cpu_percent: List[float]
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
# Global CPU loader instance, MQTT publisher, and WebSocket connections
|
|
57
|
+
cpu_loader = None
|
|
58
|
+
mqtt_publisher: Optional[MQTTPublisher] = None
|
|
59
|
+
websocket_connections: Set[WebSocket] = set()
|
|
60
|
+
monitoring_task = None
|
|
61
|
+
|
|
62
|
+
|
|
63
|
+
async def cpu_monitoring_loop():
|
|
64
|
+
"""Background task that monitors CPU usage and broadcasts to all WebSocket clients."""
|
|
65
|
+
# Initialize psutil
|
|
66
|
+
psutil.cpu_percent(interval=None, percpu=True)
|
|
67
|
+
|
|
68
|
+
while True:
|
|
69
|
+
try:
|
|
70
|
+
# Wait for 1 second
|
|
71
|
+
await asyncio.sleep(1.0)
|
|
72
|
+
|
|
73
|
+
# Get CPU metrics (non-blocking after first call)
|
|
74
|
+
per_cpu = psutil.cpu_percent(interval=None, percpu=True)
|
|
75
|
+
total_cpu = sum(per_cpu) / len(per_cpu) if per_cpu else 0.0
|
|
76
|
+
|
|
77
|
+
# Prepare message
|
|
78
|
+
message = {
|
|
79
|
+
"type": "cpu_metrics",
|
|
80
|
+
"total_cpu_percent": round(total_cpu, 1),
|
|
81
|
+
"per_cpu_percent": [round(cpu, 1) for cpu in per_cpu],
|
|
82
|
+
}
|
|
83
|
+
|
|
84
|
+
# Broadcast to all connected clients
|
|
85
|
+
if websocket_connections:
|
|
86
|
+
disconnected = set()
|
|
87
|
+
for websocket in websocket_connections:
|
|
88
|
+
try:
|
|
89
|
+
await websocket.send_json(message)
|
|
90
|
+
except Exception:
|
|
91
|
+
disconnected.add(websocket)
|
|
92
|
+
|
|
93
|
+
# Remove disconnected clients
|
|
94
|
+
websocket_connections.difference_update(disconnected)
|
|
95
|
+
|
|
96
|
+
# Publish to MQTT if enabled
|
|
97
|
+
if mqtt_publisher:
|
|
98
|
+
mqtt_publisher.publish_cpu_metrics(total_cpu, per_cpu)
|
|
99
|
+
|
|
100
|
+
except Exception as e:
|
|
101
|
+
logger.error(f"Error in CPU monitoring loop: {e}")
|
|
102
|
+
await asyncio.sleep(1.0)
|
|
103
|
+
|
|
104
|
+
|
|
105
|
+
@asynccontextmanager
|
|
106
|
+
async def lifespan(app: FastAPI):
|
|
107
|
+
"""Lifespan context manager for startup and shutdown events."""
|
|
108
|
+
global cpu_loader, mqtt_publisher, monitoring_task
|
|
109
|
+
# Startup
|
|
110
|
+
cpu_loader = CPULoader()
|
|
111
|
+
|
|
112
|
+
# Initialize MQTT publisher with settings from arguments or environment
|
|
113
|
+
mqtt_args = getattr(app.state, "mqtt_args", {})
|
|
114
|
+
try:
|
|
115
|
+
mqtt_publisher = MQTTPublisher(**mqtt_args)
|
|
116
|
+
except ImportError:
|
|
117
|
+
logger.warning("MQTT publishing disabled: paho-mqtt not installed")
|
|
118
|
+
mqtt_publisher = None
|
|
119
|
+
except Exception as e:
|
|
120
|
+
logger.error(f"Failed to initialize MQTT publisher: {e}")
|
|
121
|
+
mqtt_publisher = None
|
|
122
|
+
|
|
123
|
+
monitoring_task = asyncio.create_task(cpu_monitoring_loop())
|
|
124
|
+
yield
|
|
125
|
+
# Shutdown
|
|
126
|
+
if monitoring_task:
|
|
127
|
+
monitoring_task.cancel()
|
|
128
|
+
try:
|
|
129
|
+
await monitoring_task
|
|
130
|
+
except asyncio.CancelledError:
|
|
131
|
+
pass
|
|
132
|
+
if mqtt_publisher:
|
|
133
|
+
mqtt_publisher.disconnect()
|
|
134
|
+
cpu_loader.shutdown()
|
|
135
|
+
|
|
136
|
+
|
|
137
|
+
# Initialize FastAPI app
|
|
138
|
+
app = FastAPI(
|
|
139
|
+
title="CPU Loader API",
|
|
140
|
+
description="Control CPU load generation with configurable threads",
|
|
141
|
+
version="1.0.0",
|
|
142
|
+
lifespan=lifespan,
|
|
143
|
+
)
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
@app.get("/", response_class=HTMLResponse)
|
|
147
|
+
async def get_webui():
|
|
148
|
+
"""Serve the WebUI."""
|
|
149
|
+
from pathlib import Path
|
|
150
|
+
|
|
151
|
+
template_path = Path(__file__).parent / "templates" / "index.html"
|
|
152
|
+
html_content = template_path.read_text()
|
|
153
|
+
return HTMLResponse(content=html_content)
|
|
154
|
+
|
|
155
|
+
|
|
156
|
+
@app.get("/api/threads", response_model=ThreadsStatusResponse)
|
|
157
|
+
async def get_threads_status():
|
|
158
|
+
"""Get the current status of all threads."""
|
|
159
|
+
return ThreadsStatusResponse(
|
|
160
|
+
num_threads=cpu_loader.get_num_threads(), loads=cpu_loader.get_all_loads()
|
|
161
|
+
)
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
@app.websocket("/ws/cpu-metrics")
|
|
165
|
+
async def websocket_cpu_metrics(websocket: WebSocket):
|
|
166
|
+
"""WebSocket endpoint for real-time CPU metrics."""
|
|
167
|
+
await websocket.accept()
|
|
168
|
+
websocket_connections.add(websocket)
|
|
169
|
+
try:
|
|
170
|
+
# Keep connection alive
|
|
171
|
+
while True:
|
|
172
|
+
# Wait for any message (ping/pong)
|
|
173
|
+
await websocket.receive_text()
|
|
174
|
+
except WebSocketDisconnect:
|
|
175
|
+
websocket_connections.discard(websocket)
|
|
176
|
+
except Exception:
|
|
177
|
+
websocket_connections.discard(websocket)
|
|
178
|
+
|
|
179
|
+
|
|
180
|
+
@app.get("/api/cpu-metrics", response_model=CPUMetricsResponse)
|
|
181
|
+
async def get_cpu_metrics():
|
|
182
|
+
"""Get current CPU utilization metrics."""
|
|
183
|
+
per_cpu = psutil.cpu_percent(interval=0.1, percpu=True)
|
|
184
|
+
total_cpu = psutil.cpu_percent(interval=0)
|
|
185
|
+
return CPUMetricsResponse(total_cpu_percent=total_cpu, per_cpu_percent=per_cpu)
|
|
186
|
+
|
|
187
|
+
|
|
188
|
+
@app.post("/api/threads")
|
|
189
|
+
async def set_thread_count(request: ThreadCountRequest):
|
|
190
|
+
"""Set the number of threads."""
|
|
191
|
+
try:
|
|
192
|
+
cpu_loader.set_num_threads(request.num_threads)
|
|
193
|
+
|
|
194
|
+
# Publish updated settings to MQTT
|
|
195
|
+
if mqtt_publisher:
|
|
196
|
+
mqtt_publisher.publish_load_settings(
|
|
197
|
+
cpu_loader.get_num_threads(), cpu_loader.get_all_loads()
|
|
198
|
+
)
|
|
199
|
+
|
|
200
|
+
return {
|
|
201
|
+
"status": "success",
|
|
202
|
+
"num_threads": request.num_threads,
|
|
203
|
+
"message": f"Thread count set to {request.num_threads}",
|
|
204
|
+
}
|
|
205
|
+
except ValueError as e:
|
|
206
|
+
raise HTTPException(status_code=400, detail=str(e))
|
|
207
|
+
|
|
208
|
+
|
|
209
|
+
@app.put("/api/threads/{thread_id}/load")
|
|
210
|
+
async def set_thread_load(thread_id: int, request: ThreadLoadRequest):
|
|
211
|
+
"""Set the CPU load for a specific thread."""
|
|
212
|
+
try:
|
|
213
|
+
cpu_loader.set_thread_load(thread_id, request.load_percent)
|
|
214
|
+
|
|
215
|
+
# Publish updated settings to MQTT
|
|
216
|
+
if mqtt_publisher:
|
|
217
|
+
mqtt_publisher.publish_load_settings(
|
|
218
|
+
cpu_loader.get_num_threads(), cpu_loader.get_all_loads()
|
|
219
|
+
)
|
|
220
|
+
|
|
221
|
+
return {
|
|
222
|
+
"status": "success",
|
|
223
|
+
"thread_id": thread_id,
|
|
224
|
+
"load_percent": request.load_percent,
|
|
225
|
+
"message": f"Thread {thread_id} load set to {request.load_percent}%",
|
|
226
|
+
}
|
|
227
|
+
except ValueError as e:
|
|
228
|
+
raise HTTPException(status_code=400, detail=str(e))
|
|
229
|
+
|
|
230
|
+
|
|
231
|
+
@app.post("/api/threads/load/all")
|
|
232
|
+
async def set_all_thread_loads(request: AllThreadsLoadRequest):
|
|
233
|
+
"""Set the same CPU load for all threads."""
|
|
234
|
+
try:
|
|
235
|
+
cpu_loader.set_all_loads(request.load_percent)
|
|
236
|
+
|
|
237
|
+
# Publish updated settings to MQTT
|
|
238
|
+
if mqtt_publisher:
|
|
239
|
+
mqtt_publisher.publish_load_settings(
|
|
240
|
+
cpu_loader.get_num_threads(), cpu_loader.get_all_loads()
|
|
241
|
+
)
|
|
242
|
+
|
|
243
|
+
return {
|
|
244
|
+
"status": "success",
|
|
245
|
+
"load_percent": request.load_percent,
|
|
246
|
+
"num_threads": cpu_loader.get_num_threads(),
|
|
247
|
+
"message": f"All threads set to {request.load_percent}%",
|
|
248
|
+
}
|
|
249
|
+
except ValueError as e:
|
|
250
|
+
raise HTTPException(status_code=400, detail=str(e))
|
|
251
|
+
|
|
252
|
+
|
|
253
|
+
def parse_args():
|
|
254
|
+
"""Parse command-line arguments."""
|
|
255
|
+
parser = argparse.ArgumentParser(
|
|
256
|
+
description="CPU Loader - Generate controllable CPU load with REST API and WebUI"
|
|
257
|
+
)
|
|
258
|
+
parser.add_argument(
|
|
259
|
+
"--host",
|
|
260
|
+
default="0.0.0.0",
|
|
261
|
+
help="Host to bind the server to (default: 0.0.0.0)",
|
|
262
|
+
)
|
|
263
|
+
parser.add_argument(
|
|
264
|
+
"--port",
|
|
265
|
+
type=int,
|
|
266
|
+
default=8000,
|
|
267
|
+
help="Port to bind the server to (default: 8000)",
|
|
268
|
+
)
|
|
269
|
+
|
|
270
|
+
# MQTT arguments
|
|
271
|
+
mqtt_group = parser.add_argument_group("MQTT settings")
|
|
272
|
+
mqtt_group.add_argument(
|
|
273
|
+
"--mqtt-broker-host",
|
|
274
|
+
help="MQTT broker hostname (env: MQTT_BROKER_HOST)",
|
|
275
|
+
)
|
|
276
|
+
mqtt_group.add_argument(
|
|
277
|
+
"--mqtt-broker-port",
|
|
278
|
+
type=int,
|
|
279
|
+
help="MQTT broker port (env: MQTT_BROKER_PORT, default: 1883)",
|
|
280
|
+
)
|
|
281
|
+
mqtt_group.add_argument(
|
|
282
|
+
"--mqtt-username",
|
|
283
|
+
help="MQTT username (env: MQTT_USERNAME)",
|
|
284
|
+
)
|
|
285
|
+
mqtt_group.add_argument(
|
|
286
|
+
"--mqtt-password",
|
|
287
|
+
help="MQTT password (env: MQTT_PASSWORD)",
|
|
288
|
+
)
|
|
289
|
+
mqtt_group.add_argument(
|
|
290
|
+
"--mqtt-topic-prefix",
|
|
291
|
+
help="MQTT topic prefix (env: MQTT_TOPIC_PREFIX, default: cpu-loader)",
|
|
292
|
+
)
|
|
293
|
+
mqtt_group.add_argument(
|
|
294
|
+
"--mqtt-client-id",
|
|
295
|
+
help="MQTT client ID (env: MQTT_CLIENT_ID, default: cpu-loader)",
|
|
296
|
+
)
|
|
297
|
+
|
|
298
|
+
return parser.parse_args()
|
|
299
|
+
|
|
300
|
+
|
|
301
|
+
def run():
|
|
302
|
+
"""Entry point for the CPU Loader application."""
|
|
303
|
+
args = parse_args()
|
|
304
|
+
|
|
305
|
+
# Prepare MQTT arguments (only non-None values)
|
|
306
|
+
mqtt_args = {}
|
|
307
|
+
if args.mqtt_broker_host:
|
|
308
|
+
mqtt_args["broker_host"] = args.mqtt_broker_host
|
|
309
|
+
if args.mqtt_broker_port:
|
|
310
|
+
mqtt_args["broker_port"] = args.mqtt_broker_port
|
|
311
|
+
if args.mqtt_username:
|
|
312
|
+
mqtt_args["username"] = args.mqtt_username
|
|
313
|
+
if args.mqtt_password:
|
|
314
|
+
mqtt_args["password"] = args.mqtt_password
|
|
315
|
+
if args.mqtt_topic_prefix:
|
|
316
|
+
mqtt_args["topic_prefix"] = args.mqtt_topic_prefix
|
|
317
|
+
if args.mqtt_client_id:
|
|
318
|
+
mqtt_args["client_id"] = args.mqtt_client_id
|
|
319
|
+
|
|
320
|
+
# Store MQTT args in app state for lifespan to access
|
|
321
|
+
app.state.mqtt_args = mqtt_args
|
|
322
|
+
|
|
323
|
+
# Run the server
|
|
324
|
+
uvicorn.run(app, host=args.host, port=args.port)
|
|
325
|
+
|
|
326
|
+
|
|
327
|
+
if __name__ == "__main__":
|
|
328
|
+
run()
|
|
@@ -1,152 +1,3 @@
|
|
|
1
|
-
"""
|
|
2
|
-
CPU Loader FastAPI Application
|
|
3
|
-
Provides REST API and WebUI for controlling CPU load.
|
|
4
|
-
"""
|
|
5
|
-
import argparse
|
|
6
|
-
import asyncio
|
|
7
|
-
import logging
|
|
8
|
-
from contextlib import asynccontextmanager
|
|
9
|
-
from typing import Dict, List, Optional, Set
|
|
10
|
-
|
|
11
|
-
import psutil
|
|
12
|
-
import uvicorn
|
|
13
|
-
from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
|
|
14
|
-
from fastapi.responses import HTMLResponse
|
|
15
|
-
from pydantic import BaseModel, Field
|
|
16
|
-
|
|
17
|
-
from cpu_loader import CPULoader
|
|
18
|
-
from mqtt_publisher import MQTTPublisher
|
|
19
|
-
|
|
20
|
-
# Configure logging
|
|
21
|
-
logging.basicConfig(
|
|
22
|
-
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
|
23
|
-
)
|
|
24
|
-
logger = logging.getLogger(__name__)
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
# Pydantic models for request/response validation
|
|
28
|
-
class ThreadLoadRequest(BaseModel):
|
|
29
|
-
load_percent: float = Field(
|
|
30
|
-
..., ge=0, le=100, description="Load percentage (0-100)"
|
|
31
|
-
)
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
class AllThreadsLoadRequest(BaseModel):
|
|
35
|
-
load_percent: float = Field(
|
|
36
|
-
..., ge=0, le=100, description="Load percentage (0-100)"
|
|
37
|
-
)
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
class ThreadCountRequest(BaseModel):
|
|
41
|
-
num_threads: int = Field(
|
|
42
|
-
..., gt=0, description="Number of threads (must be positive)"
|
|
43
|
-
)
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
class ThreadsStatusResponse(BaseModel):
|
|
47
|
-
num_threads: int
|
|
48
|
-
loads: Dict[int, float]
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
class CPUMetricsResponse(BaseModel):
|
|
52
|
-
total_cpu_percent: float
|
|
53
|
-
per_cpu_percent: List[float]
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
# Global CPU loader instance, MQTT publisher, and WebSocket connections
|
|
57
|
-
cpu_loader = None
|
|
58
|
-
mqtt_publisher: Optional[MQTTPublisher] = None
|
|
59
|
-
websocket_connections: Set[WebSocket] = set()
|
|
60
|
-
monitoring_task = None
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
async def cpu_monitoring_loop():
|
|
64
|
-
"""Background task that monitors CPU usage and broadcasts to all WebSocket clients."""
|
|
65
|
-
# Initialize psutil
|
|
66
|
-
psutil.cpu_percent(interval=None, percpu=True)
|
|
67
|
-
|
|
68
|
-
while True:
|
|
69
|
-
try:
|
|
70
|
-
# Wait for 1 second
|
|
71
|
-
await asyncio.sleep(1.0)
|
|
72
|
-
|
|
73
|
-
# Get CPU metrics (non-blocking after first call)
|
|
74
|
-
per_cpu = psutil.cpu_percent(interval=None, percpu=True)
|
|
75
|
-
total_cpu = sum(per_cpu) / len(per_cpu) if per_cpu else 0.0
|
|
76
|
-
|
|
77
|
-
# Prepare message
|
|
78
|
-
message = {
|
|
79
|
-
"type": "cpu_metrics",
|
|
80
|
-
"total_cpu_percent": round(total_cpu, 1),
|
|
81
|
-
"per_cpu_percent": [round(cpu, 1) for cpu in per_cpu],
|
|
82
|
-
}
|
|
83
|
-
|
|
84
|
-
# Broadcast to all connected clients
|
|
85
|
-
if websocket_connections:
|
|
86
|
-
disconnected = set()
|
|
87
|
-
for websocket in websocket_connections:
|
|
88
|
-
try:
|
|
89
|
-
await websocket.send_json(message)
|
|
90
|
-
except Exception:
|
|
91
|
-
disconnected.add(websocket)
|
|
92
|
-
|
|
93
|
-
# Remove disconnected clients
|
|
94
|
-
websocket_connections.difference_update(disconnected)
|
|
95
|
-
|
|
96
|
-
# Publish to MQTT if enabled
|
|
97
|
-
if mqtt_publisher:
|
|
98
|
-
mqtt_publisher.publish_cpu_metrics(total_cpu, per_cpu)
|
|
99
|
-
|
|
100
|
-
except Exception as e:
|
|
101
|
-
logger.error(f"Error in CPU monitoring loop: {e}")
|
|
102
|
-
await asyncio.sleep(1.0)
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
@asynccontextmanager
|
|
106
|
-
async def lifespan(app: FastAPI):
|
|
107
|
-
"""Lifespan context manager for startup and shutdown events."""
|
|
108
|
-
global cpu_loader, mqtt_publisher, monitoring_task
|
|
109
|
-
# Startup
|
|
110
|
-
cpu_loader = CPULoader()
|
|
111
|
-
|
|
112
|
-
# Initialize MQTT publisher with settings from arguments or environment
|
|
113
|
-
mqtt_args = getattr(app.state, "mqtt_args", {})
|
|
114
|
-
try:
|
|
115
|
-
mqtt_publisher = MQTTPublisher(**mqtt_args)
|
|
116
|
-
except ImportError:
|
|
117
|
-
logger.warning("MQTT publishing disabled: paho-mqtt not installed")
|
|
118
|
-
mqtt_publisher = None
|
|
119
|
-
except Exception as e:
|
|
120
|
-
logger.error(f"Failed to initialize MQTT publisher: {e}")
|
|
121
|
-
mqtt_publisher = None
|
|
122
|
-
|
|
123
|
-
monitoring_task = asyncio.create_task(cpu_monitoring_loop())
|
|
124
|
-
yield
|
|
125
|
-
# Shutdown
|
|
126
|
-
if monitoring_task:
|
|
127
|
-
monitoring_task.cancel()
|
|
128
|
-
try:
|
|
129
|
-
await monitoring_task
|
|
130
|
-
except asyncio.CancelledError:
|
|
131
|
-
pass
|
|
132
|
-
if mqtt_publisher:
|
|
133
|
-
mqtt_publisher.disconnect()
|
|
134
|
-
cpu_loader.shutdown()
|
|
135
|
-
|
|
136
|
-
|
|
137
|
-
# Initialize FastAPI app
|
|
138
|
-
app = FastAPI(
|
|
139
|
-
title="CPU Loader API",
|
|
140
|
-
description="Control CPU load generation with configurable threads",
|
|
141
|
-
version="1.0.0",
|
|
142
|
-
lifespan=lifespan,
|
|
143
|
-
)
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
@app.get("/", response_class=HTMLResponse)
|
|
147
|
-
async def get_webui():
|
|
148
|
-
"""Serve the WebUI."""
|
|
149
|
-
html_content = """
|
|
150
1
|
<!DOCTYPE html>
|
|
151
2
|
<html lang="en">
|
|
152
3
|
<head>
|
|
@@ -603,175 +454,3 @@ async def get_webui():
|
|
|
603
454
|
</script>
|
|
604
455
|
</body>
|
|
605
456
|
</html>
|
|
606
|
-
"""
|
|
607
|
-
return HTMLResponse(content=html_content)
|
|
608
|
-
|
|
609
|
-
|
|
610
|
-
@app.get("/api/threads", response_model=ThreadsStatusResponse)
|
|
611
|
-
async def get_threads_status():
|
|
612
|
-
"""Get the current status of all threads."""
|
|
613
|
-
return ThreadsStatusResponse(
|
|
614
|
-
num_threads=cpu_loader.get_num_threads(), loads=cpu_loader.get_all_loads()
|
|
615
|
-
)
|
|
616
|
-
|
|
617
|
-
|
|
618
|
-
@app.websocket("/ws/cpu-metrics")
|
|
619
|
-
async def websocket_cpu_metrics(websocket: WebSocket):
|
|
620
|
-
"""WebSocket endpoint for real-time CPU metrics."""
|
|
621
|
-
await websocket.accept()
|
|
622
|
-
websocket_connections.add(websocket)
|
|
623
|
-
try:
|
|
624
|
-
# Keep connection alive
|
|
625
|
-
while True:
|
|
626
|
-
# Wait for any message (ping/pong)
|
|
627
|
-
await websocket.receive_text()
|
|
628
|
-
except WebSocketDisconnect:
|
|
629
|
-
websocket_connections.discard(websocket)
|
|
630
|
-
except Exception:
|
|
631
|
-
websocket_connections.discard(websocket)
|
|
632
|
-
|
|
633
|
-
|
|
634
|
-
@app.get("/api/cpu-metrics", response_model=CPUMetricsResponse)
|
|
635
|
-
async def get_cpu_metrics():
|
|
636
|
-
"""Get current CPU utilization metrics."""
|
|
637
|
-
per_cpu = psutil.cpu_percent(interval=0.1, percpu=True)
|
|
638
|
-
total_cpu = psutil.cpu_percent(interval=0)
|
|
639
|
-
return CPUMetricsResponse(total_cpu_percent=total_cpu, per_cpu_percent=per_cpu)
|
|
640
|
-
|
|
641
|
-
|
|
642
|
-
@app.post("/api/threads")
|
|
643
|
-
async def set_thread_count(request: ThreadCountRequest):
|
|
644
|
-
"""Set the number of threads."""
|
|
645
|
-
try:
|
|
646
|
-
cpu_loader.set_num_threads(request.num_threads)
|
|
647
|
-
|
|
648
|
-
# Publish updated settings to MQTT
|
|
649
|
-
if mqtt_publisher:
|
|
650
|
-
mqtt_publisher.publish_load_settings(
|
|
651
|
-
cpu_loader.get_num_threads(), cpu_loader.get_all_loads()
|
|
652
|
-
)
|
|
653
|
-
|
|
654
|
-
return {
|
|
655
|
-
"status": "success",
|
|
656
|
-
"num_threads": request.num_threads,
|
|
657
|
-
"message": f"Thread count set to {request.num_threads}",
|
|
658
|
-
}
|
|
659
|
-
except ValueError as e:
|
|
660
|
-
raise HTTPException(status_code=400, detail=str(e))
|
|
661
|
-
|
|
662
|
-
|
|
663
|
-
@app.put("/api/threads/{thread_id}/load")
|
|
664
|
-
async def set_thread_load(thread_id: int, request: ThreadLoadRequest):
|
|
665
|
-
"""Set the CPU load for a specific thread."""
|
|
666
|
-
try:
|
|
667
|
-
cpu_loader.set_thread_load(thread_id, request.load_percent)
|
|
668
|
-
|
|
669
|
-
# Publish updated settings to MQTT
|
|
670
|
-
if mqtt_publisher:
|
|
671
|
-
mqtt_publisher.publish_load_settings(
|
|
672
|
-
cpu_loader.get_num_threads(), cpu_loader.get_all_loads()
|
|
673
|
-
)
|
|
674
|
-
|
|
675
|
-
return {
|
|
676
|
-
"status": "success",
|
|
677
|
-
"thread_id": thread_id,
|
|
678
|
-
"load_percent": request.load_percent,
|
|
679
|
-
"message": f"Thread {thread_id} load set to {request.load_percent}%",
|
|
680
|
-
}
|
|
681
|
-
except ValueError as e:
|
|
682
|
-
raise HTTPException(status_code=400, detail=str(e))
|
|
683
|
-
|
|
684
|
-
|
|
685
|
-
@app.post("/api/threads/load/all")
|
|
686
|
-
async def set_all_thread_loads(request: AllThreadsLoadRequest):
|
|
687
|
-
"""Set the same CPU load for all threads."""
|
|
688
|
-
try:
|
|
689
|
-
cpu_loader.set_all_loads(request.load_percent)
|
|
690
|
-
|
|
691
|
-
# Publish updated settings to MQTT
|
|
692
|
-
if mqtt_publisher:
|
|
693
|
-
mqtt_publisher.publish_load_settings(
|
|
694
|
-
cpu_loader.get_num_threads(), cpu_loader.get_all_loads()
|
|
695
|
-
)
|
|
696
|
-
|
|
697
|
-
return {
|
|
698
|
-
"status": "success",
|
|
699
|
-
"load_percent": request.load_percent,
|
|
700
|
-
"num_threads": cpu_loader.get_num_threads(),
|
|
701
|
-
"message": f"All threads set to {request.load_percent}%",
|
|
702
|
-
}
|
|
703
|
-
except ValueError as e:
|
|
704
|
-
raise HTTPException(status_code=400, detail=str(e))
|
|
705
|
-
|
|
706
|
-
|
|
707
|
-
def parse_args():
|
|
708
|
-
"""Parse command-line arguments."""
|
|
709
|
-
parser = argparse.ArgumentParser(
|
|
710
|
-
description="CPU Loader - Generate controllable CPU load with REST API and WebUI"
|
|
711
|
-
)
|
|
712
|
-
parser.add_argument(
|
|
713
|
-
"--host",
|
|
714
|
-
default="0.0.0.0",
|
|
715
|
-
help="Host to bind the server to (default: 0.0.0.0)",
|
|
716
|
-
)
|
|
717
|
-
parser.add_argument(
|
|
718
|
-
"--port",
|
|
719
|
-
type=int,
|
|
720
|
-
default=8000,
|
|
721
|
-
help="Port to bind the server to (default: 8000)",
|
|
722
|
-
)
|
|
723
|
-
|
|
724
|
-
# MQTT arguments
|
|
725
|
-
mqtt_group = parser.add_argument_group("MQTT settings")
|
|
726
|
-
mqtt_group.add_argument(
|
|
727
|
-
"--mqtt-broker-host",
|
|
728
|
-
help="MQTT broker hostname (env: MQTT_BROKER_HOST)",
|
|
729
|
-
)
|
|
730
|
-
mqtt_group.add_argument(
|
|
731
|
-
"--mqtt-broker-port",
|
|
732
|
-
type=int,
|
|
733
|
-
help="MQTT broker port (env: MQTT_BROKER_PORT, default: 1883)",
|
|
734
|
-
)
|
|
735
|
-
mqtt_group.add_argument(
|
|
736
|
-
"--mqtt-username",
|
|
737
|
-
help="MQTT username (env: MQTT_USERNAME)",
|
|
738
|
-
)
|
|
739
|
-
mqtt_group.add_argument(
|
|
740
|
-
"--mqtt-password",
|
|
741
|
-
help="MQTT password (env: MQTT_PASSWORD)",
|
|
742
|
-
)
|
|
743
|
-
mqtt_group.add_argument(
|
|
744
|
-
"--mqtt-topic-prefix",
|
|
745
|
-
help="MQTT topic prefix (env: MQTT_TOPIC_PREFIX, default: cpu-loader)",
|
|
746
|
-
)
|
|
747
|
-
mqtt_group.add_argument(
|
|
748
|
-
"--mqtt-client-id",
|
|
749
|
-
help="MQTT client ID (env: MQTT_CLIENT_ID, default: cpu-loader)",
|
|
750
|
-
)
|
|
751
|
-
|
|
752
|
-
return parser.parse_args()
|
|
753
|
-
|
|
754
|
-
|
|
755
|
-
if __name__ == "__main__":
|
|
756
|
-
args = parse_args()
|
|
757
|
-
|
|
758
|
-
# Prepare MQTT arguments (only non-None values)
|
|
759
|
-
mqtt_args = {}
|
|
760
|
-
if args.mqtt_broker_host:
|
|
761
|
-
mqtt_args["broker_host"] = args.mqtt_broker_host
|
|
762
|
-
if args.mqtt_broker_port:
|
|
763
|
-
mqtt_args["broker_port"] = args.mqtt_broker_port
|
|
764
|
-
if args.mqtt_username:
|
|
765
|
-
mqtt_args["username"] = args.mqtt_username
|
|
766
|
-
if args.mqtt_password:
|
|
767
|
-
mqtt_args["password"] = args.mqtt_password
|
|
768
|
-
if args.mqtt_topic_prefix:
|
|
769
|
-
mqtt_args["topic_prefix"] = args.mqtt_topic_prefix
|
|
770
|
-
if args.mqtt_client_id:
|
|
771
|
-
mqtt_args["client_id"] = args.mqtt_client_id
|
|
772
|
-
|
|
773
|
-
# Store MQTT args in app state for lifespan to access
|
|
774
|
-
app.state.mqtt_args = mqtt_args
|
|
775
|
-
|
|
776
|
-
# Run the server
|
|
777
|
-
uvicorn.run(app, host=args.host, port=args.port)
|
|
@@ -0,0 +1,19 @@
|
|
|
1
|
+
LICENSE
|
|
2
|
+
MANIFEST.in
|
|
3
|
+
README.md
|
|
4
|
+
pyproject.toml
|
|
5
|
+
setup.py
|
|
6
|
+
src/cpu_loader/__init__.py
|
|
7
|
+
src/cpu_loader/__main__.py
|
|
8
|
+
src/cpu_loader/cpu_loader.py
|
|
9
|
+
src/cpu_loader/cpu_loader_core.c
|
|
10
|
+
src/cpu_loader/example.py
|
|
11
|
+
src/cpu_loader/main.py
|
|
12
|
+
src/cpu_loader/mqtt_publisher.py
|
|
13
|
+
src/cpu_loader.egg-info/PKG-INFO
|
|
14
|
+
src/cpu_loader.egg-info/SOURCES.txt
|
|
15
|
+
src/cpu_loader.egg-info/dependency_links.txt
|
|
16
|
+
src/cpu_loader.egg-info/entry_points.txt
|
|
17
|
+
src/cpu_loader.egg-info/requires.txt
|
|
18
|
+
src/cpu_loader.egg-info/top_level.txt
|
|
19
|
+
src/cpu_loader/templates/index.html
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
cpu_loader
|
|
@@ -1,16 +0,0 @@
|
|
|
1
|
-
LICENSE
|
|
2
|
-
MANIFEST.in
|
|
3
|
-
README.md
|
|
4
|
-
pyproject.toml
|
|
5
|
-
setup.py
|
|
6
|
-
src/cpu_loader.py
|
|
7
|
-
src/cpu_loader_core.c
|
|
8
|
-
src/example.py
|
|
9
|
-
src/main.py
|
|
10
|
-
src/mqtt_publisher.py
|
|
11
|
-
src/cpu_loader.egg-info/PKG-INFO
|
|
12
|
-
src/cpu_loader.egg-info/SOURCES.txt
|
|
13
|
-
src/cpu_loader.egg-info/dependency_links.txt
|
|
14
|
-
src/cpu_loader.egg-info/entry_points.txt
|
|
15
|
-
src/cpu_loader.egg-info/requires.txt
|
|
16
|
-
src/cpu_loader.egg-info/top_level.txt
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|