Merge pull request #12 from fengidri/main

Adjust the barrier api, change default mode to all instance barrier; Add barrier in server side in bmk.

Author: niehao100

fserver/csrc/public.hpp

@@ -148,7 +148,7 @@ void wait(int handler) {
   fworker_->Wait(handler);
 }
 
-void barrier(bool include_server, bool include_worker) {
+void barrier(bool include_server, bool include_worker, bool instrance_barrier=true) {
   int node_group = 0;
   if (include_server) {
     node_group += ps::kServerGroup;
@@ -158,9 +158,9 @@ void barrier(bool include_server, bool include_worker) {
   }
 
   if (role_ == Node::WORKER && include_worker) {
-    ps::Postoffice::GetWorker(instance_id_)->Barrier(0, node_group);
+    ps::Postoffice::GetWorker(instance_id_)->DoBarrier(0, node_group, instrance_barrier);
   } else if (role_ == Node::SERVER && include_server) {
-    ps::Postoffice::GetServer(instance_id_)->Barrier(0, node_group);
+    ps::Postoffice::GetServer(instance_id_)->DoBarrier(0, node_group, instrance_barrier);
   } else {
     ps::Postoffice::Get()->Barrier(0, node_group);
   }
@@ -261,7 +261,11 @@ void pybind_public(py::module &m){
   // fetch_trace needs gil_scoped_release
   m.def("fetch_trace", &fetch_trace, py::call_guard<py::gil_scoped_release>());
   m.def("get_all_handlers", &get_all_handlers, py::call_guard<py::none>());
-  m.def("barrier", &barrier, py::call_guard<py::none>());
+  m.def("barrier", &barrier, 
+    py::arg("include_server"),
+    py::arg("include_client"),
+    py::arg("instance_barrier") = true,
+    py::call_guard<py::none>());
   m.def("get_nanosecond", &get_nanosecond, py::call_guard<py::none>());
 }
 

include/ps/internal/postoffice.h

@@ -236,6 +236,13 @@ class Postoffice {
    * \param node_id the barrier group id
    */
   void Barrier(int customer_id, int node_group);
+  /**
+   * \brief barrier for all postoffice instances or groups
+   * \param customer_id the id of the customer
+   * \param node_id the barrier group id
+   * \param instance_barrier whether it's for all postoffice instances or groups
+   */
+  void DoBarrier(int customer_id, int node_group, bool instance_barrier);
   /**
    * \brief process a control message, called by van
    * \param the received message
@@ -267,14 +274,6 @@ class Postoffice {
   explicit Postoffice(int instance_idx);
   ~Postoffice() { delete van_; }
 
-  /**
-   * \brief barrier for all postoffice instances or groups
-   * \param customer_id the id of the customer
-   * \param node_id the barrier group id
-   * \param instance_barrier whether it's for all postoffice instances or groups
-   */
-  void DoBarrier(int customer_id, int node_group, bool instance_barrier);
-
   static Postoffice* po_scheduler_;
   static std::mutex init_mu_;
   // the group of postoffices for workers

tests/benchmark/bmk_comm_latency_multiserver.py

@@ -171,7 +171,7 @@ def _str_formater(data):
 if is_worker:
     th = threading.Thread(target=print_thread)
     th.start()
-    f.barrier(True, True)
+    f.barrier(False, True)
     q = Queue()
     time_list = []
     net_cost_list = [[] for _ in range(bsz + 1 + bsz)]
@@ -202,7 +202,7 @@ def worker():
 elif is_server:
     ret_buffer = torch.rand([65535, dim], dtype=torch.bfloat16, device='cuda')
     count = 0
-
+    f.barrier(True, False)
     def server():
         global count
         iter_count = 0