使用Python对CPU进行压测

背景：测试ECS实例弹性伸缩功能，使用单一业务指标达不到压测效果，使用单独脚本对CPU进行压力测试

基于centos自带的python2版本进行压力测试

#!/usr/bin/env python2
# -*- coding: utf-8 -*-

import threading
import multiprocessing
import time
import signal
import sys
import argparse

def cpu_bound_task(memory_consumption_mb=100):
    """执行一个更复杂的无限循环的计算任务，并分配内存"""
    # 分配大块内存
    memory_block = bytearray(memory_consumption_mb * 1024 * 1024)  # 默认分配100MB内存

    while True:
        # 更复杂的数学运算以消耗更多 CPU 资源
        x = 0
        for i in range(1000000):
            x += (i ** 2) * (i ** 2)  # 增加计算复杂度

        # 确保内存被使用
        memory_block[0] = (memory_block[0] + 1) % 256

def start_stress_test(num_workers=None, use_multiprocessing=False, memory_per_worker_mb=100):
    """启动指定数量的线程或进程来执行 CPU 密集型任务并消耗内存"""
    if num_workers is None:
        num_workers = multiprocessing.cpu_count()  # 默认使用所有可用的核心数

    workers = []
    print "Starting stress test with %d workers..." % num_workers

    if use_multiprocessing:
        # 使用多进程
        PoolClass = multiprocessing.Pool
        pool = PoolClass(processes=num_workers, initializer=lambda: signal.signal(signal.SIGINT, signal.SIG_IGN))
        for _ in range(num_workers):
            pool.apply_async(cpu_bound_task, args=(memory_per_worker_mb,))
        workers.append(pool)
    else:
        # 使用多线程
        for _ in range(num_workers):
            t = threading.Thread(target=cpu_bound_task, args=(memory_per_worker_mb,))
            t.daemon = True  # 设置为守护线程/进程，主线程结束时自动退出
            t.start()
            workers.append(t)

    try:
        while True:
            time.sleep(1)  # 主线程保持运行状态
    except KeyboardInterrupt:
        print "\nStopping stress test..."
        if use_multiprocessing:
            for pool in workers:
                pool.terminate()
        sys.exit(0)

if __name__ == "__main__":
    # 捕获 Ctrl+C 信号以优雅地退出
    signal.signal(signal.SIGINT, signal.default_int_handler)

    # 参数解析
    parser = argparse.ArgumentParser(description="CPU and Memory Stress Test")
    parser.add_argument('--workers', type=int, default=None, help='Number of workers to use')
    parser.add_argument('--multiprocessing', action='store_true', help='Use multiprocessing instead of threading')
    parser.add_argument('--memory-per-worker', type=int, default=100, help='Memory consumption per worker in MB')

    args = parser.parse_args()

    # 启动 CPU 和内存压力测试，默认使用所有核心和100MB内存/worker
    start_stress_test(num_workers=args.workers, use_multiprocessing=args.multiprocessing, memory_per_worker_mb=args.memory_per_worker)