Reject an unsupported ReduceOp before any message; remove the opt-in COMMUX_SPIN_US spin (#13, #12)

.github/workflows/ci.yml

@@ -94,18 +94,24 @@ jobs:
             -C cmake.define.COMMUX_UCX_PROVIDER=system \
             -C cmake.define.COMMUX_CUDA=OFF .
 
-      - name: Smoke import + end-to-end backend tests (CPU, 2-4 ranks)
+      - name: Smoke import + end-to-end backend tests (CPU, 2-7 ranks)
+        env:
+          UCX_TLS: tcp,sm,self
         run: |
           python -c "import commux; commux.register(); print('commux import OK')"
           torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/test_torch_dist.py
-          for n in 3 4; do torchrun --standalone --nnodes=1 --nproc-per-node=$n tests/test_collectives.py; done
+          for n in 3 4 5 6 7; do torchrun --standalone --nnodes=1 --nproc-per-node=$n tests/test_collectives.py; done
+          for n in 3 4; do timeout 120 torchrun --standalone --nnodes=1 --nproc-per-node=$n tests/test_reduce_op.py; done
 
       - name: C++ test (2 ranks over TCPStore)
+        env:
+          UCX_TLS: tcp,sm,self
         run: |
           # COMMUX_MT runs the concurrent multi-threaded send/recv test;
           # COMMUX_DEADLOCK runs the lost-wakeup regression test (many threads
           # parked on the shared worker must all wake). Both are 2-rank.
           export COMMUX_MT=1 COMMUX_DEADLOCK=1
           ./build/test_commux 1 2 127.0.0.1 29555 &
+          rank1=$!
           ./build/test_commux 0 2 127.0.0.1 29555
-          wait
+          wait "$rank1"  # a bare wait returns 0; this returns rank 1's status

src/process_group_ucx.cpp

@@ -5,7 +5,6 @@
 #include <dlfcn.h>
 #include <poll.h>
 
-#include <chrono>
 #include <cstdlib>
 #include <cstring>
 #include <stdexcept>
@@ -60,58 +59,9 @@ int wait_poll_ms() {
 // an event arrives, so latency is unchanged in the common case; the timeout
 // only caps idle sleep, keeping CPU near zero rather than busy-spinning
 // ucp_worker_progress().
-int spin_us() {
-  static const int us = [] {
-    const char* e = std::getenv("COMMUX_SPIN_US");
-    int v = e ? std::atoi(e) : 0;
-    return v > 0 ? v : 0;
-  }();
-  return us;
-}
-
-// One budget per wait(). Re-entries after an unrelated completion keep `end`,
-// so the total spin is COMMUX_SPIN_US however many completions arrive.
-struct SpinBudget {
-  bool on = false;
-  std::chrono::steady_clock::time_point end{};
-  const std::chrono::steady_clock::time_point* deadline() const {
-    return on ? &end : nullptr;
-  }
-  bool expired() const {
-    return on && !(std::chrono::steady_clock::now() < end);
-  }
-  static SpinBudget start() {
-    SpinBudget b;
-    int us = spin_us();
-    if (us > 0) {
-      b.on = true;
-      b.end = std::chrono::steady_clock::now() + std::chrono::microseconds(us);
-    }
-    return b;
-  }
-};
-
-// `deadline` is the end of this wait's whole spin budget. Unrelated
-// completions return early so the caller can reap, and the same deadline is
-// passed back in. A null deadline (COMMUX_SPIN_US unset or 0) does not spin.
-// Once `deadline` is already past, a UCS_ERR_BUSY return must not fall back
-// into a tight progress loop: arm refuses to sleep while events are queued,
-// which is exactly a stream of unrelated completions.
-void worker_wait_timed(ucp_worker_h worker,
-                       const std::chrono::steady_clock::time_point* deadline) {
-  if (deadline != nullptr) {
-    while (std::chrono::steady_clock::now() < *deadline) {
-      if (ucp_worker_progress(worker)) return;
-    }
-  }
+void worker_wait_timed(ucp_worker_h worker) {
   ucs_status_t s = ucp_worker_arm(worker);
-  if (s == UCS_ERR_BUSY) {
-    if (deadline != nullptr &&
-        !(std::chrono::steady_clock::now() < *deadline)) {
-      std::this_thread::sleep_for(std::chrono::milliseconds(wait_poll_ms()));
-    }
-    return;  // events arrived during arm -> progress now
-  }
+  if (s == UCS_ERR_BUSY) return;  // events arrived during arm -> progress now
   if (s != UCS_OK) {
     std::this_thread::yield();  // wakeup unsupported: fall back to a yield
     return;
@@ -280,6 +230,21 @@ void apply_reduce(at::Tensor& acc, const at::Tensor& other, c10d::ReduceOp op) {
   }
 }
 
+// Every rank checks the op before a reduce sends anything: apply_reduce runs
+// only on the ranks that receive, after the others have sent and returned.
+void check_reduce_op(c10d::ReduceOp op, const char* what) {
+  switch (op.op_) {
+    case c10d::ReduceOp::SUM:
+    case c10d::ReduceOp::PRODUCT:
+    case c10d::ReduceOp::MIN:
+    case c10d::ReduceOp::MAX:
+      return;
+    default:
+      TORCH_CHECK(false, "commux ", what, ": unsupported ReduceOp ",
+                  static_cast<int>(op.op_));
+  }
+}
+
 // c10d::Work wrapping a set of in-flight ucp requests. wait()/isCompleted()
 // drive the shared worker's progress engine (the classic "wait_req" loop). The
 // worker is UCS_THREAD_MODE_MULTI (internally thread-safe), so waiting on one
@@ -324,18 +289,12 @@ class UCXWork : public c10d::Work {
     // worker -- their progress is precisely what completes this Work. (Holding
     // the mutex across a blocking ucp_worker_wait() would serialize, and can
     // deadlock, multi-threaded consumers such as snapy.)
-    SpinBudget spin = SpinBudget::start();
     while (!reap()) {
-      if (ucp_worker_progress(worker_)) {
-        // Advanced. A stream of unrelated completions stays on this branch and
-        // would never consult the one per-wait deadline below.
-        if (spin.expired()) worker_wait_timed(worker_, spin.deadline());
-        continue;
-      }
+      if (ucp_worker_progress(worker_)) continue;  // advanced; re-check
       // Idle: sleep on the worker's wakeup fd, but only until the next event or
       // a short timeout -- a bounded sleep so concurrently parked threads can't
       // miss a lost wakeup on the shared worker (see worker_wait_timed).
-      worker_wait_timed(worker_, spin.deadline());
+      worker_wait_timed(worker_);
     }
     if (err_) {
       std::rethrow_exception(err_);
@@ -557,14 +516,9 @@ void ProcessGroupUCX::wait_request(void* req) {
   // wakeup fd (bounded, to avoid the lost-wakeup deadlock -- see
   // worker_wait_timed) instead of busy-spinning. Caller holds worker_mu_.
   ucs_status_t st;
-  SpinBudget spin = SpinBudget::start();
   while ((st = ucp_request_check_status(req)) == UCS_INPROGRESS) {
-    if (ucp_worker_progress(worker_)) {
-      // Same budget as UCXWork::wait. Unrelated completions must not skip it.
-      if (spin.expired()) worker_wait_timed(worker_, spin.deadline());
-      continue;
-    }
-    worker_wait_timed(worker_, spin.deadline());
+    if (ucp_worker_progress(worker_)) continue;
+    worker_wait_timed(worker_);
   }
   ucp_request_free(req);
   TORCH_CHECK(st == UCS_OK, "commux request failed: ", ucs_status_string(st));
@@ -802,6 +756,7 @@ void ProcessGroupUCX::broadcast_locked(std::vector<at::Tensor>& tensors,
 
 c10::intrusive_ptr<c10d::Work> ProcessGroupUCX::allreduce(
     std::vector<at::Tensor>& tensors, const c10d::AllreduceOptions& opts) {
+  check_reduce_op(opts.reduceOp, "allreduce");
   std::lock_guard<std::mutex> lock(worker_mu_);
   if (size_ > 1) {
     reduce_locked(tensors, opts.reduceOp, /*root=*/0);
@@ -812,6 +767,7 @@ c10::intrusive_ptr<c10d::Work> ProcessGroupUCX::allreduce(
 
 c10::intrusive_ptr<c10d::Work> ProcessGroupUCX::reduce(
     std::vector<at::Tensor>& tensors, const c10d::ReduceOptions& opts) {
+  check_reduce_op(opts.reduceOp, "reduce");
   std::lock_guard<std::mutex> lock(worker_mu_);
   if (size_ > 1) {
     reduce_locked(tensors, opts.reduceOp, static_cast<int>(opts.rootRank));

tests/test_collectives.py

@@ -7,6 +7,8 @@
 other rank's input unchanged, and broadcast gives every rank the root's tensor.
 all_reduce (SUM) gives every rank the sum. World sizes that are not a power of
 two and roots other than 0 cover any relabelling of ranks inside a collective.
+PRODUCT, MIN and MAX run at every root too; a 2 MiB all_reduce takes the
+rendezvous path, and a three-tensor all_reduce gives each tensor its own tag.
 The inputs are small integers, so a sum is exact in any order. Each rank counts
 its failures and they are agreed with a final all_reduce, so no rank leaves a
 collective early.
@@ -47,6 +49,35 @@ def check(name, got, want):
 dist.all_reduce(t, op=dist.ReduceOp.SUM)
 check("all_reduce", t, total)
 
+# The other supported ops. The tree combines in a different order from a
+# linear loop; products of these inputs stay exact in float64 up to 7 ranks.
+for op, fold in ((dist.ReduceOp.PRODUCT, torch.mul),
+                 (dist.ReduceOp.MIN, torch.minimum),
+                 (dist.ReduceOp.MAX, torch.maximum)):
+    want = mine(0)
+    for r in range(1, world):
+        want = fold(want, mine(r))
+    for root in range(world):
+        t = mine(rank)
+        dist.reduce(t, dst=root, op=op)
+        check(f"reduce {op} root={root}", t, want if rank == root else mine(rank))
+    t = mine(rank)
+    dist.all_reduce(t, op=op)
+    check(f"all_reduce {op}", t, want)
+
+# 2 MiB, past the eager limit, so the message takes the rendezvous path.
+big = torch.arange(1 << 18, dtype=torch.float64)
+t = big + rank
+dist.all_reduce(t, op=dist.ReduceOp.SUM)
+if not torch.equal(t, world * big + world * (world - 1) / 2):
+    failures.append(f"all_reduce 2 MiB: rank {rank} differs")
+
+# Several tensors in one call, each on its own tag.
+ts = [mine(rank) * (k + 1) for k in range(3)]
+dist.distributed_c10d._get_default_group().allreduce(ts).wait()
+for k in range(3):
+    check(f"all_reduce list[{k}]", ts[k], total * (k + 1))
+
 for msg in failures:
     print(f"FAIL {msg}", flush=True)
 count = torch.tensor([float(len(failures))])

tests/test_commux.cpp

@@ -14,7 +14,6 @@
 //   test_commux <rank> <world_size> [master_host=127.0.0.1] [master_port=29555]
 //   UCXPG_DEVICE=cuda  runs the same checks on GPU tensors.
 
-#include <sys/resource.h>
 #include <torch/torch.h>
 
 #include <atomic>
@@ -336,139 +335,6 @@ int main(int argc, char** argv) {
     watchdog.join();
   }
 
-  // COMMUX_SPIN_US only. A peer drips unrelated messages for longer than the
-  // spin window, then sends the message this rank is blocked on. The whole
-  // wait gets one window, so CPU time stays near that budget. Default-off
-  // runs skip this.
-  if (const char* spin_env = std::getenv("COMMUX_SPIN_US");
-      spin_env != nullptr && std::atoi(spin_env) > 0 && size == 2) {
-    const int peer = 1 - rank;
-    const int n_noise = 40;
-    const int gap_ms = 3;
-    const int target_tag = 9100;
-    std::vector<at::Tensor> noise_bufs;
-    std::vector<c10::intrusive_ptr<c10d::Work>> noise_works;
-    at::Tensor target;
-    c10::intrusive_ptr<c10d::Work> target_work;
-    if (rank == 0) {
-      target = torch::zeros({8}, f64);
-      std::vector<at::Tensor> tv{target};
-      target_work = pg->recv(tv, peer, target_tag);
-      for (int i = 0; i < n_noise; ++i) {
-        noise_bufs.push_back(torch::zeros({8}, f64));
-        std::vector<at::Tensor> v{noise_bufs.back()};
-        noise_works.push_back(pg->recv(v, peer, 9200 + i));
-      }
-    }
-    pg->barrier()->wait();
-    if (rank == 0) {
-      struct rusage ru0{};
-      struct rusage ru1{};
-      getrusage(RUSAGE_THREAD, &ru0);
-      auto wall0 = std::chrono::steady_clock::now();
-      target_work->wait();
-      long wall_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
-                         std::chrono::steady_clock::now() - wall0)
-                         .count();
-      getrusage(RUSAGE_THREAD, &ru1);
-      auto usec = [](const timeval& t) {
-        return static_cast<long>(t.tv_sec) * 1000000L +
-               static_cast<long>(t.tv_usec);
-      };
-      long cpu_us = (usec(ru1.ru_utime) - usec(ru0.ru_utime)) +
-                    (usec(ru1.ru_stime) - usec(ru0.ru_stime));
-      for (auto& w : noise_works) w->wait();
-      bool value_ok = target.eq(42.0).all().item<bool>();
-      long budget_us = std::atol(spin_env);
-      long drip_us = static_cast<long>(n_noise) * gap_ms * 1000L;
-      // A window longer than the drip spins until the message arrives.
-      bool parked = budget_us * 2 >= drip_us || (wall_ms * 1000 > drip_us / 2 &&
-                                                 cpu_us * 2 < wall_ms * 1000);
-      std::printf("  spin check: wall %ld ms, cpu %ld us, budget %s us\n",
-                  wall_ms, cpu_us, spin_env);
-      check(value_ok && parked,
-            "bounded spin parks while unrelated traffic flows");
-    } else {
-      for (int i = 0; i < n_noise; ++i) {
-        auto s = torch::full({8}, 1.0, f64);
-        std::vector<at::Tensor> v{s};
-        pg->send(v, peer, 9200 + i)->wait();
-        std::this_thread::sleep_for(std::chrono::milliseconds(gap_ms));
-      }
-      auto s = torch::full({8}, 42.0, f64);
-      std::vector<at::Tensor> v{s};
-      pg->send(v, peer, target_tag)->wait();
-    }
-
-    // Tight unrelated sends, no gap, held for longer than the budget. A short
-    // burst finishes inside the budget and both libraries spin through it.
-    // These recvs are posted only after the target arrives, so during the wait
-    // the messages are unexpected and ucp_worker_progress stays non-zero. The
-    // outer `if (progress) continue` has to read the deadline itself.
-    const int flood_ms = 40;
-    const int flood_max = 20000;
-    const int flood_target = 9400;
-    at::Tensor flood_got;
-    c10::intrusive_ptr<c10d::Work> flood_work;
-    if (rank == 0) {
-      flood_got = torch::empty({8}, f64);
-      std::vector<at::Tensor> tv{flood_got};
-      flood_work = pg->recv(tv, peer, flood_target);
-    }
-    pg->barrier()->wait();
-    if (rank == 0) {
-      struct rusage ru0{};
-      struct rusage ru1{};
-      getrusage(RUSAGE_THREAD, &ru0);
-      auto wall0 = std::chrono::steady_clock::now();
-      flood_work->wait();
-      long wall_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
-                         std::chrono::steady_clock::now() - wall0)
-                         .count();
-      getrusage(RUSAGE_THREAD, &ru1);
-      auto usec = [](const timeval& t) {
-        return static_cast<long>(t.tv_sec) * 1000000L +
-               static_cast<long>(t.tv_usec);
-      };
-      long cpu_us = (usec(ru1.ru_utime) - usec(ru0.ru_utime)) +
-                    (usec(ru1.ru_stime) - usec(ru0.ru_stime));
-      int n = static_cast<int>(flood_got[0].item<double>());
-      std::vector<at::Tensor> flood_bufs;
-      std::vector<c10::intrusive_ptr<c10d::Work>> flood_works;
-      if (n > 0) flood_bufs.reserve(static_cast<size_t>(n));
-      for (int i = 0; i < n; ++i) {
-        flood_bufs.push_back(torch::empty({8}, f64));
-        std::vector<at::Tensor> v{flood_bufs.back()};
-        flood_works.push_back(pg->recv(v, peer, 9500 + i));
-      }
-      for (auto& w : flood_works) w->wait();
-      long budget_us = std::atol(spin_env);
-      long wall_us = wall_ms * 1000L;
-      // A flood that finishes inside the budget cannot fail a full spin.
-      bool long_enough = wall_ms + 5 >= flood_ms;
-      bool parked = long_enough && cpu_us * 2 < wall_us;
-      std::printf(
-          "  spin flood: wall %ld ms, cpu %ld us, budget %ld us, n %d\n",
-          wall_ms, cpu_us, budget_us, n);
-      check(
-          n > 0 && parked,
-          "bounded spin parks while unrelated completions arrive back to back");
-    } else {
-      int n = 0;
-      auto end = std::chrono::steady_clock::now() +
-                 std::chrono::milliseconds(flood_ms);
-      while (std::chrono::steady_clock::now() < end && n < flood_max) {
-        auto s = torch::full({8}, 1.0, f64);
-        std::vector<at::Tensor> v{s};
-        pg->send(v, peer, 9500 + n)->wait();
-        ++n;
-      }
-      auto s = torch::full({8}, static_cast<double>(n), f64);
-      std::vector<at::Tensor> v{s};
-      pg->send(v, peer, flood_target)->wait();
-    }
-  }
-
   pg->barrier()->wait();
   if (rank == 0) {
     std::printf(