[Disco] Implement num_workers property for disco.Session

Prior to this commit, while the `num_workers` argument was provided to
the `disco.Session` object, it could not be determined from an
existing `disco.Session` object.  As a result, functions that
interacted with a multi-GPU setup frequently required separate
`num_workers` and `disco_session` argument, which could erroneously be
out-of-sync (e.g. passing the incorrect `num_workers`, or omitting the
`disco_session` argument when `num_workers>1`).

To remove this class of errors, this commit adds a
`disco.Session.num_workers` property.  The separate `num_workers`
argument is no longer necessary, as it can be determined from the
`disco.Session` instance.

Author: Lunderberg

include/tvm/runtime/disco/session.h

@@ -197,6 +197,8 @@ class SessionObj : public Object {
    * The thirtd element is the function to be called.
    */
   TVM_DLL virtual DRef CallWithPacked(const TVMArgs& args) = 0;
+  /*! \brief Get the number of workers in the session. */
+  TVM_DLL virtual int64_t GetNumWorkers() = 0;
   /*! \brief Get a global functions on workers. */
   TVM_DLL virtual DRef GetGlobalFunc(const std::string& name) = 0;
   /*!

python/tvm/runtime/disco/session.py

@@ -142,6 +142,11 @@ def empty(
         func = self._get_cached_method("runtime.disco.empty")
         return func(ShapeTuple(shape), dtype, device)
 
+    @property
+    def num_workers(self) -> int:
+        """Return the number of workers in the session"""
+        return _ffi_api.SessionGetNumWorkers(self)  # type: ignore # pylint: disable=no-member
+
     def get_global_func(self, name: str) -> DRef:
         """Get a global function on workers.
 

src/runtime/disco/process_session.cc

@@ -153,6 +153,8 @@ class ProcessSessionObj final : public BcastSessionObj {
 
   ~ProcessSessionObj() { Kill(); }
 
+  int64_t GetNumWorkers() { return workers_.size() + 1; }
+
   TVMRetValue DebugGetFromRemote(int64_t reg_id, int worker_id) {
     if (worker_id == 0) {
       this->SyncWorker(worker_id);

src/runtime/disco/session.cc

@@ -37,6 +37,8 @@ TVM_REGISTER_GLOBAL("runtime.disco.DRefDebugGetFromRemote")
     .set_body_method<DRef>(&DRefObj::DebugGetFromRemote);
 TVM_REGISTER_GLOBAL("runtime.disco.DRefDebugCopyFrom")
     .set_body_method<DRef>(&DRefObj::DebugCopyFrom);
+TVM_REGISTER_GLOBAL("runtime.disco.SessionGetNumWorkers")
+    .set_body_method<Session>(&SessionObj::GetNumWorkers);
 TVM_REGISTER_GLOBAL("runtime.disco.SessionGetGlobalFunc")
     .set_body_method<Session>(&SessionObj::GetGlobalFunc);
 TVM_REGISTER_GLOBAL("runtime.disco.SessionCopyFromWorker0")

src/runtime/disco/threaded_session.cc

@@ -154,6 +154,8 @@ class ThreadedSessionObj final : public BcastSessionObj {
     workers_.clear();
   }
 
+  int64_t GetNumWorkers() { return workers_.size(); }
+
   TVMRetValue DebugGetFromRemote(int64_t reg_id, int worker_id) {
     this->SyncWorker(worker_id);
     return this->workers_.at(worker_id).worker->register_file.at(reg_id);

tests/python/disco/test_session.py

@@ -220,6 +220,13 @@ def transpose_2(
         np.testing.assert_equal(z_nd, x_np)
 
 
+@pytest.mark.parametrize("session_kind", _all_session_kinds)
+@pytest.mark.parametrize("num_workers", [1, 2, 4])
+def test_num_workers(session_kind, num_workers):
+    sess = session_kind(num_workers=num_workers)
+    assert sess.num_workers == num_workers
+
+
 if __name__ == "__main__":
     test_int(di.ProcessSession)
     test_float(di.ProcessSession)