技术实践干货 | 初探大规模 GBDT 训练

本文作者： 字节，观远数据首席科学家。主导多个AI项目在世界500强的应用落地，多次斩获智能零售方向Hackathon冠军。曾就职于微策略，阿里云，拥有十多年的行业经验。

本文是此前评估在 Spark 上做大规模 GBDT 训练时写的一篇入门级教程与框架评估。目前市面上似乎没有多少使用 Spark 来跑 GBDT 的分享，故分享出来看看是否有做过类似场景的同学可以一道交流。

背景

在服务一些客户做商业问题的机器学习建模时，我们会碰到不少拥有非常大量数据且对模型 pipeline 运行有一定要求的情况。相比直接的单机 Python 建模，这类项目有一些难点：

1. 数据量大。 由于预测粒度较细，导致历史数据量非常巨大。一些场景的 pilot 项目中已经达到近千万级别的训练数据量，后续拓展到整个业务线，数据量会超过十亿甚至百亿行级别。

2. 整体流程运行时间有一定要求。 一般模型所依赖的上游数据会在半夜开始通过一系列 ETL 任务从业务系统导入到 Hive 数仓中，大约在凌晨 3 点后，各类预测所需数据会准备就绪。接下来运行整个取数，清洗，特征，训练，预测，业务系统对接全流程，需要在早上 8 点前完成并下发到业务系统中，整体运行时间必须控制在 5 小时以内。

3. 对监控运维等方面的高要求。 海量数据细粒度的预测，覆盖非常多业务人员的日常工作需求，因而也会受到更多的审视与挑战。如何确保模型预测输出的稳定性，在业务反馈问题后又如何快速定位排查，遇到数据不可用，服务器 down 机等突发情况，有什么样的备选方案确保整体流程的稳定运行，都是需要考虑的问题。

从前两点来看，我们之前习惯的单机 Pandas + lgb/xgb 建模思路已经难以适用（除非搞台神威·太湖之光之类的机器），所以我们需要引入目前大数据界的当红炸子鸡 – Spark 来协助完成此类项目。

部署 Spark

要玩 Spark，第一步是部署。如果是本机测试运行，一般跑一个 pip install pyspark 就能把一个 local 节点跑起来了，非常的方便。如果想部署一个相对完整一点的 standalone 集群，可以参考以下步骤：

1. 到 Spark 官方网站[1] 下载 Spark。当时最新的稳定版本是 2.4.5，下载 pre-built for Apache Hadoop 2.7 就好。
2. 解压缩，做一些简单的配置文件配置。在解压开的 spark 目录下，进入到 conf 目录里，会看到一系列配置的 template。把需要自定义的配置 copy 一份，例如：cp spark-env.sh.template spark-env.sh，然后进行编辑。我改的一些配置具体如下：

spark-env.sh
# 配置 master 和 worker

SPARK_MASTER_HOST=0.0.0.0
SPARK_DAEMON_MEMORY=4g
SPARK_WORKER_CORES=6
SPARK_WORKER_MEMORY=36g

slaves
# 指定 slaves 机器的列表，这里就选了本机

localhost

spark-defaults.conf
# 这个文件很多教程都会让你改，说是 spark-submit 命令会默认从这里读取相关配置
# 但要注意我们写的 PySpark 程序很多时候并不是通过 spark-submit 命令提交的，所以这里改了可能没用

spark.driver.memory 4g

3. 启动集群。直接运行 sbin/start-all.sh 即可。或者也可以分别起 master 和 slave，运行./sbin/start-master.sh 和 ./sbin/start-slave.sh spark://127.0.0.1:7077 -c 6 -m 36G 即可。

4. 停止集群。命令与上面非常类似，sbin/stop-all.sh ，或者分别停 slave 和 master 都行。

这样就算部署完了！其中 spark master 会有一个监听 8080 端口的 web-ui，worker 会监听 8081，后面提交 application 就会有监听 4040 端口的管理界面，功能强大，用户友好度强。

跑第一个 PySpark 程序

直接上代码：

from pyspark.sql import SparkSession

spark = (SparkSession.builder
         .master('spark://127.0.0.1:7077')
         .appName('zijie')
         .getOrCreate())
df = spark.read.parquet('data/the_only_data_i_ever_wanted.parquet')
df.show()

我们的大数据平台就跑起来了。

Spark 与 Pandas 的一些不同之处

在网上看一些 Spark 相关的介绍应该很快会有一些认识。有几个比较明显的区别点我大致列一下：

1. Spark 里对 DataFrame 的操作大多是 lazy 的，也就是所谓的 transformation，只有少数的 action，例如 take, count, collect 等会真实进行计算返回结果。而 pandas 只要做了操作就会立刻执行。

