Python API to allocator

[jakirkham]

Would be great to have a Python API for configuring NumPy's memory allocator

[seberg]

@jakirkham the C-API allows you to write an allocator that exposes a Python interface that has a with statement. There is even an example implementation (sorry, gotta scroll through the PR to find it) which creates one using ctypes.
So I am not sure what there would be for NumPy to do, unless you are suggesting to provide specific, alternative allocators in NumPy itself?

[mattip]

Here is the sample repo https://github.com/inaccel/numpy-allocator

[leofang]

Sorry guys, but isn't the ctype thing a hack/proof of concept/prototype/whatever that, as @seberg has emphasized in various occasions, "should be done outside of NumPy until considered valuable/feasible to merge into the main project"? (Sorry, this may not be the precise quote, but I am pretty sure I heard this more than once from @seberg.) I consider @jakirkham's proposal to kick off this discussion. C'mon, we can do better than ctypes.CDLL, can't we?

Specifically, I think that -- John correct me if I am wrong -- what's being asked for is merely an interface for hooking up a user-provided allocator, instead of an alternative allocator provided by NumPy. Why would NumPy need to implement two allocators...?

[seberg]

what's being asked for is merely an interface for hooking up a user-provided allocator

But what is that API supposed to look like (or make easy)? The repo has a full implementation for a context managers, which allows API like:

with my_allocator():
   do_numpy_stuff()

Now, I guess the whole context manager stuff may be enough trouble to write correctly, that we want need to make it easier write such allocator?
Or maybe I am missing something and just stealing that context manager (or writing one in cython) won't work?

I don't immediately like the thought that my_allocator is some capsule and we add:

with np.use_allocator(my_allocator):
    do_numpy_stuff()

The first seems neater to me, but maybe the function has its merit because it is clearer that it affects NumPy?
The first is always possible of course, so we could just add the second, I just am not sure if exposing those bare capsules to users is nice to begin with.

[jakirkham]

Yeah how this API works from the Python side does deserve a fair bit of thought. It is possible it is a similar API in Python, but not the exact same as the C API.

For example the concept of Python code working with raw pointers (void*) is a bit unsettling. Though maybe memory could be wrapped in memoryviews for Python. This would be a bit safer and have RAII semantics. Additionally ctx can be added to a class (extension type?) with methods (malloc, calloc, etc.). Third-parties could then inherit from this class to make their own or maybe use a factory function to create one quickly.

Being able to set the allocator locally (with a context manager) or globally is definitely useful.

Other useful features from an end user perspective are inspecting memory usage by the allocator, freeing unused memory, etc.

Maybe for more inspiration we could look at how CuPy or Numba handle similar things with device memory.

[seberg]

I am still confused how this would look like and be used. I.e. I would have thought that:

• Very few people will want to write an allocator (so there is no need for super-easy extending by inheritance maybe).
• Those who extend will typically want very low overhead, so need to use the C-API.

So, I am surprised that we would want to channel array allocations through Python.

Looking at the cupy API, it looks a bit like it stores the "pointer" object which owns the memory? The current API in NumPy did not even anticipate such use unfortunately, so I think that would require extension even.
You could channel a pointer through Python of course, but that would come with no ownership model/reference counting.

Other useful features from an end user perspective are inspecting memory usage by the allocator, freeing unused memory, etc.

Yes, but it seems to me that those can be achieved by creating an allocator that provides them? I.e. it requires a Python API for the end-user, but not necessarily for the allocator author?

[seberg]

To be honest, I am starting to have some things that I would like to improve here. There are two lines of thought:

1. PyMemHandler_Set() seems like an unfortunate API. We are using a context vars but numpy-allocator has some weird additional internal context var. That doesn't smell right. The way things should probably be is embracing the context var as public API:

   • PyMemHandler_Set() would IMO ideally return the contextvar token (i.e. opaque object). Probably we can't change the return value, so that needs a new function.
   • PyMemHandler_Restore(token)
   • EDIT: While tokens are nice and safer, they don't actually matter.

2. In some thoughts, I had wishes for metadata. But that metadata (and knowledge about which allocator) should also be available on views!

   • We currently have to walk all the .base attribute, which doesn't seem ideal even if it usually/mostly works in practice.
   • If you would tag the allocator onto the view, you would still not solve the problem of accessing the metadata, since you still need the original allocation pointer to find it!
     Walking the base is at least mostly workable, but it does make me wonder if the choice that metadata can always be put into the original allocation is actually great.

Both of these can probably be addressed. I am not quite sure how to do metadata best (it begs the question if the allocator should have always been slightly fatter). OTOH, I can imagine we can tag it by expanding the API a little (although no thought on how).