first
commit
490406e3cd
|
|
@ -0,0 +1,532 @@
|
|||
#!/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()
|
||||
Loading…
Reference in New Issue