2. Pandas 里对性能方面的关注主要是这个操作能不能利用底层的计算库做 vectorize，而在 Spark 里需要关注的点就太多了，可能比较主要的是看怎么尽量减少 shuffle 这类宽依赖吧。当然还有什么数据倾斜等相关高级话题。

3. 用 Spark 来做算法相关的应用时，要非常注意整体的计算逻辑（数据 lineage），对需要反复用到的数据集，一定要记得 cache/persist/checkpoint 才行（这条不知是否过时了）。

从实际操作来看，在 PySpark 中其实有很多操作长得跟 Pandas 非常类似，比如我们常用的 df[df['date'] > '2020-01-01'] 之类的写法。当然区别也有不少，所以后来 Databricks 干脆推出了一个 Koalas 的库来支持更平滑的切换。

Spark 特征工程

这里主要记录几个在项目过程中写的感觉比较好玩的，并对比 pandas 的版本方便大家理解。

日期填充

pandas version：

# 对每家店每个 SKU 历史无销售情况进行填零处理
def fill_dates(df):
    new_df = []
    for store_id in df.store_id.unique():
        for sku in df.query('store_id == @store_id').sku.unique():
            tmp = pd.DataFrame()
            cond = (df.store_id == store_id) & (df.sku == sku)
            min_date = df.loc[cond, 'date'].min()
            max_date = df.loc[cond, 'date'].max()
            dates_in_between = daterange(min_date, max_date)
            tmp['date'] = dates_in_between
            tmp['sku'] = sku
            tmp['store_id'] = store_id
            new_df.append(tmp)
    new_df = pd.concat(new_df)
    new_df = new_df.merge(df, on=['date', 'sku', 'store_id'], how='left').fillna(0)
    return new_df

可以看到整体逻辑就是取所有 store, sku 的组合，然后找到每个组合最小最大的售卖日期，把中间的日期都填上。

PS: 这段代码应该效率不高，后续我们又迭代了几个版本。

Spark version：

from pyspark.sql import functions as F

def fill_dates_spark(df):
    tmp = df.groupby(['store_id', 'sku']).agg(F.min('date').cast('date').alias('min_date'),
                                              F.max('date').cast('date').alias('max_date'))
    tmp = tmp.withColumn('date', F.explode(F.sequence('min_date', 'max_date'))).select(
        ['date', 'store_id', 'sku'])
    new_df = tmp.join(df, ['date', 'store_id', 'sku'], 'left').fillna(0, subset=['y'])
    return new_df

用了 sequence+explode 操作，代码简洁很多。其中 sequence 会自动生成从 start 到 end 的序列（时间，数字都支持），explode 操作直接把一行“炸开”成多行，省去了 join 操作，性能也更好。

Lag 特征

这个是我们最常用的一种特征了，在 pandas 里主要就是做循环 join：

def shift_daily_data(df, delay, shift_by='date', shift_value='y'):
    groupby_df = [x for x in df.columns if (x != shift_by) and (x != shift_value)]
    shift_df = df.copy()
    shift_df[shift_by] = shift_df[shift_by].apply(lambda x: x + relativedelta(days=delay))
    shift_df = shift_df.rename(columns={shift_value: '%s_%s_day_lag_%d' % ('_'.join(groupby_df), shift_value, delay)})
    return shift_df

def add_daily_shifts(df, days, categories, shift_by='date', shift_value='y'):
    merge_df = df.copy()
    for base_categories in categories:
        feat_cols = base_categories + [shift_by]
        base_df = df.groupby(feat_cols, as_index=False).agg({shift_value: sum})
        for i in days:
            delay_df = shift_daily_data(base_df, i, shift_by, shift_value)
            merge_df = pd.merge(left=merge_df, right=delay_df, how='left', on=feat_cols, sort=False).reset_index(
                drop=True).fillna(0)
            gc.collect()
    return merge_df

# 按照不同维度生成 lag 自回归时序特征
def add_lag_features(all_data_df, fcst_type):
    lag_days = list(range(1, 11)) + [14, 21, 28, 29, 30, 31]
    lag_days = [x for x in lag_days if x >= fcst_type]
    groupby_cats = [['sku'], ['store_id'], ['sku', 'store_id']]
    all_data_df = add_daily_shifts(all_data_df, lag_days, groupby_cats)
    return all_data_df

在迁移到 Spark 时第一版我也采用了类似的写法，不过发现性能比较差，而且随着 lag 数的增多，join 次数也增多了，数据血缘关系会拉得非常长。

第二版我们采用了 window function 的写法：

from pyspark.sql import functions as F
from pyspark.sql import Window

def add_date_index(df, date_col, start_day='2016-01-01'):
    df = df.withColumn(f'{date_col}_index'