XGBoost 分布式训练上 Kubernetes：基于 Kubeflow Trainer 的多节点训练完整指南

XGBoost 分布式训练上 Kubernetes：基于 Kubeflow Trainer 的多节点训练完整指南

【免费下载链接】xgboostScalable, Portable and Distributed Gradient Boosting (GBDT, GBRT or GBM) Library, for Python, R, Java, Scala, C and more. Runs on single machine, Hadoop, Spark, Dask, Flink and DataFlow项目地址: https://gitcode.com/gh_mirrors/xg/xgboost

本指南以 XGBoost 官方教程 doc/tutorials/kubernetes.rst 为核心，系统讲解如何借助 Kubeflow Trainer 的xgboost-distributed运行时，在 Kubernetes 集群上编排多节点分布式 XGBoost 训练任务。读完本文，你将掌握分布式训练的架构原理（Collective 协议与 Rabit Tracker）、TrainJob与ClusterTrainingRuntime两种资源的配置方法、Python SDK 与kubectl两种作业提交方式，以及内存优化、早停、断点续训、日志治理等一系列生产级最佳实践。

概述：XGBoost 如何在多节点上协同训练

XGBoost 原生支持通过Collective通信协议（历史上称为 Rabit）进行分布式训练。在分布式场景下，多个 worker 进程各自持有数据集的一个分片（shard），并通过 AllReduce 同步直方图（histogram）的统计信息，从而在所有 worker 上达成一致的树分裂决策。Kubeflow Trainer 的 XGBoost 运行时（runtime）将这一过程在 Kubernetes 上自动化，具体职责包括：

• 将 worker Pod 以JobSet的形式部署；
• 自动注入 XGBoost Collective 通信层所需的DMLC_*环境变量；
• 向 rank-0 Pod 提供 tracker 地址，使你的训练代码能够启动RabitTracker协调各 worker；
• 同时支持 CPU 与 GPU 训练负载。

从仓库源码可以看到 Collective 通信层的具体实现：CommunicatorContext在进入上下文时调用init()、退出时调用finalize()，并对外暴露get_rank()、get_world_size()、broadcast()、communicator_print()等分布式原语（见 collective.py）。RabitTracker类则封装了 tracker 的创建、启动与等待逻辑（见 tracker.py），本文后面会逐一用到它们。

架构：四大组件与作业流转

Kubernetes 上的分布式 XGBoost 训练架构由以下组件构成：

1. TrainJob：Kubernetes 自定义资源，声明训练作业的配置（节点数、每节点资源、训练代码）。
2. ClusterTrainingRuntime：集群级资源，定义 XGBoost 运行时模板（容器镜像、ML 策略、默认设置）。内置运行时名为xgboost-distributed。
3. Trainer Controller：将TrainJob与引用的运行时进行解析，执行 XGBoost ML 策略（注入环境变量），并创建底层的JobSet。
4. Worker Pods：每个 Pod 运行相同的训练脚本。rank-0 Pod 上的用户训练代码负责启动RabitTracker用于协调。

整个作业流转过程可以用下面的图表示：

┌─────────────────────────────────────────────────────────────────┐ │ User submits TrainJob (SDK or kubectl) │ └──────────────────────────┬──────────────────────────────────────┘ │ ▼ ┌─────────────────────────────────────────────────────────────────┐ │ Trainer Controller │ │ • Resolves ClusterTrainingRuntime (xgboost-distributed) │ │ • Enforces XGBoost MLPolicy (injects DMLC_* env vars) │ │ • Creates JobSet with worker pods │ └──────────────────────────┬──────────────────────────────────────┘ │ ▼ ┌─────────────────────────────────────────────────────────────────┐ │ Kubernetes Cluster (Headless Service) │ │ │ │ ┌────────────────┐ ┌──────────┐ ┌──────────┐ │ │ │ Pod: node-0-0 │ │ node-0-1 │ │ node-0-2 │ ... │ │ │ TASK_ID=0 │ │ TASK_ID=1│ │ TASK_ID=2│ │ │ │ (Tracker) │ │ (Worker) │ │ (Worker) │ │ │ └───────┬────────┘ └────┬─────┘ └────┬─────┘ │ │ │ │ │ │ │ └──── Collective Protocol ───────┘ │ └─────────────────────────────────────────────────────────────────┘

