fix: Fix memory leak with dlpack when using python Tensor objects

[nnshah1]

Update:

The root cause - thanks to @tanmayv25 - was both the dynamic memory for the shape array and the reference to the owning tensor. While the reference count of the object was being incremented - it wasn't being passed correctly via the void_p and thus the DecrRef didn't release the object.

In order to increment the reference count and pass the object as a void pointer we follow a pattern outlined here:

  https://groups.google.com/g/dev-python/c/QRRqVC7gkf4/m/zH7l1gTXBwAJ

Also updated to create the shape array as part of a new object used for the management - to simply the deletion logic.

Updated the tests to include memory leak tests and tests for basic reference counting.

---

Initial investigation found that the mechanism for setting the shape field of the managed tensor was not being cleaned up correctly. This caused tensors to not be released even when the garbage collector was forced to run.

Once the shape was changed to be allocated via malloc and freed via malloc - the tests still require the garbage collector to be forced to run to remove tensors deterministically. That indicates there is a circular reference - but exact RCA is not known.

The Tensor object stores a reference to the original object (for example numpy).
When a numpy array is created from the Tensor object - a reference is created to the Tensor object (and implicitly the original numpy array).

When the objects are deleted or dereferenced - they should automatically dereference completely.

 for index in range(50):

    # Create original source in this case on GPU
    tensor = cupy.ones(2**27)

   # create a Triton Tensor object
   # is zero copy - contains reference to tensor object

    dl_pack_tensor = tritonserver.Tensor.from_dlpack(tensor)

    # create second array from Triton tensor
    # is zero copy - contains reference to Triton tensor
    array = cupy.from_dlpack(dl_pack_tensor)

    #            print(index, index*torch.numel(tensor)*tensor.element_size())
   
    # Delete array - also deletes capsule
    # should dereference Triton Tensor
    del array

    #  Delete Triton Tensor
    # should deference original tensor
    del dl_pack_tensor

    # Delete original storage
    # Should free actual memory
    del tensor

    print(index)

[kthui]

If I understand correctly, the issue is calling

ctypes.pythonapi.Py_DecRef(...)
ctypes.pythonapi.Py_IncRef(...)

on

dl_managed_tensor.dl_tensor.shape

object directly has no effect, but it will work if the shape variable is referencing another object, i.e.

manager_ctx.shape

and the Inc/Dec ref is called on the managed_ctx?

[nnshah1]

If I understand correctly, the issue is calling

ctypes.pythonapi.Py_DecRef(...)
ctypes.pythonapi.Py_IncRef(...)

on

dl_managed_tensor.dl_tensor.shape

object directly has no effect, but it will work if the shape variable is referencing another object, i.e.

manager_ctx.shape

and the Inc/Dec ref is called on the managed_ctx?

yes essentially - tensor.shape itself is a pointer to the array directly and not a python object. There are a couple ways to solve that - but in this case the simplest seemed to create a single object that held references for all required objects - and then incr / decr that as a whole.