533 lines
17 KiB
Python
Executable File
533 lines
17 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
Heat Generator - CPU & NPU Stress Test
|
|
For testing cooling solutions on Rockchip boards
|
|
With real-time temperature graphing
|
|
"""
|
|
|
|
import os
|
|
import sys
|
|
import time
|
|
import glob
|
|
import signal
|
|
import argparse
|
|
import multiprocessing
|
|
from multiprocessing import Process, Value
|
|
from ctypes import c_bool
|
|
from collections import deque
|
|
from io import StringIO
|
|
|
|
# Check for numpy (used for heavier CPU load)
|
|
try:
|
|
import numpy as np
|
|
HAS_NUMPY = True
|
|
except ImportError:
|
|
HAS_NUMPY = False
|
|
|
|
# Check for RKNN Lite (NPU)
|
|
try:
|
|
from rknnlite.api import RKNNLite
|
|
HAS_RKNN = True
|
|
except ImportError:
|
|
HAS_RKNN = False
|
|
|
|
# Check for plotext (graphing)
|
|
try:
|
|
import plotext as plt
|
|
HAS_PLOTEXT = True
|
|
except ImportError:
|
|
HAS_PLOTEXT = False
|
|
|
|
|
|
# ANSI escape codes
|
|
RED = "\033[91m"
|
|
GREEN = "\033[92m"
|
|
YELLOW = "\033[93m"
|
|
CYAN = "\033[96m"
|
|
RESET = "\033[0m"
|
|
BOLD = "\033[1m"
|
|
|
|
# Flicker-free display
|
|
CURSOR_HOME = "\033[H"
|
|
CLEAR_SCREEN = "\033[2J"
|
|
CLEAR_TO_END = "\033[J"
|
|
HIDE_CURSOR = "\033[?25l"
|
|
SHOW_CURSOR = "\033[?25h"
|
|
# Alternate screen buffer (like vim/htop use)
|
|
ALT_SCREEN_ON = "\033[?1049h"
|
|
ALT_SCREEN_OFF = "\033[?1049l"
|
|
|
|
|
|
def get_terminal_size() -> tuple[int, int]:
|
|
"""Get terminal size (columns, rows)"""
|
|
try:
|
|
import shutil
|
|
size = shutil.get_terminal_size()
|
|
return size.columns, size.lines
|
|
except:
|
|
return 80, 24
|
|
|
|
# Graph colors
|
|
GRAPH_COLORS = ['red', 'cyan', 'green', 'yellow', 'magenta']
|
|
|
|
|
|
class TemperatureHistory:
|
|
"""Track temperature history for graphing"""
|
|
|
|
def __init__(self, max_samples: int = 120):
|
|
self.max_samples = max_samples
|
|
self.history: dict[str, deque] = {}
|
|
self.time_points: deque = deque(maxlen=max_samples)
|
|
self.start_time = time.time()
|
|
self.start_temps: dict[str, float] = {}
|
|
self.max_temps: dict[str, float] = {}
|
|
self.min_temps: dict[str, float] = {}
|
|
|
|
def update(self, temps: dict[str, float]):
|
|
"""Add new temperature readings"""
|
|
current_time = time.time() - self.start_time
|
|
self.time_points.append(current_time)
|
|
|
|
for name, temp in temps.items():
|
|
if name not in self.history:
|
|
self.history[name] = deque(maxlen=self.max_samples)
|
|
self.start_temps[name] = temp
|
|
self.max_temps[name] = temp
|
|
self.min_temps[name] = temp
|
|
|
|
self.history[name].append(temp)
|
|
self.max_temps[name] = max(self.max_temps[name], temp)
|
|
self.min_temps[name] = min(self.min_temps[name], temp)
|
|
|
|
def get_graph(self, height: int = 12) -> str:
|
|
"""Generate temperature graph as string"""
|
|
if not HAS_PLOTEXT or len(self.time_points) < 2:
|
|
return self._get_ascii_graph()
|
|
|
|
try:
|
|
plt.clear_figure()
|
|
plt.clear_data()
|
|
|
|
plt.theme('dark')
|
|
plt.title("🌡️ Temperature History")
|
|
plt.xlabel("Time (seconds)")
|
|
plt.ylabel("Temp (°C)")
|
|
|
|
time_list = list(self.time_points)
|
|
|
|
for i, (name, temps) in enumerate(self.history.items()):
|
|
temps_list = list(temps)
|
|
friendly = name.replace('-thermal', '').replace('_thermal', '').upper()
|
|
color = GRAPH_COLORS[i % len(GRAPH_COLORS)]
|
|
|
|
min_len = min(len(time_list), len(temps_list))
|
|
if min_len > 1:
|
|
plt.plot(
|
|
time_list[-min_len:],
|
|
temps_list[-min_len:],
|
|
label=friendly,
|
|
color=color
|
|
)
|
|
|
|
# Set y-axis range
|
|
all_temps = [t for temps in self.history.values() for t in temps]
|
|
if all_temps:
|
|
min_temp = max(0, min(all_temps) - 5)
|
|
max_temp = min(105, max(all_temps) + 5)
|
|
plt.ylim(min_temp, max_temp)
|
|
|
|
plt.plotsize(None, height)
|
|
plt.build()
|
|
return plt.active().build()
|
|
except Exception:
|
|
return self._get_ascii_graph()
|
|
|
|
def _get_ascii_graph(self) -> str:
|
|
"""Fallback ASCII sparkline graph"""
|
|
lines = []
|
|
lines.append("\n📊 Temperature History")
|
|
lines.append("-" * 50)
|
|
|
|
chars = " ▁▂▃▄▅▆▇█"
|
|
|
|
for name, temps in self.history.items():
|
|
temps_list = list(temps)[-50:]
|
|
if not temps_list:
|
|
continue
|
|
|
|
friendly = name.replace('-thermal', '').replace('_thermal', '').upper()
|
|
min_t, max_t = min(temps_list), max(temps_list)
|
|
range_t = max_t - min_t if max_t > min_t else 1
|
|
|
|
sparkline = ""
|
|
for t in temps_list:
|
|
idx = int((t - min_t) / range_t * (len(chars) - 1))
|
|
sparkline += chars[idx]
|
|
|
|
lines.append(f"{friendly:8s} [{min_t:5.1f}-{max_t:5.1f}°C] {sparkline}")
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
def cpu_stress_worker(worker_id: int, running: Value):
|
|
"""CPU stress worker - performs intensive calculations."""
|
|
if HAS_NUMPY:
|
|
size = 500
|
|
while running.value:
|
|
a = np.random.rand(size, size).astype(np.float32)
|
|
b = np.random.rand(size, size).astype(np.float32)
|
|
c = np.dot(a, b)
|
|
np.sin(a)
|
|
np.cos(b)
|
|
np.exp(c / 1000)
|
|
else:
|
|
def is_prime(n):
|
|
if n < 2:
|
|
return False
|
|
for i in range(2, int(n ** 0.5) + 1):
|
|
if n % i == 0:
|
|
return False
|
|
return True
|
|
|
|
n = 2
|
|
while running.value:
|
|
is_prime(n)
|
|
n += 1
|
|
if n > 1000000:
|
|
n = 2
|
|
|
|
|
|
def check_npu_available() -> bool:
|
|
"""Check if NPU is actually available"""
|
|
if os.path.exists('/dev/rknpu'):
|
|
return True
|
|
|
|
npu_paths = [
|
|
'/sys/devices/platform/fde40000.npu',
|
|
'/sys/devices/platform/fdab0000.npu',
|
|
]
|
|
for path in npu_paths:
|
|
if os.path.exists(path):
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def gpu_stress_worker(running: Value):
|
|
"""GPU stress worker - heavy compute operations"""
|
|
size = 1024
|
|
while running.value:
|
|
a = np.random.rand(size, size).astype(np.float32)
|
|
b = np.random.rand(size, size).astype(np.float32)
|
|
c = np.dot(a, b)
|
|
np.fft.fft2(a)
|
|
np.linalg.svd(a[:256, :256], compute_uv=False)
|
|
|
|
|
|
def npu_stress_worker(model_path: str, running: Value):
|
|
"""NPU stress worker - runs continuous inference on the NPU."""
|
|
if not HAS_RKNN:
|
|
return
|
|
|
|
if not check_npu_available():
|
|
gpu_stress_worker(running)
|
|
return
|
|
|
|
# Suppress ALL RKNN output (it writes to both stdout and stderr at C level)
|
|
# Save original file descriptors
|
|
stdout_fd = sys.stdout.fileno()
|
|
stderr_fd = sys.stderr.fileno()
|
|
saved_stdout = os.dup(stdout_fd)
|
|
saved_stderr = os.dup(stderr_fd)
|
|
|
|
# Redirect both to /dev/null
|
|
devnull_fd = os.open(os.devnull, os.O_WRONLY)
|
|
os.dup2(devnull_fd, stdout_fd)
|
|
os.dup2(devnull_fd, stderr_fd)
|
|
os.close(devnull_fd)
|
|
|
|
try:
|
|
rknn = RKNNLite()
|
|
|
|
if not model_path or not os.path.exists(model_path):
|
|
return
|
|
|
|
ret = rknn.load_rknn(model_path)
|
|
if ret != 0:
|
|
return
|
|
|
|
ret = rknn.init_runtime()
|
|
if ret != 0:
|
|
rknn.release()
|
|
return
|
|
|
|
input_shapes = [
|
|
(1, 320, 320, 3),
|
|
(1, 3, 320, 320),
|
|
(320, 320, 3),
|
|
(1, 224, 224, 3),
|
|
(1, 3, 224, 224),
|
|
]
|
|
|
|
input_data = None
|
|
for shape in input_shapes:
|
|
try:
|
|
test_input = np.random.randint(0, 255, shape, dtype=np.uint8)
|
|
rknn.inference(inputs=[test_input])
|
|
input_data = test_input
|
|
break
|
|
except:
|
|
continue
|
|
|
|
if input_data is None:
|
|
rknn.release()
|
|
return
|
|
|
|
while running.value:
|
|
input_data = np.random.randint(0, 255, input_data.shape, dtype=np.uint8)
|
|
rknn.inference(inputs=[input_data])
|
|
|
|
rknn.release()
|
|
|
|
finally:
|
|
# Restore stdout and stderr
|
|
os.dup2(saved_stdout, stdout_fd)
|
|
os.dup2(saved_stderr, stderr_fd)
|
|
os.close(saved_stdout)
|
|
os.close(saved_stderr)
|
|
|
|
|
|
def get_temperatures() -> dict[str, float]:
|
|
"""Get current temperatures"""
|
|
temps = {}
|
|
thermal_path = "/sys/class/thermal"
|
|
|
|
for zone_path in glob.glob(f"{thermal_path}/thermal_zone*"):
|
|
try:
|
|
with open(f"{zone_path}/type") as f:
|
|
zone_type = f.read().strip()
|
|
with open(f"{zone_path}/temp") as f:
|
|
temp = int(f.read().strip()) / 1000.0
|
|
temps[zone_type] = temp
|
|
except:
|
|
pass
|
|
|
|
return temps
|
|
|
|
|
|
def get_temp_color(temp: float) -> str:
|
|
"""Get color based on temperature"""
|
|
if temp >= 80:
|
|
return RED
|
|
elif temp >= 60:
|
|
return YELLOW
|
|
return GREEN
|
|
|
|
|
|
def get_status_icon(temp: float) -> str:
|
|
"""Get status icon based on temperature"""
|
|
if temp >= 80:
|
|
return "🔴"
|
|
elif temp >= 60:
|
|
return "🟡"
|
|
return "🟢"
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="🔥 Heat Generator - Stress test CPU & NPU with live temperature graph"
|
|
)
|
|
parser.add_argument('-c', '--cpu-only', action='store_true', help='Only stress CPU')
|
|
parser.add_argument('-n', '--npu-only', action='store_true', help='Only stress NPU')
|
|
parser.add_argument('-w', '--workers', type=int, default=0, help='Number of CPU workers (default: all cores)')
|
|
parser.add_argument('-m', '--model', type=str, default='/home/radxa/little_sophia_brain/RetinaFace.rknn', help='RKNN model path')
|
|
parser.add_argument('-t', '--time', type=int, default=0, help='Run for N seconds (0 = until Ctrl+C)')
|
|
parser.add_argument('--no-graph', action='store_true', help='Disable temperature graph')
|
|
|
|
args = parser.parse_args()
|
|
|
|
num_cpus = multiprocessing.cpu_count()
|
|
num_workers = args.workers if args.workers > 0 else num_cpus
|
|
npu_available = check_npu_available()
|
|
|
|
# Shared flag for stopping workers
|
|
running = Value(c_bool, True)
|
|
processes = []
|
|
|
|
# Signal handler
|
|
def signal_handler(sig, frame):
|
|
running.value = False
|
|
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
|
|
# Get terminal size for graph height
|
|
term_cols, term_rows = get_terminal_size()
|
|
|
|
# Calculate available space for graph:
|
|
# Header: 4 lines (═══, title, ═══, blank)
|
|
# Status: 2 lines (info, blank)
|
|
# Table: 5 lines (header, ─── , sensor1, sensor2, ───)
|
|
# Footer: 2 lines (blank, Ctrl+C)
|
|
# Total fixed: 13 lines
|
|
# Graph overhead: ~4 lines (title, y-label area, x-axis, x-label)
|
|
# So graph plot area = term_rows - 13 - 4 = term_rows - 17
|
|
|
|
# Function to get graph height based on current terminal size
|
|
def get_current_graph_height():
|
|
_, rows = get_terminal_size()
|
|
return max(6, rows - 17)
|
|
|
|
# Initialize display - use alternate screen buffer for clean full-screen display
|
|
sys.stdout.write(ALT_SCREEN_ON + HIDE_CURSOR + CLEAR_SCREEN + CURSOR_HOME)
|
|
sys.stdout.flush()
|
|
|
|
# Temperature history for graphing (~30 seconds at 20 updates/sec)
|
|
history = TemperatureHistory(max_samples=600)
|
|
|
|
start_time = time.time()
|
|
|
|
try:
|
|
# Start workers silently
|
|
if not args.npu_only:
|
|
for i in range(num_workers):
|
|
p = Process(target=cpu_stress_worker, args=(i, running))
|
|
p.start()
|
|
processes.append(p)
|
|
|
|
if not args.cpu_only and HAS_RKNN:
|
|
p = Process(target=npu_stress_worker, args=(args.model, running))
|
|
p.start()
|
|
processes.append(p)
|
|
|
|
# Main display loop
|
|
end_time = start_time + args.time if args.time > 0 else float('inf')
|
|
|
|
# Graph caching - only update graph every 0.2s (5fps) for performance
|
|
cached_graph = ""
|
|
last_graph_time = 0
|
|
graph_update_interval = 0.2
|
|
|
|
while running.value and time.time() < end_time:
|
|
temps = get_temperatures()
|
|
history.update(temps)
|
|
elapsed = time.time() - start_time
|
|
|
|
# Build entire output in memory first
|
|
lines = []
|
|
|
|
# Header
|
|
lines.append(f"{BOLD}{RED}{'═' * 60}{RESET}")
|
|
lines.append(f"{BOLD}{RED} 🔥 HEAT GENERATOR - STRESS TEST 🔥{RESET}")
|
|
lines.append(f"{BOLD}{RED}{'═' * 60}{RESET}")
|
|
lines.append("")
|
|
|
|
# System info
|
|
status_line = f" {CYAN}CPU Workers:{RESET} {num_workers if not args.npu_only else 0} "
|
|
status_line += f"{CYAN}NPU:{RESET} {'✅ Active' if (not args.cpu_only and npu_available) else '❌ Off'} "
|
|
status_line += f"{CYAN}Elapsed:{RESET} {elapsed:.1f}s"
|
|
if args.time > 0:
|
|
remaining = args.time - elapsed
|
|
status_line += f" {CYAN}Remaining:{RESET} {remaining:.1f}s"
|
|
lines.append(status_line)
|
|
lines.append("")
|
|
|
|
# Current temperatures table
|
|
lines.append(f" {'Sensor':<12} {'Current':>10} {'Start':>10} {'Min':>8} {'Max':>8}")
|
|
lines.append(f" {'-' * 52}")
|
|
|
|
for name, temp in temps.items():
|
|
friendly = name.replace('-thermal', '').replace('_thermal', '').upper()
|
|
icon = get_status_icon(temp)
|
|
color = get_temp_color(temp)
|
|
start_t = history.start_temps.get(name, temp)
|
|
min_t = history.min_temps.get(name, temp)
|
|
max_t = history.max_temps.get(name, temp)
|
|
|
|
lines.append(f" {icon} {friendly:<10} {color}{temp:>7.1f}°C{RESET} {start_t:>7.1f}°C {min_t:>6.1f}° {max_t:>6.1f}°")
|
|
|
|
lines.append(f" {'-' * 52}")
|
|
|
|
# Temperature graph - only regenerate every 0.2s for performance
|
|
current_time = time.time()
|
|
if not args.no_graph and len(history.time_points) > 2:
|
|
if current_time - last_graph_time >= graph_update_interval:
|
|
current_graph_height = get_current_graph_height()
|
|
cached_graph = history.get_graph(height=current_graph_height)
|
|
last_graph_time = current_time
|
|
|
|
graph_str = cached_graph
|
|
|
|
# Footer
|
|
lines.append("")
|
|
lines.append(f" Press {BOLD}Ctrl+C{RESET} to stop")
|
|
|
|
# Build output: header content first
|
|
header_content = "\n".join(lines[:len(lines)-2]) # Everything except last 2 lines
|
|
footer_content = "\n".join(lines[-2:]) # Last 2 lines
|
|
|
|
# Calculate total lines
|
|
header_lines = header_content.count('\n') + 1
|
|
graph_lines = graph_str.count('\n') + 1 if graph_str else 0
|
|
footer_lines = footer_content.count('\n') + 1
|
|
|
|
total_content_lines = header_lines + graph_lines + footer_lines
|
|
|
|
# Build complete output
|
|
if graph_str:
|
|
output = header_content + "\n" + graph_str + "\n" + footer_content
|
|
else:
|
|
output = header_content + "\n" + footer_content
|
|
|
|
# Pad to exactly fill terminal (add lines at the end)
|
|
current_lines = output.count('\n') + 1
|
|
if current_lines < term_rows:
|
|
output += "\n" * (term_rows - current_lines)
|
|
|
|
# Single atomic write: home + content + clear remainder
|
|
sys.stdout.write(CURSOR_HOME + output + CLEAR_TO_END)
|
|
sys.stdout.flush()
|
|
|
|
time.sleep(0.05) # 20 updates per second
|
|
|
|
running.value = False
|
|
|
|
finally:
|
|
running.value = False
|
|
|
|
# Exit alternate screen and restore cursor
|
|
sys.stdout.write(ALT_SCREEN_OFF + SHOW_CURSOR)
|
|
sys.stdout.flush()
|
|
|
|
# Wait for workers
|
|
for p in processes:
|
|
p.join(timeout=2)
|
|
if p.is_alive():
|
|
p.terminate()
|
|
|
|
# Final summary on clean screen
|
|
elapsed = time.time() - start_time
|
|
temps = get_temperatures()
|
|
|
|
print(f"{GREEN}{'═' * 60}{RESET}")
|
|
print(f"{GREEN} ✅ Heat generation stopped after {elapsed:.1f}s{RESET}")
|
|
print(f"{GREEN}{'═' * 60}{RESET}\n")
|
|
|
|
print(f" {'Sensor':<12} {'Final':>10} {'Start':>10} {'Peak':>10} {'Rise':>10}")
|
|
print(f" {'-' * 54}")
|
|
|
|
for name, temp in temps.items():
|
|
friendly = name.replace('-thermal', '').replace('_thermal', '').upper()
|
|
color = get_temp_color(temp)
|
|
start_t = history.start_temps.get(name, temp)
|
|
max_t = history.max_temps.get(name, temp)
|
|
rise = max_t - start_t
|
|
|
|
print(f" {friendly:<12} {color}{temp:>7.1f}°C{RESET} {start_t:>7.1f}°C {max_t:>7.1f}°C {'+' if rise >= 0 else ''}{rise:>6.1f}°C")
|
|
|
|
print()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|