|
8 | 8 | """ |
9 | 9 |
|
10 | 10 | import ctypes |
| 11 | +import gc |
| 12 | +import weakref |
11 | 13 | from pathlib import Path |
12 | 14 |
|
13 | 15 | try: |
@@ -105,6 +107,41 @@ def test_dlpack_view_of_buffer_reuses_exporting_buffer(init_cuda): |
105 | 107 | buffer.close() |
106 | 108 |
|
107 | 109 |
|
| 110 | +class _EphemeralDLPackExporter: |
| 111 | + """CPU exporter whose capsule is the only owner of a freshly created array.""" |
| 112 | + |
| 113 | + def __init__(self): |
| 114 | + self.array_ref = None |
| 115 | + |
| 116 | + def __dlpack_device__(self): |
| 117 | + return (int(DLDeviceType.kDLCPU), 0) |
| 118 | + |
| 119 | + def __dlpack__(self, **kwargs): |
| 120 | + array = np.arange(4, dtype=np.int32) |
| 121 | + self.array_ref = weakref.ref(array) |
| 122 | + return array.__dlpack__(**kwargs) |
| 123 | + |
| 124 | + |
| 125 | +@pytest.mark.agent_authored(model="claude-opus-5-5") |
| 126 | +def test_dlpack_derived_view_keeps_capsule_owned_storage_alive(init_cuda): |
| 127 | + """A derived view keeps capsule-owned storage alive until the last view is released.""" |
| 128 | + exporter = _EphemeralDLPackExporter() |
| 129 | + original = StridedMemoryView.from_dlpack(exporter, stream_ptr=-1) |
| 130 | + derived = original.view(dtype=original.dtype) |
| 131 | + ptr = original.ptr |
| 132 | + del original |
| 133 | + gc.collect() |
| 134 | + |
| 135 | + # Check liveness before reading, so a regression fails without touching freed memory. |
| 136 | + assert exporter.array_ref() is not None |
| 137 | + assert derived.ptr == ptr |
| 138 | + assert np.from_dlpack(derived).tolist() == [0, 1, 2, 3] |
| 139 | + |
| 140 | + del derived |
| 141 | + # No gc.collect(): the capsule must be released by reference counting, not a cycle. |
| 142 | + assert exporter.array_ref() is None |
| 143 | + |
| 144 | + |
108 | 145 | def test_dlpack_export_unversioned_capsule_and_deleter(): |
109 | 146 | """``__dlpack__()`` with no ``max_version`` yields an *unversioned* unused |
110 | 147 | DLPack capsule; dropping it unconsumed runs ``_smv_pycapsule_deleter`` on |
|
0 commit comments