Python script useful for stress testing systems

stress_test.py

"""
Produces load on all available CPU cores.
Requires system environment var STRESS_MINS to be set.
"""

from multiprocessing import Pool
from multiprocessing import cpu_count
import time
import os

def f(x):
    set_time = os.environ['STRESS_MINS']
    timeout = time.time() + 60*float(set_time)  # X minutes from now
    while True:
        if time.time() > timeout:
            break
        x*x

if __name__ == '__main__':
    processes = cpu_count()
    print ('utilizing %d cores\n' % processes)
    pool = Pool(processes)
    pool.map(f, range(processes))

Here's my modified version which uses command line arguments and has some basic throttling available.

As with the original one above, hitting Ctrl-C while it's running can lead to some zombie processes - let it exit normally or kill the parent python process directly if you want to exit early.

'''
Examples:

Run for 15 seconds maxing out all processors:
  stress.py 15

Run for 15 seconds with each subprocess sleeping for 0.01s every 100,000 cycles across all processors (on my machine it's about a 50% duty cycle):
  stress.py 15 0.01 100000

Run for 15 seconds, sleep 0.01s every 100,000 cycles, but only use a max of 8 processors:
  stress.py 15 0.01 100000 8
'''

import time
import sys
from itertools import repeat
from multiprocessing import Pool, cpu_count

def f(x, runtime=1, sleeptime=0, busycycles=100000):
    timeout = time.time() + runtime
    cnt = 0
    while True:
        if time.time() > timeout:
            break
        if sleeptime and cnt % busycycles == 0:
            time.sleep(sleeptime)
        x*x
        cnt += 1

if __name__ == '__main__':
    runtime = 5 if len(sys.argv) <= 1 else float(sys.argv[1])
    sleeptime = 0.0 if len(sys.argv) <= 2 else float(sys.argv[2])
    busycycles = 100000 if len(sys.argv) <= 3 else int(sys.argv[3])
    processes = cpu_count() if len(sys.argv) <= 4 else int(sys.argv[4]) if 0 < int(sys.argv[4]) <= cpu_count() else cpu_count()
    print(f'Running for {runtime}s with sleep time of {sleeptime}s per {busycycles} cycles utilizing {processes} cores')
    pool = Pool(processes)
    pool.starmap(f, zip(range(processes), repeat(runtime), repeat(sleeptime), repeat(busycycles)))