gemini-code-assist[bot] commented on code in PR #128:
URL: https://github.com/apache/tvm-ffi/pull/128#discussion_r2432759240
##########
python/tvm_ffi/module.py:
##########
@@ -54,6 +54,32 @@ class Module(core.Object):
--------
:py:func:`tvm_ffi.load_module`
+ Notes
+ -----
+ If you load a module within a local scope, be careful when any called
function
+ creates and returns an object. That object's deleter resides in the loaded
library.
+ If the module is unloaded before the object is destroyed, the deleter may
call an
+ invalid address. Keep the module loaded until all returned objects are
deleted.
+ You can safely use returned objects inside a nested function that finishes
before
+ the module goes out of scope. When possible, consider keeping the module
alive in
+ a long-lived/global scope (for example, in a global state) to avoid
premature unloading.
+
+ .. code-block:: python
+
+ def bad_pattern(x):
+ # Bad: unload order of `tensor` and `mod` is not guaranteed
+ mod = tvm_ffi.load_module("path/to/library.so")
+ # ... do something with the tensor
+ tensor = mod.func_create_and_return_tensor(x)
+
+ def good_pattern(x):
+ # Good: `tensor` is freed before `mod` goes out of scope
+ mod = tvm_ffi.load_module("path/to/library.so")
+ def run_some_tests():
+ tensor = mod.func_create_and_return_tensor(x)
+ # ... do something with the tensor
+ run_some_tests()
+
"""
Review Comment:
This addition to the documentation is helpful in clarifying the lifecycle
issues related to dynamically loaded modules. However, consider adding a brief
explanation of *why* the deleter resides in the loaded library. This will
provide more context to the user.
For example, you could add a sentence like: "This is because the memory
allocation and deallocation routines are part of the library's code."
##########
tests/python/test_load_inline.py:
##########
@@ -243,7 +243,14 @@ def test_load_inline_with_env_tensor_allocator() -> None:
""",
functions=["return_add_one"],
)
- if torch is not None:
+ assert torch is not None
+
+ def run_check() -> None:
+ """Must run in a separate function to ensure deletion happens before
mod unloads.
+
+ When a module returns an object, the object deleter address is part of
the
+ loaded library. We need to keep the module loaded until the object is
deleted.
+ """
Review Comment:
This docstring is helpful, but it could be improved by explicitly mentioning
the role of the garbage collector in this scenario. The garbage collector might
unload the module before the tensor is deleted if the tensor's lifetime extends
beyond the scope where the module was loaded.
Consider adding a sentence like: "Without this separation, the Python
garbage collector might release the module before the tensor, leading to a
crash."
##########
tests/python/test_load_inline.py:
##########
@@ -341,11 +350,16 @@ def test_cuda_memory_alloc_noleak() -> None:
""",
functions=["return_tensor"],
)
- x = torch.arange(1024 * 1024, dtype=torch.float32, device="cuda")
- current_allocated = torch.cuda.memory_allocated()
- repeat = 8
- for i in range(repeat):
- mod.return_tensor(x)
- diff = torch.cuda.memory_allocated() - current_allocated
- # memory should not grow as we loop over
- assert diff <= 1024**2 * 8
+
+ def run_check() -> None:
+ """Must run in a separate function to ensure deletion happens before
mod unloads."""
Review Comment:
Consider adding a comment explaining *why* the `run_check` function is
necessary. This will help readers understand the purpose of this refactoring.
For example, you could add a sentence like: "This function ensures that the
tensor `x` is deleted before the module `mod` is unloaded, preventing potential
memory errors."
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
