SageMaker上XGBoost自定义脚本分布式训练启用Pipe模式的问题求助

一、解决分布式训练主机未利用的警告 #

该警告说明分布式集群未正确初始化，核心原因是Pipe模式下自定义脚本未适配SageMaker的分布式环境配置，需做以下修改：

1. 读取SageMaker分布式环境变量 #

SageMaker会自动注入分布式训练相关环境变量，脚本中必须解析这些变量来配置XGBoost的分布式参数：

import os
import json

# 解析分布式集群信息
hosts = json.loads(os.environ['SM_HOSTS'])
current_host = os.environ['SM_CURRENT_HOST']
rank = hosts.index(current_host)
world_size = len(hosts)

2. 配置XGBoost分布式训练参数 #

初始化XGBoost训练器时，必须传入分布式相关配置，确保集群节点能互相通信：

import xgboost as xgb

# 定义训练参数，重点配置分布式部分
params = {
    'objective': 'binary:logistic',
    'eval_metric': 'auc',
    'tree_method': 'hist',  # Pipe模式推荐使用hist/gpu_hist，避免内存溢出
    'nthread': os.cpu_count(),
    'distributed_training': {
        'enabled': True,
        'master_addr': hosts[0],
        'master_port': '6000',
        'rank': rank,
        'world_size': world_size
    }
}

# 加载数据并启动训练
dtrain = xgb.DMatrix(train_data, label=train_labels)
model = xgb.train(params, dtrain, num_boost_round=100)

3. 适配Pipe模式的数据读取逻辑 #

Pipe模式下，SageMaker通过命名管道传递数据，而非直接挂载S3文件，需修改数据读取代码：

import pyarrow.parquet as pq

# 从SageMaker指定的管道路径读取Parquet数据
train_pipe_path = os.environ['SM_CHANNEL_TRAIN']
with open(train_pipe_path, 'rb') as f:
    parquet_file = pq.ParquetFile(f)
    train_data = parquet_file.read().to_pandas()

注：Pipe模式下SageMaker会自动完成数据分片，无需手动拆分数据。

二、解决模型未上传至S3的问题 #

模型未上传的核心原因是未将模型保存到SageMaker指定的输出目录，需做以下调整：

1. 保存模型到指定目录 #

训练完成后，必须将模型保存到SM_MODEL_DIR环境变量指向的路径，SageMaker会自动将该目录内容同步到S3：

import os

# 保存XGBoost原生模型
model.save_model(os.path.join(os.environ['SM_MODEL_DIR'], 'xgboost-model'))

# 若使用joblib保存模型，同样需指定该路径
import joblib
joblib.dump(model, os.path.join(os.environ['SM_MODEL_DIR'], 'model.joblib'))

2. 检查训练作业配置 #

创建SageMaker训练作业时，需确保指定了正确的S3输出路径，且IAM角色拥有该路径的写入权限：

from sagemaker.xgboost import XGBoost

estimator = XGBoost(
    entry_point='your_train_script.py',
    role='your_sagemaker_execution_role',
    instance_count=2,  # 分布式训练需至少2个节点
    instance_type='ml.m5.xlarge',
    input_mode='Pipe',  # 明确启用Pipe模式
    output_path='s3://your-bucket/sagemaker-model-output/',  # 合法的S3输出路径
    hyperparameters={
        'num_round': 100,
        'objective': 'binary:logistic'
    }
)

三、Pipe模式启用的关键修改总结 #

• 数据读取：从SM_CHANNEL_*指定的管道路径读取，替代本地挂载文件的读取方式
• 分布式配置：利用SageMaker注入的环境变量初始化XGBoost的rank、world_size等参数
• 模型保存：必须将模型写入SM_MODEL_DIR目录
• 作业配置：明确设置input_mode='Pipe'，且实例数≥2

内容的提问来源于stack exchange，提问作者Arvs