环境变量：运行时自动注入的DMLC_*

XGBoost 运行时插件会自动向每个 worker Pod 注入以下环境变量。它们是 XGBoost Collective 协议的原生变量：

这些环境变量是保留变量，用户不能在TrainJobspec 中手动设置。运行时插件会做校验，任何试图覆盖它们的TrainJob都会被拒绝。

从源码看，xgboost.collective.init()正是通过dmlc_tracker_uri、dmlc_tracker_port、dmlc_task_id等参数初始化通信组（见 collective.py 中init的参数说明），这与你从环境变量中读到的值一一对应。

Worker 数量计算：DMLC_NUM_WORKER如何确定

worker 总数（DMLC_NUM_WORKER）的计算公式为：

DMLC_NUM_WORKER = numNodes × workersPerNode

其中workersPerNode由训练类型决定：

• CPU 训练：每节点 1 个 worker。XGBoost不会为 CPU 训练派生多个 worker 进程，而是由单个 worker 进程使用 OpenMP 在节点所有可用 CPU 核上并行构建树（直方图构建、分裂评估等）。也就是说，如果一个 Pod 有 8 个 CPU 核，那么 1 个 XGBoost worker 会使用全部 8 核做进程内并行。

  线程数可以通过 Booster 参数nthread控制：

  # 默认情况下，XGBoost 使用所有可用的 CPU 核。 # 设置 nthread 可限制每个 worker 的 OpenMP 线程数。 params = { "objective": "binary:logistic", "nthread": 4, # 只使用可用核中的 4 个 "tree_method": "hist", }

  DMatrix构造函数中的nthread参数控制数据加载阶段的并行度，而 Booster 参数中的nthread控制训练阶段的并行度。若都不设置，两者默认取机器上可用的最大线程数。

  提示：在TrainJob中设置resourcesPerNode的 CPU requests 时，请将nthread与 CPU requests 对齐以避免超额订阅（over-subscription）。例如，如果你申请了cpu: "4"，就在训练参数中设置"nthread": 4。

• GPU 训练：每 GPU 1 个 worker。GPU 数量从TrainJob或运行时模板的resourcesPerNodelimits 中推导。分布式环境下请使用device="cuda"（而非"cuda:<ordinal>"）；GPU ordinal 的选择由分布式框架处理，指定 ordinal 会报错。

常见配置下的 worker 数量对照：

前置条件

在 Kubernetes 上运行分布式 XGBoost 作业前，请确保满足以下条件：

1. Kubernetes 集群：一个正在运行的 Kubernetes 集群（v1.27+）。可以使用kind、minikube，或托管 Kubernetes 服务（GKE、EKS、AKS）。

2. kubectl：Kubernetes CLI 工具，并已配置连接你的集群。

3. Kubeflow Trainer：在集群上安装 Kubeflow Trainer 及其依赖（JobSet）。安装控制面（包含 JobSet）：

   kubectl apply --server-side -k "github.com/kubeflow/trainer/manifests/overlays/standalone"

4. Kubeflow Python SDK（可选，用于以编程方式提交作业）：

   pip install kubeflow

5. GPU 支持（可选，用于 GPU 训练）：确保集群上安装了 NVIDIA GPU Operator 或等效的 device plugin。

验证安装

安装 Kubeflow Trainer 后，验证 XGBoost 运行时是否可用：

kubectl get clustertrainingruntime

你应该能看到xgboost-distributed运行时：

NAME AGE xgboost-distributed 1m

XGBoost ClusterTrainingRuntime

xgboost-distributedClusterTrainingRuntime随 Kubeflow Trainer 安装一起部署，它定义了默认的 XGBoost 运行时模板：

