8000 BUG: Factor out slow `getenv` call used for memory policy warning · numpy/numpy@1d0487b · GitHub
[go: up one dir, main page]
Skip to content

Commit 1d0487b

Browse files
sebergseberg
committed
BUG: Factor out slow getenv call used for memory policy warning
Using `getenv` regularly is probably not great anyway, but it seems very slow on windows which leads to a large overhead for every array deallocation here. Refactor it out to only check on first import and add helper because the tests are set up slightly differently. (Manually checked that the startup works, tests run with policy set to 1, not modifying it and passing.)
1 parent 0aaa5d3 commit 1d0487b

File tree

5 files changed

+49
-13
lines changed

5 files changed

+49
-13
lines changed

‎doc/source/reference/global_state.rst

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -82,3 +82,6 @@ memory allocation policy, the default will be to call ``free``. If
8282
``NUMPY_WARN_IF_NO_MEM_POLICY`` is set to ``"1"``, a ``RuntimeWarning`` will
8383
be emitted. A better alternative is to use a ``PyCapsule`` with a deallocator
8484
and set the ``ndarray.base``.
85+
86+
.. versionchanged:: 1.25.2
87+
This variable is only checked on the first import.

‎numpy/core/src/multiarray/arrayobject.c

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,9 @@ maintainer email: oliphant.travis@ieee.org
6262
#include "binop_override.h"
6363
#include "array_coercion.h"
6464

65+
66+
NPY_NO_EXPORT npy_bool numpy_warn_if_no_mem_policy = 0;
67+
6568
/*NUMPY_API
6669
Compute the size of an array (in number of items)
6770
*/
@@ -460,9 +463,8 @@ array_dealloc(PyArrayObject *self)
460463
}
461464
}
462465
if (fa->mem_handler == NULL) {
463-
char *env = getenv("NUMPY_WARN_IF_NO_MEM_POLICY");
464-
if ((env != NULL) && (strncmp(env, "1", 1) == 0)) {
465-
char const * msg = "Trying to dealloc data, but a memory policy "
466+
if (numpy_warn_if_no_mem_policy) {
467+
char const *msg = "Trying to dealloc data, but a memory policy "
466468
"is not set. If you take ownership of the data, you must "
467469
"set a base owning the data (e.g. a PyCapsule).";
468470
WARN_IN_DEALLOC(PyExc_RuntimeWarning, msg);

‎numpy/core/src/multiarray/arrayobject.h

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
#ifndef NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_
66
#define NUMPY_CORE_SRC_MULTIARRAY_ARRAYOBJECT_H_
77

8+
extern NPY_NO_EXPORT npy_bool numpy_warn_if_no_mem_policy;
9+
810
NPY_NO_EXPORT PyObject *
911
_strings_richcompare(PyArrayObject *self, PyArrayObject *other, int cmp_op,
1012
int rstrip);

‎numpy/core/src/multiarray/multiarraymodule.c

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4378,6 +4378,24 @@ normalize_axis_index(PyObject *NPY_UNUSED(self),
43784378
}
43794379

43804380

4381+
static PyObject *
4382+
_set_numpy_warn_if_no_mem_policy(PyObject *NPY_UNUSED(self), PyObject *arg)
4383+
{
4384+
int res = PyObject_IsTrue(arg);
4385+
if (res < 0) {
4386+
return NULL;
4387+
}
4388+
int old_value = numpy_warn_if_no_mem_policy;
4389+
numpy_warn_if_no_mem_policy = res;
4390+
if (old_value) {
4391+
Py_RETURN_TRUE;
4392+
}
4393+
else {
4394+
Py_RETURN_FALSE;
4395+
}
4396+
}
4397+
4398+
43814399
static PyObject *
43824400
_reload_guard(PyObject *NPY_UNUSED(self), PyObject *NPY_UNUSED(args)) {
43834401
static int initialized = 0;
@@ -4625,6 +4643,9 @@ static struct PyMethodDef array_module_methods[] = {
46254643
METH_O, "Set the NEP 50 promotion state. This is not thread-safe.\n"
46264644
"The optional warnings can be safely silenced using the \n"
46274645
"`np._no_nep50_warning()` context manager."},
4646+
{"_set_numpy_warn_if_no_mem_policy",
4647+
(PyCFunction)_set_numpy_warn_if_no_mem_policy,
4648+
METH_O, "Change the warn if no mem policy flag for testing."},
46284649
{"_add_newdoc_ufunc", (PyCFunction)add_newdoc_ufunc,
46294650
METH_VARARGS, NULL},
46304651
{"_get_sfloat_dtype",
@@ -4913,6 +4934,14 @@ initialize_static_globals(void)
49134934
return -1;
49144935
}
49154936

4937+
char *env = getenv("NUMPY_WARN_IF_NO_MEM_POLICY");
4938+
if ((env != NULL) && (strncmp(env, "1", 1) == 0)) {
4939+
numpy_warn_if_no_mem_policy = 1;
4940+
}
4941+
else {
4942+
numpy_warn_if_no_mem_policy = 0;
4943+
}
4944+
49164945
return 0;
49174946
}
49184947

‎numpy/core/tests/test_mem_policy.py

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -409,16 +409,19 @@ def test_switch_owner(get_module, policy):
409409
a = get_module.get_array()
410410
assert np.core.multiarray.get_handler_name(a) is None
411411
get_module.set_own(a)
412-
oldval = os.environ.get('NUMPY_WARN_IF_NO_MEM_POLICY', None)
412+
413413
if policy is None:
414-
if 'NUMPY_WARN_IF_NO_MEM_POLICY' in os.environ:
415-
os.environ.pop('NUMPY_WARN_IF_NO_MEM_POLICY')
414+
# See what we expect to be set based on the env variable
415+
policy = os.getenv("NUMPY_WARN_IF_NO_MEM_POLICY", "0") == "1"
416+
oldval = None
416417
else:
417-
os.environ['NUMPY_WARN_IF_NO_MEM_POLICY'] = policy
418+
policy = policy == "1"
419+
oldval = np.core._multiarray_umath._set_numpy_warn_if_no_mem_policy(
420+
policy)
418421
try:
419422
# The policy should be NULL, so we have to assume we can call
420423
# "free". A warning is given if the policy == "1"
421-
if policy == "1":
424+
if policy:
422425
with assert_warns(RuntimeWarning) as w:
423426
del a
424427
gc.collect()
@@ -427,11 +430,8 @@ def test_switch_owner(get_module, policy):
427430
gc.collect()
428431

429432
finally:
430-
if oldval is None:
431-
if 'NUMPY_WARN_IF_NO_MEM_POLICY' in os.environ:
432-
os.environ.pop('NUMPY_WARN_IF_NO_MEM_POLICY')
433-
else:
434-
os.environ['NUMPY_WARN_IF_NO_MEM_POLICY'] = oldval
433+
if oldval is not None:
434+
np.core._multiarray_umath._set_numpy_warn_if_no_mem_policy(oldval)
435435

436436

437437
@pytest.mark.skipif(sys.version_info >= (3, 12), reason="no numpy.distutils")

0 commit comments

Comments
 (0)
0