def test_import_fail_windows() -> None:
"""
Expects
-------
* Should fail to import if the memory limit is too low
Note
----
For sklearn on windows, this triggers an OSError with code 22 and winerr 1455
and so we have to be aware of it before hand.
There's no automatic way to identfiy this from a regular OSError so we better
be explicit about it.
"""
m = Monitor()
print("Before job ", m.memory("mb"))
try:
with limit(
import_sklearn,
memory=(100, "mb"),
wrap_errors={"memory": [(OSError, 22, 1455)]},
) as lf:
lf()
except MemoryLimitException:
pass
except Exception as e:
print(e, type(e))
raise e
def import_sklearn() -> bool:
"""Just imports sklearn"""
m = Monitor()
print("Before import ", m.memory("mb"))
import sklearn # noqa
print("After import ", m.memory("mb"))
return True
def train_svr(model: Any, X: np.ndarray, y: np.ndarray, limit: Any) -> bool:
"""Trains an svr model"""
m = Monitor()
print("Memory before fit", m.memory("mb"))
print("Limit ", limit)
# This will cause a coredump with sigdev
with open(os.devnull, "rb") as f:
faulthandler.enable(file=f)
model.fit(X, y)
print("Memory after fit", m.memory("mb"))
return True
def test_import_fail_all() -> None:
"""
Expects
-------
* Should fail to import but give a PynisherException as we can't properly
identify that it's root cause is due to memory
"""
m = Monitor()
print("Before job ", m.memory("mb"))
try:
with limit(import_sklearn, wrap_errors=True, memory=(100, "mb")) as lf:
lf()
except MemoryLimitException as e:
raise e
except PynisherException:
pass
except Exception as e:
print(e, type(e))
raise e
def test_train_svr_memory(context: str) -> None:
"""
Expects
-------
* Should raise a MemoryError if it ran out of memory during fit
Note
----
What will actually happen is a segmentation fault, we deal with this
in `pynisher`
"""
m = Monitor()
from sklearn.datasets import make_regression
from sklearn.svm import SVR
model = SVR()
X, y = make_regression(n_samples=30_000, n_features=128)
# Seem fit will consume about 28mb extra, see __main__
# Add 1MB
too_little_mem = round(m.memory("MB") + 1)
lf = limit(train_svr, memory=(too_little_mem, "MB"))
completed = False
with pytest.raises(MemoryLimitException):
completed = lf(model, X, y, too_little_mem)
assert completed is not True
assert lf._process is not None and not lf._process.is_running()
try:
children = lf._process.children()
assert not any(children)
except psutil.NoSuchProcess:
pass
def test_success(limit: int, context: str) -> None:
"""Using less than the allocated memory should be fine
Processes take up some amount of memory natively, e.g. 37MB was preallocated
for my own run on Linx where "fork" is used.
Hence, we skip if the test limit is not enough
"""
allocate = int(limit / 3)
current_usage = Monitor().memory("MB")
expected_usage = current_usage + allocate
if expected_usage > limit:
msg = (
f"Limit {limit}MB is too low as the current usage is {current_usage}MB"
f" with another {allocate}MB being allocated. This will total "
f" {expected_usage}MB, over the limit.")
pytest.skip(msg)
rf = Pynisher(usememory, memory=(limit, "MB"), context=context)
rf((allocate, "MB"))
children = lf._process.children()
assert not any(children)
except psutil.NoSuchProcess:
pass
if __name__ == "__main__":
"""
For checking memory usage with the SVM
My machine:
Before model and data 689.390625
After model and data 836.6640625
After fit 864.01953125
Fitting duration in cputime 89.290726962
It seems it takes about 28mb extra to fit with (30_000, 128) data
"""
from sklearn.datasets import make_regression
from sklearn.svm import SVR
m = Monitor()
print("Before model and data", m.memory("mb"))
X, y = make_regression(n_samples=30_000, n_features=128)
model = SVR()
print("After model and data", m.memory("mb"))
model.fit(X, y)
print("After fit", m.memory("mb"))
def __call__(self, *args: Any, **kwargs: Any) -> None:
"""Set process limits and then call the function with the given arguments.
Note
----
This should usually not be called by the end user, it's main use is in
Pynisher::run().
Parameters
----------
*args, **kwargs
Arguments to the function
Returns
-------
None
"""
try:
if self.cpu_time is not None:
self.limit_cpu_time(self.cpu_time)
if self.memory is not None:
# We should probably warn if we exceed the memory usage before
# any limitation is set
memusage = Monitor().memory("B")
if self.memory <= memusage:
self._raise_warning(
f"Current memory usage in new process is {memusage}B but "
f" setting limit to {self.memory}B. Likely to fail, try"
f" increasing the memory limit"
)
self.limit_memory(self.memory)
# Call our function
result = self.func(*args, **kwargs)
error = None
tb = None
except Exception as e:
result = None
error = e
tb = "".join(traceback.format_exception(*sys.exc_info()))
if error is not None:
error = self._wrap_error(error, *args, **kwargs)
# Now let's try to send the result back
response = (result, error, tb)
try:
self.output.send(response)
self.output.close()
if self.terminate_child_processes is True:
terminate_process(timeout=2, children=True, parent=False)
return
except Exception as e:
failed_send_tb = "".join(traceback.format_exception(*sys.exc_info()))
failed_send_response = (None, e, failed_send_tb)
# We failed to send the result back, lets try to send the errors we got from
# doing so
try:
self.output.send(failed_send_response)
self.output.close()
if self.terminate_child_processes is True:
terminate_process(timeout=2, children=True, parent=False)
return
except Exception:
pass
# One last effot to send the smallest None through
last_response_attempt = None
try:
self.output.send(last_response_attempt)
self.output.close()
if self.terminate_child_processes is True:
terminate_process(timeout=2, children=True, parent=False)
return
except Exception:
pass
# Look, everything failed, try kill any child processes and just kill this
# process is we can't for memory reasons
self.output.close()
try:
if self.terminate_child_processes is True:
terminate_process(timeout=2, children=True, parent=False)
except MemoryError:
os.kill(os.getpid(), signal.SIGSEGV)
finally:
return
print(e, type(e))
raise e
@pytest.mark.skipif(not supports("memory"),
reason="System doesn't support memory")
def test_import_with_enough_memory() -> None:
"""
Expects
-------
* Should import without a problem given enough memory
"""
with limit(import_sklearn, memory=(1, "GB")) as lf:
assert lf() is True
if __name__ == "__main__":
"""
Use this to estimate memory limits. Output on my machine:
* Before import = 17.83984375
* After Import sklearn = 671.28515625
"""
m = Monitor()
print("Before import = ", m.memory("mb"))
import sklearn # noqa
print("After Import sklearn = ", m.memory("mb"))