apiVersion: trainer.kubeflow.org/v1alpha1 kind: ClusterTrainingRuntime metadata: name: xgboost-distributed labels: trainer.kubeflow.org/framework: xgboost spec: mlPolicy: numNodes: 1 xgboost: {} template: spec: replicatedJobs: - name: node template: metadata: labels: trainer.kubeflow.org/trainjob-ancestor-step: trainer spec: template: spec: containers: - name: node image: ghcr.io/kubeflow/trainer/xgboost-runtime:latest

关键点：

• mlPolicy.xgboost: {}激活 XGBoost 运行时插件，该插件负责注入DMLC_*环境变量；
• numNodes默认为1，可被每个TrainJob覆盖；
• 容器镜像ghcr.io/kubeflow/trainer/xgboost-runtime:latest基于nvidia/cuda:12.4.0-runtime-ubuntu22.04，包含支持 CUDA 12 的 XGBoost 3.0.2、NumPy 和 scikit-learn。

实战示例：分布式 XGBoost 训练

本节演示两种运行分布式 XGBoost 训练的方式：使用 Python SDK（推荐用于交互式使用）和使用kubectl+ YAML 清单。

方式一：Python SDK

Kubeflow Python SDK 提供了TrainerClient，可以编程方式简化作业的提交与管理。

第 1 步：定义训练函数

编写将在每个 worker 节点上被序列化并执行的训练函数。DMLC_*环境变量由运行时自动注入：

def xgboost_train_classification(): """ Distributed XGBoost training function using the Collective API. DMLC_* env vars are injected by the Kubeflow Trainer XGBoost plugin: - DMLC_TRACKER_URI: DNS name of the rank-0 pod running the tracker - DMLC_TRACKER_PORT: Port for tracker communication (default: 29500) - DMLC_TASK_ID: Worker rank (0, 1, 2, ...) - DMLC_NUM_WORKER: Total number of workers """ import os import xgboost as xgb from xgboost import collective as coll from xgboost.tracker import RabitTracker from sklearn.datasets import make_classification from sklearn.model_selection import train_test_split from sklearn.metrics import accuracy_score # Read injected environment variables. rank = int(os.environ["DMLC_TASK_ID"]) world_size = int(os.environ["DMLC_NUM_WORKER"]) tracker_uri = os.environ["DMLC_TRACKER_URI"] tracker_port = int(os.environ["DMLC_TRACKER_PORT"]) # Rank 0 starts the Rabit tracker (required for coordination). tracker = None if rank == 0: tracker = RabitTracker( host_ip="0.0.0.0", n_workers=world_size, port=tracker_port ) tracker.start() # All workers connect to the tracker via the Collective communicator. with coll.CommunicatorContext( dmlc_tracker_uri=tracker_uri, dmlc_tracker_port=tracker_port, dmlc_task_id=str(rank), ): # Generate synthetic classification data. # In practice, each worker would load its own data shard. X, y = make_classification( n_samples=10000, n_features=20, n_informative=10, n_classes=2, random_state=42 + rank, ) X_train, X_valid, y_train, y_valid = train_test_split( X, y, test_size=0.2, random_state=42, ) # NOTE: DMatrix construction MUST be inside the communicator context # because it involves cross-worker synchronization for quantization. # # Use QuantileDMatrix instead of DMatrix for the hist tree method # (the default). QuantileDMatrix quantizes data on-the-fly, avoiding # an intermediate dense copy and significantly reducing memory usage. dtrain = xgb.QuantileDMatrix(X_train, label=y_train) # Validation QuantileDMatrix must reference the training matrix # so that the same quantile bins are reused. dvalid = xgb.QuantileDMatrix(X_valid, label=y_valid, ref=dtrain) # Training parameters. params = { "objective": "binary:logistic", "max_depth": 6, "eta": 0.1, "eval_metric": "logloss", } # Distributed training - workers synchronize histogram stats via collective ops. # early_stopping_rounds activates early stopping based on the validation metric. # verbose_eval=10 prints evaluation results every 10 rounds (rank 0 only). model = xgb.train( params, dtrain, num_boost_round=100, evals=[(dvalid, "validation")], early_stopping_rounds=10, verbose_eval=10, ) # Note: early_stopping_rounds returns the *last* model, not the best. # Use bst.best_iteration to slice the model to the best round. if hasattr(model, "best_iteration"): model = model[: model.best_iteration + 1] # Evaluate on validation set. preds = model.predict(dvalid) predictions = [1 if p > 0.5 else 0 for p in preds] accuracy = accuracy_score(y_valid, predictions) # Only perform logging and model saving from rank 0 # to avoid duplicate output and file write conflicts. if coll.get_rank() == 0: print(f"Validation Accuracy: {accuracy:.4f}") model.save_model("/workspace/xgboost_model.json") print("Model saved to /workspace/xgboost_model.json") # Wait for tracker to finish (rank 0 only). if tracker is not None: tracker.wait_for()

