[coll] Support worker port. (#12010)

Author: trivialfis

python-package/xgboost/collective.py

@@ -6,7 +6,7 @@
 import pickle
 from dataclasses import dataclass
 from enum import IntEnum, unique
-from typing import Any, Dict, Optional, TypeAlias, Union
+from typing import Any, Callable, Dict, Optional, TypeAlias, Union
 
 import numpy as np
 
@@ -46,6 +46,21 @@ class Config:
 
     tracker_timeout : See :py:class:`~xgboost.tracker.RabitTracker`.
 
+    worker_port :
+
+        The port each worker listens to for peer-to-peer connections. By default,
+        workers use an available port assigned by the OS. This option can be used in
+        restricted network environments where only specific ports are open.
+
+        This can be an integer for a fixed port used by all workers, or a callback
+        function that takes no arguments and returns a port number. The callback is
+        invoked per-worker at the worker side.
+
+        .. note::
+
+            The option does not affect the NCCL communicator group, which must be
+            configured via NCCL's own environment variables.
+
     """
 
     retry: Optional[int] = None
@@ -55,6 +70,18 @@ class Config:
     tracker_port: Optional[int] = None
     tracker_timeout: Optional[int] = None
 
+    worker_port: Optional[Union[Callable[[], int], int]] = None
+
+    def update_worker_args(self, args: _Conf) -> _Conf:
+        """Worker side arguments resolution."""
+        if self.worker_port is None:
+            return args
+        if callable(self.worker_port):
+            args["dmlc_worker_port"] = self.worker_port()
+        else:
+            args["dmlc_worker_port"] = self.worker_port
+        return args
+
     def get_comm_config(self, args: _Conf) -> _Conf:
         """Update the arguments for the communicator."""
         if self.retry is not None:

python-package/xgboost/dask/__init__.py

@@ -761,6 +761,9 @@ def do_train(  # pylint: disable=too-many-positional-arguments
         local_history: TrainingCallback.EvalsLog = {}
         global_config.update({"nthread": n_threads})
 
+        if coll_cfg is not None:
+            coll_args = coll_cfg.update_worker_args(coll_args)
+
         with CommunicatorContext(**coll_args), config.config_context(**global_config):
             Xy, evals = _get_dmatrices(
                 train_ref,

python-package/xgboost/spark/core.py

@@ -1120,6 +1120,9 @@ def _train_booster(
             if not launch_tracker_on_driver:
                 _rabit_args = json.loads(messages[0])["rabit_msg"]
 
+            if conf is not None:
+                _rabit_args = conf.update_worker_args(_rabit_args)
+
             evals_result: Dict[str, Any] = {}
             with (
                 config_context(verbosity=verbosity, use_rmm=use_rmm),
@@ -1641,7 +1644,11 @@ def saveMetadata(
         if instance.isDefined("coll_cfg"):
             conf: Config = instance.getOrDefault("coll_cfg")
             if conf is not None:
-                extraMetadata["coll_cfg"] = asdict(conf)
+                extraMetadata["coll_cfg"] = {
+                    k: v for k, v in asdict(conf).items() if not callable(v)
+                }
+            if callable(conf.worker_port):
+                logger.warning("The `worker_port` is not serialized.")
 
         DefaultParamsWriter.saveMetadata(
             instance, path, sc, extraMetadata=extraMetadata, paramMap=jsonParams

src/collective/comm.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2023-2024, XGBoost Contributors
+ * Copyright 2023-2026, XGBoost Contributors
  */
 #include "comm.h"
 
@@ -108,7 +108,9 @@ Result ConnectTrackerImpl(proto::PeerInfo info, std::chrono::seconds timeout, st
 
   rc = std::move(rc) << [&] {
     return cpu_impl::RingAllgather(comm, s_buffer, HOST_NAME_MAX, 0, prev_ch, next_ch);
-  } << [&] { return block(); };
+  } << [&] {
+    return block();
+  };
   if (!rc.OK()) {
     return Fail("Failed to get host names from peers.", std::move(rc));
   }
@@ -119,7 +121,9 @@ Result ConnectTrackerImpl(proto::PeerInfo info, std::chrono::seconds timeout, st
     auto s_ports = common::Span{reinterpret_cast<std::int8_t*>(peers_port.data()),
                                 peers_port.size() * sizeof(ninfo.port)};
     return cpu_impl::RingAllgather(comm, s_ports, sizeof(ninfo.port), 0, prev_ch, next_ch);
-  } << [&] { return block(); };
+  } << [&] {
+    return block();
+  };
   if (!rc.OK()) {
     return Fail("Failed to get the port from peers.", std::move(rc));
   }
@@ -143,9 +147,11 @@ Result ConnectTrackerImpl(proto::PeerInfo info, std::chrono::seconds timeout, st
   for (std::int32_t r = (comm.Rank() + 1); r < comm.World(); ++r) {
     auto const& peer = peers[r];
     auto worker = std::make_shared<TCPSocket>();
-    rc = std::move(rc)
-         << [&] { return Connect(peer.host, peer.port, retry, timeout, worker.get()); }
-         << [&] { return worker->RecvTimeout(timeout); };
+    rc = std::move(rc) << [&] {
+      return Connect(peer.host, peer.port, retry, timeout, worker.get());
+    } << [&] {
+      return worker->RecvTimeout(timeout);
+    };
     if (!rc.OK()) {
       return rc;
     }
@@ -204,17 +210,18 @@ std::string InitLog(std::string task_id, std::int32_t rank) {
 
 RabitComm::RabitComm(std::string const& tracker_host, std::int32_t tracker_port,
                      std::chrono::seconds timeout, std::int32_t retry, std::string task_id,
-                     StringView nccl_path)
+                     StringView nccl_path, std::int32_t worker_port)
     : HostComm{tracker_host, tracker_port, timeout, retry, std::move(task_id)},
-      nccl_path_{std::move(nccl_path)} {
+      nccl_path_{std::move(nccl_path)},
+      worker_port_{worker_port} {
   if (this->TrackerInfo().host.empty()) {
     // Not in a distributed environment.
     LOG(CONSOLE) << InitLog(task_id_, rank_);
     return;
   }
 
   loop_.reset(new Loop{std::chrono::seconds{timeout_}});  // NOLINT
-  auto rc = this->Bootstrap(timeout_, retry_, task_id_);
+  auto rc = this->Bootstrap(timeout_, retry_, task_id_, worker_port_);
   if (!rc.OK()) {
     this->ResetState();
     SafeColl(Fail("Failed to bootstrap the communication group.", std::move(rc)));
@@ -230,7 +237,7 @@ Comm* RabitComm::MakeCUDAVar(Context const*, std::shared_ptr<Coll>) const {
 #endif  //  !defined(XGBOOST_USE_NCCL)
 
 [[nodiscard]] Result RabitComm::Bootstrap(std::chrono::seconds timeout, std::int32_t retry,
-                                          std::string task_id) {
+                                          std::string task_id, std::int32_t worker_port) {
   TCPSocket tracker;
   std::int32_t world{-1};
   auto rc = ConnectTrackerImpl(this->TrackerInfo(), timeout, retry, task_id, &tracker, this->Rank(),
@@ -243,8 +250,14 @@ Comm* RabitComm::MakeCUDAVar(Context const*, std::shared_ptr<Coll>) const {
 
   // Start command
   TCPSocket listener = TCPSocket::Create(tracker.Domain());
-  std::int32_t lport{0};
+  std::int32_t lport{worker_port};
   rc = std::move(rc) << [&] {
+    if (lport > 0) {
+      // User-specified port, bind to INADDR_ANY with the given port.
+      auto addr = (tracker.Domain() == SockDomain::kV6) ? "::" : "0.0.0.0";
+      return listener.Bind(addr, &lport);
+    }
+    // Default: let the OS pick an available port.
     return listener.BindHost(&lport);
   } << [&] {
     return listener.Listen();
@@ -306,8 +319,11 @@ Comm* RabitComm::MakeCUDAVar(Context const*, std::shared_ptr<Coll>) const {
   error_worker_.detach();
 
   proto::Start start;
-  rc = std::move(rc) << [&] { return start.WorkerSend(lport, &tracker, eport); }
-                     << [&] { return start.WorkerRecv(&tracker, &world); };
+  rc = std::move(rc) << [&] {
+    return start.WorkerSend(lport, &tracker, eport);
+  } << [&] {
+    return start.WorkerRecv(&tracker, &world);
+  };
   if (!rc.OK()) {
     return rc;
   }
@@ -418,8 +434,11 @@ RabitComm::~RabitComm() noexcept(false) {
   }
   TCPSocket out;
   proto::Print print;
-  return Success() << [&] { return this->ConnectTracker(&out); }
-                   << [&] { return print.WorkerSend(&out, msg); };
+  return Success() << [&] {
+    return this->ConnectTracker(&out);
+  } << [&] {
+    return print.WorkerSend(&out, msg);
+  };
 }
 
 [[nodiscard]] Result RabitComm::SignalError(Result const& res) {

src/collective/comm.h

@@ -1,5 +1,5 @@
 /**
- * Copyright 2023-2024, XGBoost Contributors
+ * Copyright 2023-2026, XGBoost Contributors
  */
 #pragma once
 #include <chrono>   // for seconds
@@ -127,16 +127,19 @@ class HostComm : public Comm {
 
 class RabitComm : public HostComm {
   std::string nccl_path_ = std::string{DefaultNcclName()};
+  // User-specified port for the worker listener socket. 0 means the OS picks an available
+  // port.
+  std::int32_t worker_port_{0};
 
   [[nodiscard]] Result Bootstrap(std::chrono::seconds timeout, std::int32_t retry,
-                                 std::string task_id);
+                                 std::string task_id, std::int32_t worker_port);
 
  public:
   // bootstrapping construction.
   RabitComm() = default;
   RabitComm(std::string const& tracker_host, std::int32_t tracker_port,
             std::chrono::seconds timeout, std::int32_t retry, std::string task_id,
-            StringView nccl_path);
+            StringView nccl_path, std::int32_t worker_port);
   ~RabitComm() noexcept(false) override;
 
   [[nodiscard]] bool IsFederated() const override { return false; }

src/collective/comm_group.cc

@@ -1,5 +1,5 @@
 /**
- * Copyright 2023-2024, XGBoost Contributors
+ * Copyright 2023-2026, XGBoost Contributors
  */
 #include "comm_group.h"
 
@@ -84,11 +84,17 @@ CommGroup::CommGroup()
     auto tracker_host = get_param("dmlc_tracker_uri", std::string{}, String{});
     auto tracker_port = get_param("dmlc_tracker_port", static_cast<std::int64_t>(0), Integer{});
     auto nccl = get_param("dmlc_nccl_path", std::string{DefaultNcclName()}, String{});
-    auto ptr = new CommGroup{
-        std::shared_ptr<RabitComm>{new RabitComm{  // NOLINT
-            tracker_host, static_cast<std::int32_t>(tracker_port), std::chrono::seconds{timeout},
-            static_cast<std::int32_t>(retry), task_id, nccl}},
-        std::shared_ptr<Coll>(new Coll{})};  // NOLINT
+    auto worker_port = get_param("dmlc_worker_port", static_cast<std::int64_t>(0), Integer{});
+    CHECK_LE(worker_port, std::numeric_limits<in_port_t>::max());
+    CHECK_GE(worker_port, 0);
+    CHECK_LE(tracker_port, std::numeric_limits<in_port_t>::max());
+    CHECK_GE(tracker_port, 0);
+    auto ptr = new CommGroup{std::shared_ptr<RabitComm>{new RabitComm{
+                                 // NOLINT
+                                 tracker_host, static_cast<std::int32_t>(tracker_port),
+                                 std::chrono::seconds{timeout}, static_cast<std::int32_t>(retry),
+                                 task_id, nccl, static_cast<std::int32_t>(worker_port)}},
+                             std::shared_ptr<Coll>(new Coll{})};  // NOLINT
     return ptr;
   } else if (type == "federated") {
 #if defined(XGBOOST_USE_FEDERATED)

src/predictor/interpretability/shap.cc

@@ -7,16 +7,16 @@
 #include <type_traits>  // for remove_const_t
 #include <vector>       // for vector
 
-#include "../../common/threading_utils.h"     // for ParallelFor
-#include "../../gbm/gbtree_model.h"           // for GBTreeModel
-#include "../../tree/tree_view.h"             // for ScalarTreeView
-#include "../data_accessor.h"                 // for GHistIndexMatrixView
-#include "../predict_fn.h"                    // for GetTreeLimit
-#include "../treeshap.h"                      // for CalculateContributions
-#include "dmlc/omp.h"                         // for omp_get_thread_num
-#include "xgboost/base.h"                     // for bst_omp_uint
-#include "xgboost/logging.h"                  // for CHECK
-#include "xgboost/multi_target_tree_model.h"  // for MTNotImplemented
+#include "../../common/threading_utils.h"  // for ParallelFor
+#include "../../gbm/gbtree_model.h"        // for GBTreeModel
+#include "../../tree/tree_view.h"          // for ScalarTreeView
+#include "../data_accessor.h"              // for GHistIndexMatrixView
+#include "../predict_fn.h"                 // for GetTreeLimit
+#include "../treeshap.h"                   // for CalculateContributions
+#include "dmlc/omp.h"                      // for omp_get_thread_num
+#include "xgboost/base.h"                  // for bst_omp_uint
+#include "xgboost/logging.h"               // for CHECK
+#include "xgboost/tree_model.h"            // for MTNotImplemented
 
 namespace xgboost::interpretability {
 namespace {

src/predictor/interpretability/shap.cu

@@ -25,7 +25,6 @@
 #include "../../common/cuda_context.cuh"  // for CUDAContext
 #include "../../common/cuda_rt_utils.h"   // for SetDevice
 #include "../../common/device_helpers.cuh"
-#include "../../common/error_msg.h"
 #include "../../common/nvtx_utils.h"
 #include "../../data/batch_utils.h"      // for StaticBatch
 #include "../../data/cat_container.cuh"  // for EncPolicy, MakeCatAccessor

tests/cpp/collective/test_worker.h

@@ -1,17 +1,18 @@
 /**
- * Copyright 2023-2024, XGBoost Contributors
+ * Copyright 2023-2026, XGBoost Contributors
  */
 #pragma once
 #include <gtest/gtest.h>
 #include <xgboost/global_config.h>  // for InitNewThread
 
-#include <chrono>   // for seconds
-#include <cstdint>  // for int32_t
-#include <fstream>  // for ifstream
-#include <string>   // for string
-#include <thread>   // for thread
-#include <utility>  // for move
-#include <vector>   // for vector
+#include <algorithm>  // for max
+#include <chrono>     // for seconds
+#include <cstdint>    // for int32_t
+#include <fstream>    // for ifstream
+#include <string>     // for string
+#include <thread>     // for thread
+#include <utility>    // for move
+#include <vector>     // for vector
 
 #include "../../../src/collective/comm.h"              // for RabitComm
 #include "../../../src/collective/communicator-inl.h"  // for Init, Finalize
@@ -42,7 +43,7 @@ class WorkerForTest {
         tracker_port_{port},
         world_size_{world},
         task_id_{"t:" + std::to_string(rank)},
-        comm_{tracker_host_, tracker_port_, timeout, retry_, task_id_, DefaultNcclName()} {
+        comm_{tracker_host_, tracker_port_, timeout, retry_, task_id_, DefaultNcclName(), 0} {
     CHECK_EQ(world_size_, comm_.World());
   }
   virtual ~WorkerForTest() noexcept(false) { SafeColl(comm_.Shutdown()); }

tests/python/test_tracker.py

@@ -8,6 +8,7 @@
 from hypothesis import HealthCheck, given, settings, strategies
 from xgboost import RabitTracker, collective
 from xgboost import testing as tm
+from xgboost.testing.collective import get_avail_port
 
 
 def test_rabit_tracker() -> None:
@@ -53,6 +54,32 @@ def test_socket_error() -> None:
         tracker.free()
 
 
+@pytest.mark.skipif(**tm.no_loky())
+def test_worker_port() -> None:
+    from loky import get_reusable_executor
+
+    n_workers = 4
+
+    tracker = RabitTracker(host_ip="127.0.0.1", n_workers=n_workers)
+    tracker.start()
+    args = tracker.worker_args()
+
+    def local_test(worker_id: int, rabit_args: dict) -> int:
+        cfg = collective.Config(worker_port=get_avail_port)
+        cfg.update_worker_args(rabit_args)
+        with collective.CommunicatorContext(**rabit_args):
+            a = np.array([1])
+            result = collective.allreduce(a, collective.Op.SUM)
+            assert result[0] == n_workers
+
+            return 1
+
+    fn = update_wrapper(partial(local_test, rabit_args=args), local_test)
+    with get_reusable_executor(max_workers=n_workers) as pool:
+        results = pool.map(fn, range(n_workers))
+        assert sum(results) == n_workers
+
+
 def run_rabit_ops(pool, n_workers: int, address: str) -> None:
     tracker = RabitTracker(host_ip=address, n_workers=n_workers)
     tracker.start()