了解 Waze 如何利用 TFX 扩大可以投入生产的 ML 规模

发布人:Waze 数据科学家 Gal Moran、Iris Shmuel 和 Daniel Marcous

Waze

Waze 是世界上最大的基于社区的交通和导航应用。该应用借助实时数据来帮助用户避开路上的障碍,享受无忧旅途。除了移动导航功能,Waze 还提供网络平台、拼车、合作服务、广告投放等功能。如此广泛的服务组合带来了多种技术挑战和丰富案例。

Waze 的 ML 应用

Waze 依赖许多 ML 解决方案,其中包括:

  • 预测 ETA
  • 匹配乘客和司机(拼车)
  • 推送适当的广告

但是,要把这些事情做好并达到“生产级别”并不是那么容易。这类项目通常需要 复杂的相关基础架构,以便将其投入生产,因此需要多位工程师(数据科学家、软件工程师和软件可靠性工程师)一起投入大量时间。如果将大规模数据、低延迟(实际上为实时)推理、多样性案例和大量地理空间数据等 Waze 的特殊要求结合在一起,则难度更甚。

上述理由很好地解释了,为什么不经周密安排就开始在 Waze 中实施 ML 会造成混乱局面。对我们来说,这表现为:

  • 使用多个 ML 框架——你能想到的都有(sklearn、XGBoost、TensorFlow、fbprophet、Java PMML、人工设定的框架等等)
  • ML 和运算脱节——模型和特征工程被工程师嵌入 (Java) 后端服务器,监控和验证能力有限
  • 培训、验证和部署使用半手动操作
  • 从想法到投入生产的开发周期异常漫长

总的来说,数据科学家最终把大量的时间耗费在运算和监控上,而不能专注于实际的建模和数据处理。

业务发展到一定程度时,我们决定整改这混乱的局面,在自动化和流程方面投资,以便能够更快地扩大业务规模。我们决定采用全周期数据科学理念,重点投资能够大幅提高速度和质量的方法。这意味着,在我们想要建立的这个新方法中,只需一个数据科学家就能完成从研究到生产级服务的产品周期。

在新方法中,数据科学家可以直接为生产做出贡献,发挥自己最大优势。他们可专注于建模和数据处理,获得许多基础架构和开箱即用的运算。虽然上述愿景还没有完全实现,但我们觉得本文所阐述的内容对实现最终目的来说至关重要。

Waze 的 ML 堆栈

为将上述理念转化为技术规范,我们决心创建一个简单、稳定、自动化和统一的方法来构建 Waze 的 ML 流水线。

深入研究技术要求后,我们得出了以下评判标准:

  • 简单——易于理解、使用和操作
  • 可托管——没有服务器,没有硬件,只有代码
  • 可自定义——免费获得简单的东西,但也有足够的灵活性,可以为 5% 需要跳出局限的东西而激发奇思妙想
  • 可扩容——自动可扩容的数据处理、训练、推理
  • 利用 Python——我们需要一些已经可以投入生产的工具,可以兼容当今大多数工具和代码,并且适合一般的数据科学家。现在除了 Python,几乎没有其他选择

基于上述原因,我们选择 TFX 及其内置组件来提供这些功能,大部分是开箱即用的。

值得一提的是,Waze 是在 Google Cloud Platform (GCP) 上运行其技术栈。

GCP 恰好可提供一套名为 Vertex AI 的工具。Waze 就是在此 ML 基础架构平台上构建的。虽然我们使用了 Vertex AI 托管式服务的许多组件,但我们将在本文中重点介绍 Vertex Pipelines,这是 ML 流水线的一个框架,能够帮助我们封装 TFX(或任何流水线)的复杂性和设置。

连同我们的数据技术栈,Waze 的整体 ML 架构(均符合可托管、可扩容,且利用 Python 等标准)如下:

细心的读者会在这里注意到所谓的警告:我们所有流程都通过 TensorFlow 进行。

TFX 指的是 TensorFlow(尽管这种说法 已经不再完全准确了,但我们假设它是这个意思)。

当你处理许多不同的案例时,这些警告一开始可能看起来会有点吓人。

幸运的是,TF 生态系统很丰富,而且 Waze 的优点是有规模足够大的数据,神经网络会收敛。

自从开始这项工作以来,我们还没有发现在哪个案例中, TF 不能像其他框架那样更好地或充分地解决问题(在这里不讨论微观的百分点,不进行 Kaggle 比赛,而是想得到一些能够投入生产的东西)。

Waze TFX

您可能认为选择 TFX 和 Vertex 流水线已解决我们所有的问题,但这并不完全正确。

为了让这项工作真正变得简单,我们不得不写一些“胶水代码”(整合上述架构图中的各种产品),并提取足够的细节,以便普通的数据科学家能够有效和快速地使用这些东西。

这能带来以下好处:

  • 可以淘汰样板文件
  • 可隐藏所有通用的 TFX 组件,这让数据科学家可以只专注于特征工程和建模,并免费获得整个流水线
  • 可生成基于 BigQuery 的训练/评估分块
  • 可提供预先实现的可选通用特征转换(例如,扩容、归一化、插补)
  • 可提供预先实现的 Keras 模型(例如 DNN/RNN 模型。其类似 TF Estimator,但在具有 TFX 特征的 Keras 中)
  • 效用函数(例如,TF 列的准备)
  • tf.transform 特征工程代码的单元测试框架
  • 使用安装了所有 TFX 软件包的云运行实例,从 Airflow 中协调和调度流水线运行(无需在 Airflow Composer 上安装)

