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