This is an automated email from the ASF dual-hosted git repository.
masahi pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new fd820ade5f [Disco] Expose disco.Session.shutdown through the python
API (#16979)
fd820ade5f is described below
commit fd820ade5fd68db3b4b4caa2e3d5bf3f0f48018c
Author: Eric Lunderberg <lunderb...@users.noreply.github.com>
AuthorDate: Mon May 13 14:27:24 2024 -0500
[Disco] Expose disco.Session.shutdown through the python API (#16979)
Prior to this commit, the `SessionObj::Shutdown` method could be
called from the C++ API, but could not be called through the Python
API. While it is implicitly called when the `SessionObj` is
destructed, Python's garbage collection may result in the destruction
occurring later than expected.
This commit exposes `SessionObj::Shutdown` through the Python API as
`disco.Session.shutdown`, allowing it to be closed cleanly.
---
python/tvm/runtime/disco/session.py | 4 ++++
src/runtime/disco/session.cc | 2 ++
2 files changed, 6 insertions(+)
diff --git a/python/tvm/runtime/disco/session.py
b/python/tvm/runtime/disco/session.py
index b8f74bacb0..ee151db716 100644
--- a/python/tvm/runtime/disco/session.py
+++ b/python/tvm/runtime/disco/session.py
@@ -142,6 +142,10 @@ class Session(Object):
func = self._get_cached_method("runtime.disco.empty")
return func(ShapeTuple(shape), dtype, device)
+ def shutdown(self):
+ """Shut down the Disco session"""
+ _ffi_api.SessionShutdown(self) # type: ignore # pylint:
disable=no-member
+
def get_global_func(self, name: str) -> DRef:
"""Get a global function on workers.
diff --git a/src/runtime/disco/session.cc b/src/runtime/disco/session.cc
index 12339c4fa5..e74d3819fe 100644
--- a/src/runtime/disco/session.cc
+++ b/src/runtime/disco/session.cc
@@ -52,6 +52,8 @@
TVM_REGISTER_GLOBAL("runtime.disco.SessionCallPacked").set_body([](TVMArgs args,
*rv = SessionObj::FFI::CallWithPacked(
self, TVMArgs(args.values + 1, args.type_codes + 1, args.num_args - 1));
});
+TVM_REGISTER_GLOBAL("runtime.disco.SessionShutdown")
+ .set_body_method<Session>(&SessionObj::Shutdown);
} // namespace runtime
} // namespace tvm
