[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.

Author: Lunderberg

python/tvm/runtime/disco/session.py

@@ -142,6 +142,10 @@ def empty(
         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.
 

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