分布式 Scikit-learn / Joblib#

Ray 通过为 joblib 实现一个 Ray 后端来支持运行分布式 scikit-learn 程序,该后端使用 Ray Actors 而非本地进程。这使得将使用 scikit-learn 的现有应用从单节点扩展到集群变得容易。

注意

此 API 是新的,未来 Ray 版本可能会进行修订。如果遇到任何 bug,请在 GitHub 上提交 Issue

快速入门#

首先安装 Ray,然后使用 from ray.util.joblib import register_ray 并运行 register_ray() 即可开始。这将把 Ray 注册为 scikit-learn 使用的 joblib 后端。然后将你的原始 scikit-learn 代码运行在 with joblib.parallel_backend('ray') 块中。这将启动一个本地 Ray 集群。有关在多节点 Ray 集群上运行的说明,请参阅下面的在集群上运行部分。

import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import RandomizedSearchCV
from sklearn.svm import SVC
digits = load_digits()
param_space = {
    'C': np.logspace(-6, 6, 30),
    'gamma': np.logspace(-8, 8, 30),
    'tol': np.logspace(-4, -1, 30),
    'class_weight': [None, 'balanced'],
}
model = SVC(kernel='rbf')
search = RandomizedSearchCV(model, param_space, cv=5, n_iter=300, verbose=10)

import joblib
from ray.util.joblib import register_ray
register_ray()
with joblib.parallel_backend('ray'):
    search.fit(digits.data, digits.target)

你还可以在 parallel_backend 中设置 ray_remote_args 参数来配置 构成 Pool Ray Actor。这可以用于例如为 Actor 分配资源,例如 GPU

# Allows to use GPU-enabled estimators, such as cuML
with joblib.parallel_backend('ray', ray_remote_args=dict(num_gpus=1)):
    search.fit(digits.data, digits.target)

在集群上运行#

本节假设你已经有一个正在运行的 Ray 集群。要启动 Ray 集群,请参阅集群设置说明。

要将 scikit-learn 连接到正在运行的 Ray 集群,你需要通过设置 RAY_ADDRESS 环境变量来指定头节点的地址。

你也可以在调用 with joblib.parallel_backend('ray') 之前,通过手动调用 ray.init() (带任何支持的配置选项) 来启动 Ray。

警告

如果你不设置 RAY_ADDRESS 环境变量,也不在 ray.init(address=<address>) 中提供 address,那么 scikit-learn 将在单节点上运行!