diff options
| -rw-r--r-- | numpy/conftest.py | 18 | ||||
| -rw-r--r-- | numpy/lib/tests/test_io.py | 7 | ||||
| -rw-r--r-- | numpy/testing/_private/decorators.py | 4 | ||||
| -rw-r--r-- | numpy/testing/_private/utils.py | 70 |
4 files changed, 94 insertions, 5 deletions
diff --git a/numpy/conftest.py b/numpy/conftest.py index 18d5d1ce9..1baf4adda 100644 --- a/numpy/conftest.py +++ b/numpy/conftest.py @@ -3,6 +3,8 @@ Pytest configuration and fixtures for the Numpy test suite. """ from __future__ import division, absolute_import, print_function +import os + import pytest import numpy @@ -22,6 +24,22 @@ def pytest_configure(config): "slow: Tests that are very slow.") +def pytest_addoption(parser): + parser.addoption("--available-memory", action="store", default=None, + help=("Set amount of memory available for running the " + "test suite. This can result to tests requiring " + "especially large amounts of memory to be skipped. " + "Equivalent to setting environment variable " + "NPY_AVAILABLE_MEM. Default: determined" + "automatically.")) + + +def pytest_sessionstart(session): + available_mem = session.config.getoption('available_memory') + if available_mem is not None: + os.environ['NPY_AVAILABLE_MEM'] = available_mem + + #FIXME when yield tests are gone. @pytest.hookimpl() def pytest_itemcollected(item): diff --git a/numpy/lib/tests/test_io.py b/numpy/lib/tests/test_io.py index 6e2291ca3..66d6a5b00 100644 --- a/numpy/lib/tests/test_io.py +++ b/numpy/lib/tests/test_io.py @@ -24,6 +24,7 @@ from numpy.testing import ( assert_allclose, assert_array_equal, temppath, tempdir, IS_PYPY, HAS_REFCOUNT, suppress_warnings, assert_no_gc_cycles, assert_no_warnings ) +from numpy.testing._private.utils import requires_memory class TextIO(BytesIO): @@ -575,13 +576,9 @@ class TestSaveTxt(object): @pytest.mark.skipif(sys.platform=='win32', reason="large files cause problems") @pytest.mark.slow + @requires_memory(7e9) def test_large_zip(self): # The test takes at least 6GB of memory, writes a file larger than 4GB - try: - a = 'a' * 6 * 1024 * 1024 * 1024 - del a - except (MemoryError, OverflowError): - pytest.skip("Cannot allocate enough memory for test") test_data = np.asarray([np.random.rand(np.random.randint(50,100),4) for i in range(800000)]) with tempdir() as tmpdir: diff --git a/numpy/testing/_private/decorators.py b/numpy/testing/_private/decorators.py index 24c4e385d..eab40e7c9 100644 --- a/numpy/testing/_private/decorators.py +++ b/numpy/testing/_private/decorators.py @@ -15,6 +15,10 @@ function name, setup and teardown functions and so on - see """ from __future__ import division, absolute_import, print_function +import sys +import os +import re + try: # Accessing collections abstract classes from collections # has been deprecated since Python 3.3 diff --git a/numpy/testing/_private/utils.py b/numpy/testing/_private/utils.py index b14c776d9..409ed142f 100644 --- a/numpy/testing/_private/utils.py +++ b/numpy/testing/_private/utils.py @@ -2380,3 +2380,73 @@ def break_cycles(): gc.collect() # one more, just to make sure gc.collect() + + +def requires_memory(free_bytes): + """Decorator to skip a test if not enough memory is available""" + import pytest + + env_var = 'NPY_AVAILABLE_MEM' + env_value = os.environ.get(env_var) + if env_value is not None: + try: + mem_free = _parse_size(env_value) + except ValueError as exc: + raise ValueError('Invalid environment variable {}: {!s}'.format( + env_var, exc)) + + msg = ('{0} GB memory required, but environment variable ' + 'NPY_AVAILABLE_MEM={1} set'.format( + free_bytes/1e9, env_value)) + else: + mem_free = _get_mem_available() + + if mem_free is None: + msg = ("Could not determine available memory; set NPY_AVAILABLE_MEM " + "environment variable (e.g. NPY_AVAILABLE_MEM=16GB) to run " + "the test.") + mem_free = -1 + else: + msg = '{0} GB memory required, but {1} GB available'.format( + free_bytes/1e9, mem_free/1e9) + + return pytest.mark.skipif(mem_free < free_bytes, reason=msg) + + +def _parse_size(size_str): + """Convert memory size strings ('12 GB' etc.) to float""" + suffixes = {'': 1.0, 'b': 1.0, + 'k': 1e3, 'm': 1e6, 'g': 1e9, 't': 1e12, + 'kb': 1e3, 'mb': 1e6, 'gb': 1e9, 'tb': 1e12} + + size_re = re.compile(r'^\s*(\d+|\d+\.\d+)\s*({0})\s*$'.format( + '|'.join(suffixes.keys())), re.I) + + m = size_re.match(size_str.lower()) + if not m or m.group(2) not in suffixes: + raise ValueError("value {!r} not a valid size".format(size_str)) + return float(m.group(1)) * suffixes[m.group(2)] + + +def _get_mem_available(): + """Return available memory in bytes, or None if unknown.""" + try: + import psutil + return psutil.virtual_memory().available + except (ImportError, AttributeError): + pass + + if sys.platform.startswith('linux'): + info = {} + with open('/proc/meminfo', 'r') as f: + for line in f: + p = line.split() + info[p[0].strip(':').lower()] = float(p[1]) * 1e3 + + if 'memavailable' in info: + # Linux >= 3.14 + return info['memavailable'] + else: + return info['memfree'] + info['cached'] + + return None |