Memory & CPU Stress Test in Python3

stress_test.py

#!/usr/bin/env python
# MIT License
# Copyright (c) 2019 Lucian Maly, <[email protected]>

# Tested on Python v3.7.3 (default, Jun 14 2019, 13:46:58) and Python v2.7.5 (default, Mar 26 2019, 22:13:06)
# Usage: stress_test.py <AMOUNT_OF_MEMORY_YOU_WANT_TO_CONSUME_IN_GB>
# Example: /usr/local/bin/python3.7 stress_test.py 10
#    -------------------------
#    Filling up 10GB of memory...
#    Running load on CPUs...
#    Utilizing 8 cores
#    -------------------------
#    Traceback (most recent call last):
#    File "stress3.py", line 41, in <module>
#        pool = Pool(processes)
#    File "/usr/local/lib/python3.7/multiprocessing/context.py", line 119, in Pool
#        context=self.get_context())
#    File "/usr/local/lib/python3.7/multiprocessing/pool.py", line 176, in __init__
#        self._repopulate_pool()
#    File "/usr/local/lib/python3.7/multiprocessing/pool.py", line 241, in _repopulate_pool
#        w.start()
#    File "/usr/local/lib/python3.7/multiprocessing/process.py", line 112, in start
#        self._popen = self._Popen(self)
#    File "/usr/local/lib/python3.7/multiprocessing/context.py", line 277, in _Popen
#        return Popen(process_obj)
#    File "/usr/local/lib/python3.7/multiprocessing/popen_fork.py", line 20, in __init__
#        self._launch(process_obj)
#    File "/usr/local/lib/python3.7/multiprocessing/popen_fork.py", line 70, in _launch
#        self.pid = os.fork()
#    OSError: [Errno 12] Cannot allocate memory

"""
Produces load on all available CPUs and fills up memory with X amount of GBs.
"""

# Imports
from multiprocessing import Pool
from multiprocessing import cpu_count
import signal
import sys

# Variables
stop_loop = 0
GB = 1024*1024*1024
gib = int(sys.argv[1])

def exit_chld(x,y):
    global stop_loop
    stop_loop = 1

def f(x):
    global stop_loop
    while not stop_loop:
        x*x

signal.signal(signal.SIGINT, exit_chld)

if __name__ == '__main__':
    processes = cpu_count()
    print('-' * 25)
    print('Filling up ' + str(gib) + 'GB of memory...')
    a = "1" * (gib * GB)
    print('Running load on CPUs...')
    print('Utilizing %d cores' % processes)
    print('-' * 25)
    pool = Pool(processes)
    pool.map(f, range(processes))