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))

For the stress minutes glitch I replace the part where is says os.environ(‘STRESSMIN’)
with the number of minutes.

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)))