[runtime][python] Add debug sink to bindings #19013
Conversation
We don't support custom debug sinks in the Runtime Python bindings. In particular the ability to register a custom callback when tracing tensors. This change makes it possible to create a HAL module with a Python function as a callback. A special care is taken to avoid leaking due to circular references where the callback reference the HAL module itself. Signed-off-by: Boian Petkantchin <[email protected]>
|
This PR plumbs through the Python bindings the trace callback recently added by #18966. |
Signed-off-by: Boian Petkantchin <[email protected]>
sogartar
force-pushed
the
hal-module-sink-python-bindings
branch
from
8b2ec54
to
c060bb2
Compare
sogartar
requested review from
benvanik and
stellaraccident
and removed request for
stellaraccident
Signed-off-by: Boian Petkantchin <[email protected]>
|
In its current state, I found out that one needs to hold the HAL module python object in order for the debug sink context to not get deleted, which is not what I wanted. I wanted once registered in IREE, the HAL module to continue be a shared owner of the debug sink. I guess the C++ objects for the bindings gets destroyed and created on demand from the underlying IREE runtime objects. |
| // Retain a reference. We want the callback to be valid after | ||
| // the user has dropped its reference and not burden the user | ||
| // with lifetime management. | ||
| vm_module.SetHalModuleDebugSink(*debug_sink); |
There was a problem hiding this comment.
Need to think a bit carefully about this: The underlying VM module can have a lifetime that exceeds the Python level wrapper for it. Reviewing the rest of the patch and will see if I can spot a better way.
There was a problem hiding this comment.
I came up with 2 solutions.
- Burden the user to guarantee that the callback function outlives the HAL module by holding a reference to it. In general not a good solution in Python.
- In order to bind the lifetime of the Python callback object to the HAL module we need to be able to register hooks with the module where it gets called on destruction. This burdens the IREE runtime with more stuff. Maybe not desirable. Instead of doing this for every VM module maybe at least the IREE runtime instance can have on-destruction hooks. The general problem is that we don't know when stuff gets destroyed.
There was a problem hiding this comment.
Another option is to add ownership of the debug sink to the VM context when the HAL module is registered there.
Then as long as the Python object of the VM context lives the debug sink will live.
This solution suffers from the problem of interop between the the Python bindings and some other interaction with the runtime. For example if we take the native handle from the Python VM context object and start using it while we destroy the Python object. I am not sure currently if the the bindings support this interaction, but I may be a future use case.
Signed-off-by: Boian Petkantchin <[email protected]>
We don't support custom debug sinks in the Runtime Python bindings. In particular the ability to register a custom callback when tracing tensors.
This change makes it possible to create a HAL module with a Python function as a callback.
A special care is taken to avoid leaking due to circular references where the callback reference the HAL module itself.