我们已经将其全部放入一个易于使用的 Python 软件包中,称为 “wze-data-tfx”。

在上文中,我们为数据科学家提供了一个超级详细的演示、使用指南和代码模板,所以常见的 DS 工作流是:创建分支,改变配置,稍微调整一下代码,部署。

作为参考,这里提供了一个简单的 waze-data-tfx 流水线:

1. 配置

_DATASET_NAME = 'tfx_examples'
_TABLE_NAME = 'simple_template_data'

_LABEL_KEY = 'label'
_CATEGORICAL_INT_FEATURES = {
  "categorical_calculated": 2,
}
_DENSE_FLOAT_FEATURE_KEYS = ["numeric_feature1", "numeric_feature2"]
_BUCKET_FEATURES = {
  "numeric_feature1": 5,
}
_VOCAB_FEATURES = {
  "categorical_feature": {
      'top_k': 5,
      'num_oov_buckets': 3
  }
}

_TRAIN_BATCH_SIZE = 128
_EVAL_BATCH_SIZE = 128
_NUM_EPOCHS = 250

_TRAINING_ARGS = {
  'dnn_hidden_units': [6, 3],
  'optimizer': tf.keras.optimizers.Adam,
  'optimizer_kwargs': {
      'learning_rate': 0.01
  },
  'layer_activation': None,
  'metrics': ["Accuracy"]
}

_EVAL_METRIC_SPEC = create_metric_spec([
  mse_metric(upper_bound=25, absolute_change=1),
  accuracy_metric()
])

2. 特征工程

def preprocessing_fn(inputs):
  """tf.transform's callback function for preprocessing inputs.

  Args:
      inputs: map from feature keys to raw not-yet-transformedfeatures.

  Returns:
      Map from string feature key to transformed feature operations.
  """
  outputs = features_transform(
      inputs=inputs,
      label_key=_LABEL_KEY,
      dense_features=_DENSE_FLOAT_FEATURE_KEYS,
      vocab_features=_VOCAB_FEATURES,
      bucket_features=_BUCKET_FEATURES,
  )
  return outputs

3. 建模

def _build_keras_model(**training_args):
   """Build a keras model.

   Args:
       hidden_units: [int], the layer sizes of the DNN (input layer first).
       learning_rate: [float], learning rate of the Adam optimizer.

   Returns:
       A keras model
   """
   feature_columns = \
       prepare_feature_columns(
           dense_features=_DENSE_FLOAT_FEATURE_KEYS,
           vocab_features=_VOCAB_FEATURES,
           bucket_features=_BUCKET_FEATURES,
       )

   return _dnn_regressor(deep_columns=list(feature_columns.values()),
                         dnn_hidden_units=training_args.get(
                             "dnn_hidden_units"),
                         dense_features=_DENSE_FLOAT_FEATURE_KEYS,
                         vocab_features=_VOCAB_FEATURES,
                         bucket_features=_BUCKET_FEATURES,
                         )

4. 协调

pipeline_run = WazeTFXPipelineOperator(
   dag=dag,
   task_id='pipeline_run',
   model_name='basic_pipeline_template',
   package=tfx_pipeline_basic,
   pipeline_project_id=EnvConfig.get_value('gcp-project-infra'),
   table_project_id=EnvConfig.get_value('gcp-project-infra'),
   project_utils_filename='utils.py',
   gcp_conn_id=gcp_conn_id,
   enable_pusher=True,
)

很简单,对不对?

配置文件提交到代码库后,系统就会对其进行部署并设置持续训练以及一个完整的流水线,包括所有 TFX 和 Vertex AI 方法,如数据验证、部署到 Dataflow 的转换、监控等。

总结

当我们的一位数据科学家休完长假回来,不得不使用这个新的框架来处理一个案例时,我们知道本项研究取得了不错的结果。据她所说,她能够在几个小时内启动一个可投入生产的完整流水线,而在她休假之前,她需要花几周的时间才能做到这一点。

展望未来,我们计划在 waze-data-tfx 中融入更多内容。我们认为,拥有此通用基础架构的一个关键优势是:添加一个特征后,所有人都可以“免费”使用该特征。例如,我们计划在流水线中添加额外的组件,如 Infra Validator 和 Fairness Indicators。当这些组件得到支持后,每个新的或现有的 ML 流水线将以开箱即用的方式添加这些组件,无需额外的代码。

我们正在计划有关部署的其他改进工作。希望在保证部署质量的同时,尽可能地实现自动化。

我们目前正在探索的一种方式是使用 Canary 部署。数据科学家只需要配置一个评估指标,框架(使用 Vertex Prediction 数据流量分块功能和其他持续评估方法)将在生产环境中测试新模型,并根据评估指标逐步部署或回滚。

原文:How Waze Uses TFX to Scale Production-Ready ML
中文:TensorFlow 公众号