Merge remote-tracking branch 'upstream/master' into quantile-refactor

Author: RAMitchell

doc/treemethod.rst

@@ -8,7 +8,7 @@ namely ``updater`` and ``tree_method``.  XGBoost has 3 builtin tree methods, nam
 free standing updaters including ``refresh``, ``prune`` and ``sync``.  The parameter
 ``updater`` is more primitive than ``tree_method`` as the latter is just a
 pre-configuration of the former.  The difference is mostly due to historical reasons that
-each updater requires some specific configurations and might has missing features.  As we
+each updater requires some specific configurations and might have missing features.  As we
 are moving forward, the gap between them is becoming more and more irrelevant.  We will
 collectively document them under tree methods.
 

include/xgboost/predictor.h

@@ -107,7 +107,8 @@ class Predictor {
    */
   virtual void PredictBatch(DMatrix* dmat, PredictionCacheEntry* out_preds,
                             gbm::GBTreeModel const& model, bst_tree_t tree_begin,
-                            bst_tree_t tree_end = 0) const = 0;
+                            bst_tree_t tree_end = 0,
+                            std::vector<float> const* tree_weights = nullptr) const = 0;
 
   /**
    * \brief Inplace prediction.

plugin/sycl/predictor/predictor.cc

@@ -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/collective.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/dask/__init__.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

python-package/xgboost/spark/core.py

@@ -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.cc

@@ -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.h

@@ -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/collective/comm_group.cc

@@ -187,10 +187,8 @@ class EllpackFormatPolicy {
  public:
   EllpackFormatPolicy() {
     StringView msg{" The overhead of iterating through external memory might be significant."};
-    if (!has_hmm_) {
+    if (!(has_hmm_ || curt::SupportsAts())) {
       LOG(WARNING) << "CUDA heterogeneous memory management is not available." << msg;
-    } else if (!curt::SupportsAts()) {
-      LOG(WARNING) << "CUDA address translation service is not available." << msg;
     }
     if (!(GlobalConfigThreadLocalStore::Get()->use_rmm ||
           GlobalConfigThreadLocalStore::Get()->use_cuda_async_pool)) {