#! /bin/python
import time
import multiprocessing
import unittest
import os
import sys
import signal
import logging
import pynisher
import psutil
try:
import sklearn
is_sklearn_available = True
except ImportError:
print("Scikit Learn was not found!")
is_sklearn_available = False
all_tests=1
logger = multiprocessing.log_to_stderr()
logger.setLevel(logging.WARNING)
# TODO: add tests with large return value to test for deadlock!
def rogue_subprocess():
pid = os.getpid()
oldgrp = os.getpgrp()
os.setpgrp()
logger.debug("{}: Changed group id from {} to {}".format(pid, oldgrp, os.getpgrp()))
time.sleep(60)
def spawn_rogue_subprocess(num_procs = 5):
for i in range(num_procs):
p = multiprocessing.Process(target=rogue_subprocess, daemon=False)
p.start()
p = psutil.Process()
time.sleep(10)
def simulate_work(size_in_mb, wall_time_in_s, num_processes):
# allocate memory (size_in_mb) with an array
# note the actual size in memory of this process is a little bit larger
A = [42.]*((1024*size_in_mb)//8);
# try to spawn new processes
if (num_processes > 0):
# data parallelism
multiprocessing.Pool(num_processes)
# sleep for specified duration
time.sleep(wall_time_in_s+1)
return(size_in_mb, wall_time_in_s, num_processes);
def svm_example(n_samples = 10000, n_features = 100):
from sklearn.svm import SVR
from sklearn.datasets import make_regression
X,Y = make_regression(n_samples, n_features)
m = SVR()
m.fit(X,Y)
def svc_example(n_samples = 10000, n_features = 4):
from sklearn.svm import LinearSVC
from sklearn.preprocessing import PolynomialFeatures
from sklearn.datasets import make_classification
X,Y = make_classification(n_samples, n_features)
#pp = PolynomialFeatures(degree=3)
#X = pp.fit_transform(X)
m = LinearSVC()
m.fit(X,Y)
def crash_unexpectedly(signum):
print("going to receive signal {}.".format(signum))
pid = os.getpid()
time.sleep(1)
os.kill(pid, signum)
time.sleep(1)
def return_big_array(num_elements):
return([1]*num_elements)
def cpu_usage():
i = 1
while True:
i += 1
def nested_pynisher(level=2, cputime=5, walltime=5, memlimit = 10e24, increment = -1, grace_period = 1):
print("this is level {}".format(level))
if level == 0:
spawn_rogue_subprocess(10)
else:
func = pynisher.enforce_limits(mem_in_mb=memlimit, cpu_time_in_s=cputime, wall_time_in_s=walltime, grace_period_in_s = grace_period)(nested_pynisher)
func(level-1, None, walltime+increment, memlimit, increment)
class test_limit_resources_module(unittest.TestCase):
@unittest.skipIf(not all_tests, "skipping successful tests")
def test_success(self):
print("Testing unbounded function call which have to run through!")
local_mem_in_mb = None
local_wall_time_in_s = None
local_cpu_time_in_s = None
local_grace_period = None
wrapped_function = pynisher.enforce_limits(mem_in_mb = local_mem_in_mb, wall_time_in_s=local_wall_time_in_s, cpu_time_in_s = local_cpu_time_in_s, grace_period_in_s = local_grace_period)(simulate_work)
for mem in [1,2,4,8,16]:
self.assertEqual((mem,0,0),wrapped_function(mem,0,0))
self.assertEqual(wrapped_function.exit_status, 0)
@unittest.skipIf(not all_tests, "skipping out_of_memory test")
def test_out_of_memory(self):
print("Testing memory constraint.")
local_mem_in_mb = 32
local_wall_time_in_s = None
local_cpu_time_in_s = None
local_grace_period = None
wrapped_function = pynisher.enforce_limits(mem_in_mb = local_mem_in_mb, wall_time_in_s=local_wall_time_in_s, cpu_time_in_s = local_cpu_time_in_s, grace_period_in_s = local_grace_period)(simulate_work)
for mem in [1024, 2048, 4096]:
self.assertIsNone(wrapped_function(mem,0,0))
self.assertEqual(wrapped_function.exit_status, pynisher.MemorylimitException)
@unittest.skipIf(not all_tests, "skipping time_out test")
def test_time_out(self):
print("Testing wall clock time constraint.")
local_mem_in_mb = None
local_wall_time_in_s = 1
local_cpu_time_in_s = None
local_grace_period = None
wrapped_function = pynisher.enforce_limits(mem_in_mb = local_mem_in_mb, wall_time_in_s=local_wall_time_in_s, cpu_time_in_s = local_cpu_time_in_s, grace_period_in_s = local_grace_period)(simulate_work)
for mem in range(1,10):
self.assertIsNone(wrapped_function(mem,10,0))
self.assertEqual(wrapped_function.exit_status, pynisher.TimeoutException)
@unittest.skipIf(not all_tests, "skipping too many processes test")
def test_num_processes(self):
print("Testing number of processes constraint.")
local_mem_in_mb = None
local_num_processes = 1
local_wall_time_in_s = None
local_grace_period = None
wrapped_function = pynisher.enforce_limits(mem_in_mb = local_mem_in_mb, wall_time_in_s=local_wall_time_in_s,num_processes = local_num_processes, grace_period_in_s = local_grace_period)(simulate_work)
for processes in [2,15,50,100,250]:
self.assertIsNone(wrapped_function(0,0, processes))
self.assertEqual(wrapped_function.exit_status, pynisher.SubprocessException)
@unittest.skipIf(not all_tests, "skipping unexpected signal test")
def test_crash_unexpectedly(self):
print("Testing an unexpected signal simulating a crash.")
wrapped_function = pynisher.enforce_limits()(crash_unexpectedly)
self.assertIsNone(wrapped_function(signal.SIGQUIT))
self.assertEqual(wrapped_function.exit_status, pynisher.AnythingException)
@unittest.skipIf(not all_tests, "skipping unexpected signal test")
def test_high_cpu_percentage(self):
print("Testing cpu time constraint.")
cpu_time_in_s = 2
grace_period = 1
wrapped_function = pynisher.enforce_limits(cpu_time_in_s = cpu_time_in_s, grace_period_in_s = grace_period)(cpu_usage)
self.assertEqual(None,wrapped_function())
self.assertEqual(wrapped_function.exit_status, pynisher.CpuTimeoutException)
@unittest.skipIf(not all_tests, "skipping big data test")
def test_big_return_data(self):
print("Testing big return values")
wrapped_function = pynisher.enforce_limits()(return_big_array)
for num_elements in [4,16,64, 256, 1024, 4096, 16384, 65536, 262144]:
bla = wrapped_function(num_elements)
self.assertEqual(len(bla), num_elements)
@unittest.skipIf(not all_tests, "skipping subprocess changing process group")
def test_kill_subprocesses(self):
wrapped_function = pynisher.enforce_limits(wall_time_in_s = 1)(spawn_rogue_subprocess)
wrapped_function(5)
time.sleep(1)
p = psutil.Process()
self.assertEqual(len(p.children(recursive=True)), 0)
@unittest.skipIf(not is_sklearn_available, "test requires scikit learn")
@unittest.skipIf(not all_tests, "skipping fitting an SVM to see how C libraries are handles")
def test_busy_in_C_library(self):
global logger
wrapped_function = pynisher.enforce_limits(wall_time_in_s = 2)(svm_example)
start = time.time()
wrapped_function(16384, 128)
duration = time.time()-start
time.sleep(1)
p = psutil.Process()
self.assertEqual(len(p.children(recursive=True)), 0)
self.assertTrue(duration < 2.1)
@unittest.skipIf(not is_sklearn_available, "test requires scikit learn")
@unittest.skipIf(not all_tests, "skipping fitting an SVM to see how C libraries are handles")
def test_liblinear_svc(self):
global logger
time_limit = 2
grace_period = 1
wrapped_function = pynisher.enforce_limits(cpu_time_in_s = time_limit, mem_in_mb=None, grace_period_in_s=grace_period, logger=logger)(svc_example)
start = time.time()
wrapped_function(16384, 1000)
duration = time.time()-start
time.sleep(1)
p = psutil.Process()
self.assertEqual(len(p.children(recursive=True)), 0)
# fails with pynisher.AnythingException for some reason
#self.assertTrue(wrapped_function.exit_status == pynisher.CpuTimeoutException)
self.assertTrue(duration > time_limit-0.1)
self.assertTrue(duration < time_limit+grace_period+0.1)
@unittest.skipIf(not all_tests, "skipping nested pynisher test")
def test_nesting(self):
tl = 2 #time limit
gp = 1 #grace period
start = time.time()
nested_pynisher(level=2, cputime = 2, walltime = 2, memlimit = None, increment = 1, grace_period = gp)
duration = time.time()-start
print(duration)
time.sleep(1)
p = psutil.Process()
self.assertEqual(len(p.children(recursive=True)), 0)
self.assertTrue(duration > tl-0.1)
self.assertTrue(duration < tl+gp+0.1)
@unittest.skipIf(not all_tests, "skipping capture stdout test")
def test_capture_output(self):
print("Testing capturing of output.")
global logger
time_limit = 2
grace_period = 1
def print_and_sleep(t):
for i in range(t):
print(i)
time.sleep(1)
wrapped_function = pynisher.enforce_limits(wall_time_in_s = time_limit, mem_in_mb=None, grace_period_in_s=grace_period, logger=logger, capture_output=True)(print_and_sleep)
wrapped_function(5)
self.assertTrue('0' in wrapped_function.stdout)
self.assertTrue(wrapped_function.stderr == '')
def print_and_fail():
print(0)
raise RuntimeError()
wrapped_function = pynisher.enforce_limits(wall_time_in_s = time_limit, mem_in_mb=None, grace_period_in_s=grace_period, logger=logger, capture_output=True)(print_and_fail)
wrapped_function()
self.assertTrue('0' in wrapped_function.stdout)
self.assertTrue('RuntimeError' in wrapped_function.stderr)
unittest.main()