这里有几个值得注意的源码细节：

• RabitTracker.start()启动 tracker 服务，而wait_for()会阻塞等待 tracker 完成全部工作后关闭（见 tracker.py），因此必须在所有 worker 训练结束后再调用；
• coll.CommunicatorContext进入时调用init()、退出时调用finalize()，且get_rank()/get_world_size()只在上下文内有效（见 collective.py），这就是为何数据矩阵构造、训练与模型保存都必须放在上下文内部的根本原因。

第 2 步：提交训练作业

使用TrainerClient将训练函数作为分布式作业提交：

from kubeflow.trainer import CustomTrainer, TrainerClient client = TrainerClient() # Submit a distributed XGBoost training job on 3 nodes. job_name = client.train( trainer=CustomTrainer( func=xgboost_train_classification, num_nodes=3, resources_per_node={"cpu": 3}, ), runtime="xgboost-distributed", ) print(f"TrainJob '{job_name}' submitted")

GPU 训练时，在资源中指定 GPU：

job_name = client.train( trainer=CustomTrainer( func=xgboost_train_classification, num_nodes=2, resources_per_node={ "cpu": 4, "gpu": 4, # 4 GPUs per node → 8 total workers }, ), runtime="xgboost-distributed", )

注意：GPU 训练时，请在训练函数的 XGBoostparams字典中加入"device": "cuda"。

第 3 步：监控训练作业

检查作业状态并查看日志：

# Wait for the job to start running. client.wait_for_job_status(name=job_name, status={"Running"}) # Check the steps (one per worker node). for step in client.get_job(name=job_name).steps: print(f"Step: {step.name}, Status: {step.status}") # Stream logs from each worker node. num_nodes = 3 for i in range(num_nodes): logs = client.get_job_logs(name=job_name, follow=True, step=f"node-{i}") print(f"\n=== Node {i} ===") print("\n".join(logs))

第 4 步：清理

训练结束后删除作业：

client.delete_job(job_name)

方式二：kubectl + YAML

你也可以直接用kubectl创建TrainJob资源。

CPU 训练示例

下面的 YAML 创建一个 4 个 worker 节点的分布式 XGBoost 训练作业：

apiVersion: trainer.kubeflow.org/v1alpha1 kind: TrainJob metadata: name: xgboost-cpu-example spec: runtimeRef: name: xgboost-distributed trainer: image: ghcr.io/kubeflow/trainer/xgboost-runtime:latest command: - python - train.py numNodes: 4 resourcesPerNode: requests: cpu: "4" memory: "8Gi"

应用清单：

kubectl apply -f xgboost-cpu-trainjob.yaml

GPU 训练示例

多节点 GPU 训练时，通过resourcesPerNode指定 GPU 资源：

apiVersion: trainer.kubeflow.org/v1alpha1 kind: TrainJob metadata: name: xgboost-gpu-example spec: runtimeRef: name: xgboost-distributed trainer: image: ghcr.io/kubeflow/trainer/xgboost-runtime:latest command: - python - train.py numNodes: 2 resourcesPerNode: limits: nvidia.com/gpu: "4" requests: cpu: "4" memory: "16Gi"

使用该配置，运行时计算出的DMLC_NUM_WORKER = 2 节点 × 4 GPU = 8。每个 GPU 运行一个 XGBoost worker 进程。

使用 kubectl 监控

# 查看 TrainJob 状态 kubectl get trainjob xgboost-cpu-example # 查看某个 worker Pod 的日志 kubectl logs xgboost-cpu-example-node-0-0 # 删除 TrainJob kubectl delete trainjob xgboost-cpu-example

工作原理：运行时插件内部机制

XGBoost 运行时插件

XGBoost 运行时以 Go 插件的形式实现于 Kubeflow Trainer controller 中（见 Trainer 仓库的pkg/runtime/framework/plugins/xgboost/）。它实现了两个接口：

• EnforceMLPolicyPlugin：注入DMLC_*环境变量（见上文「环境变量」小节），并暴露容器端口29500；
• CustomValidationPlugin：拒绝任何手动设置保留DMLC_*环境变量的TrainJob。

Tracker 发现机制

worker 通过 Kubernetes headless service 发现 rank-0 上的RabitTracker。DMLC_TRACKER_URI的构造规则为：

<trainjob-name>-node-0-0.<trainjob-name>

例如，一个名为myjob、含 4 个节点的TrainJob会创建如下 Pod：

myjob-node-0-0 DMLC_TASK_ID=0 (Tracker + Worker) myjob-node-0-1 DMLC_TASK_ID=1 (Worker) myjob-node-0-2 DMLC_TASK_ID=2 (Worker) myjob-node-0-3 DMLC_TASK_ID=3 (Worker)

注意：启动 tracker 是用户的责任。运行时只负责注入环境变量，rank-0 上的训练代码必须在其他 worker 连接之前调用RabitTracker(...).start()。

最佳实践

使用 QuantileDMatrix 降低内存占用

默认树方法为hist（tree_method="auto"会解析为hist）。使用hist时，推荐用xgboost.QuantileDMatrix而非xgboost.DMatrix。QuantileDMatrix直接从输入生成分位数化数据，跳过中间稠密表示，显著降低内存消耗（类定义见 core.py）：

# 标准 DMatrix —— 先加载数据再分位数化（峰值内存更高） dtrain = xgb.DMatrix(X_train, label=y_train) # QuantileDMatrix —— 边加载边分位数化（峰值内存更低） dtrain = xgb.QuantileDMatrix(X_train, label=y_train)

构造验证集QuantileDMatrix时，务必把训练矩阵作为ref传入，让 XGBoost 复用相同的分位数桶。验证数据省略ref可能导致分位数化不一致，进而降低模型质量：

dtrain = xgb.QuantileDMatrix(X_train, label=y_train) dvalid = xgb.QuantileDMatrix(X_valid, label=y_valid, ref=dtrain) # 正确

注意：QuantileDMatrix自 XGBoost 1.7.0 引入。无需显式指定tree_method—— 默认的auto已经使用hist。

早停（Early Stopping）

向xgboost.train传入early_stopping_rounds即可激活早停。它要求evals中至少有一个验证集。当验证指标在连续指定轮数内不再提升时，训练停止：

model = xgb.train( params, dtrain, num_boost_round=500, evals=[(dvalid, "validation")], early_stopping_rounds=10, )

早停在分布式模式下同样正确工作——验证指标已经通过 collective 协议在 worker 之间同步。

重要：带early_stopping_rounds的xgb.train返回的是最后一个模型，而不是最优模型。要拿到最优模型，请使用模型切片：

# 训练后，仅保留到最佳迭代轮的轮次 if hasattr(model, "best_iteration"): model = model[: model.best_iteration + 1]

或者，直接使用xgboost.callback.EarlyStopping回调并设置save_best=True，自动只保留最优模型：

from xgboost.callback import EarlyStopping model = xgb.train( params, dtrain, num_boost_round=500, evals=[(dvalid, "validation")], callbacks=[EarlyStopping(rounds=10, save_best=True)], ) # model 现在只包含到最佳迭代轮的轮次

当evals提供了多个评估数据集时，使用最后一个条目进行早停；当指定了多个eval_metric时，使用最后一个指标。

分布式模式下的日志

分布式训练中，print()会在每个 worker 上执行，产生重复日志行。若要只从单个 worker 打印，用 rank 检查守卫：

from xgboost import collective as coll with coll.CommunicatorContext(...): # 只从 rank 0 打印 if coll.get_rank() == 0: print(f"Training complete, best score: {model.best_score}")

xgboost.collective.communicator_print是另一个选择，它把消息经 tracker 转发而非输出到 stdout。注意它不会按 rank 过滤——任何调用它的 worker 的消息都会被 tracker 打印。它主要用于内部场景（例如verbose_eval，其自带 rank-0 守卫，通过xgboost.callback.EvaluationMonitor实现，见 callback.py）。

生产环境设置 verbose_eval

在分布式 Kubernetes 作业中，把verbose_eval设为整数而非True，以降低日志量：

model = xgb.train( params, dtrain, num_boost_round=500, evals=[(dvalid, "validation")], verbose_eval=50, # 每 50 轮打印一次，而不是每轮 )

断点续训（Checkpointing）

XGBoost 提供了xgboost.callback.TrainingCheckPoint回调（见 callback.py），在训练过程中周期性地保存模型快照。该回调只会从 rank 0 保存，避免多个 worker 写入同一路径：

from xgboost.callback import TrainingCheckPoint model = xgb.train( params, dtrain, num_boost_round=500, evals=[(dvalid, "validation")], callbacks=[ TrainingCheckPoint( directory="/workspace/checkpoints", name="xgb_model", interval=50, # 每 50 轮保存一次 ), ], )

警告：XGBoost 不处理分布式文件系统。directory路径必须能从 rank-0 Pod 写入——例如，挂载到 Pod 中的 Kubernetes PersistentVolumeClaim。

从检查点恢复训练时，通过xgb_model传入已保存的模型文件：

model = xgb.train( params, dtrain, num_boost_round=500, xgb_model="/workspace/checkpoints/xgb_model_200.ubj", # 从第 200 轮恢复 evals=[(dvalid, "validation")], )

数据划分

默认情况下，分布式 XGBoost 作业中每个 worker 持有数据行的不同子集（水平划分），且只支持按行的数据拆分。这种模式下，每个 worker 加载自己的数据分片：

with coll.CommunicatorContext(...): # 每个 worker 根据其 rank 加载不同的数据分片 rank = coll.get_rank() X_shard, y_shard = load_data_shard(rank) dtrain = xgb.QuantileDMatrix(X_shard, label=y_shard)

按列拆分已经被移除；分布式训练使用按行划分。

与 rank 相关的逻辑

在通信上下文内使用xgboost.collective.get_rank和xgboost.collective.get_world_size执行与 rank 相关的操作：

with coll.CommunicatorContext(...): if coll.get_rank() == 0: model.save_model("/workspace/model.json") # 如需向所有 worker 广播结果 results = coll.broadcast(results, root=0)

xgboost.collective.broadcast可以把任意可 pickle 的 Python 对象从一个 worker 广播到所有其他 worker。这在共享 rank 0 上计算的预处理元数据（如标签编码器、特征名列表）时非常有用。

常见问题与边界情况

保留环境变量

运行时插件会拒绝任何手动设置保留DMLC_*环境变量（DMLC_TRACKER_URI、DMLC_TRACKER_PORT、DMLC_TASK_ID、DMLC_NUM_WORKER）的TrainJob。如果你在spec.trainer.env中设置了其中任何一个，webhook 会返回Forbidden错误：

spec.trainer.env[0]: Forbidden: DMLC_TRACKER_URI is reserved for the XGBoost runtime

请从TrainJobspec 中移除这些保留变量，让运行时自动注入。

trainer 为空时不注入环境变量

如果TrainJob不包含spec.trainer段，XGBoost 插件会完全跳过环境变量注入。DMLC_*变量只有在spec.trainer存在且运行时能在 Pod 模板中找到node容器时才会被注入。请确保TrainJob包含trainer字段。

资源优先级：TrainJob 覆盖运行时

当 GPU 资源同时出现在ClusterTrainingRuntime模板和TrainJob.spec.trainer.resourcesPerNode中时，TrainJob的值优先。这会直接影响workersPerNode的计算：

Runtime template: nvidia.com/gpu: 1 → workersPerNode = 1 TrainJob override: nvidia.com/gpu: 3 → workersPerNode = 3 (this wins)

如果两者都没有指定 GPU 资源，workersPerNode默认为1（CPU 模式）。

分布式模式下的 GPU 设备序号

分布式训练中，不要在 XGBoost 参数中使用device="cuda:0"或任何具体的 GPU ordinal。GPU 设备分配由 Kubernetes device plugin 和分布式框架处理。请使用device="cuda"：

# 正确 params = {"device": "cuda", "tree_method": "hist"} # 错误 —— 分布式模式下会报错 params = {"device": "cuda:0", "tree_method": "hist"}

数据矩阵必须位于 CommunicatorContext 内

在CommunicatorContext外部构造xgb.DMatrix或xgb.QuantileDMatrix，对稠密数据可能看起来正常，但行为是未定义的。构造函数会执行跨 worker 同步，用于数据形状校验和分位数素描（tree_method="hist"所需）。务必在上下文内部构造数据矩阵：

# 错误 —— 数据矩阵在上下文之外 dtrain = xgb.QuantileDMatrix(X_train, label=y_train) with coll.CommunicatorContext(...): model = xgb.train(params, dtrain, ...) # 未定义行为 # 正确 —— 数据矩阵在上下文之内 with coll.CommunicatorContext(...): dtrain = xgb.QuantileDMatrix(X_train, label=y_train) model = xgb.train(params, dtrain, ...)

单节点默认值

如果TrainJob未指定numNodes，运行时使用ClusterTrainingRuntime中的默认值（xgboost-distributed运行时默认为1）。单节点作业仍会走完整的运行时流水线——RabitTracker在 rank-0（即唯一 Pod）上启动，DMLC_NUM_WORKER设为1。这在本地扩展前测试训练函数时很有用。

CPU 超额订阅

默认情况下，XGBoost 通过 OpenMP 使用所有可用 CPU 核。在 Kubernetes Pod 中，“可用核数”由容器运行时设置的 cgroup 限制决定。如果你的 Pod 只指定了 CPUrequests（没有limits），cgroup 可能不会限制 CPU 使用，XGBoost 可能会尝试使用节点上的所有核，导致与其他 Pod 争抢资源。

要避免这种情况，可以：

• 在 XGBoost 参数中设置nthread，使其与 CPU request 匹配；
• 在resourcesPerNode中设置 CPUlimits（而不只是 requests），让容器运行时强制 cgroup 上限。

# 同时设置 requests 和 limits，确保 XGBoost 看到正确的核数 resourcesPerNode: requests: cpu: "4" limits: cpu: "4"

小结

通过 Kubeflow Trainer 的xgboost-distributed运行时，XGBoost 的分布式训练能力被完整地映射到了 Kubernetes 生态：TrainJob声明作业意图，ClusterTrainingRuntime提供运行时模板，Trainer Controller 注入DMLC_*环境变量并创建 JobSet，而用户代码只需在 rank-0 上启动RabitTracker并在CommunicatorContext内完成数据构造与训练。无论是 4 节点的 CPU 作业，还是 2 节点 8 GPU 的 GPU 作业，遵循本文的配置规则、worker 数量计算方法和最佳实践，你都能在 Kubernetes 上稳定、高效地跑通多节点分布式 XGBoost 训练。

【免费下载链接】xgboostScalable, Portable and Distributed Gradient Boosting (GBDT, GBRT or GBM) Library, for Python, R, Java, Scala, C and more. Runs on single machine, Hadoop, Spark, Dask, Flink and DataFlow项目地址: https://gitcode.com/gh_mirrors/xg/xgboost