Files
2026-07-13 13:17:40 +08:00

240 lines
7.5 KiB
Python

import os
import pickle
import sys
import time
from unittest import mock
import joblib
import numpy as np
import pytest
from sklearn.datasets import load_digits, load_iris
from sklearn.ensemble import ExtraTreesClassifier, RandomForestClassifier
from sklearn.kernel_approximation import Nystroem, RBFSampler
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import RandomizedSearchCV, cross_val_score
from sklearn.neural_network import MLPClassifier
from sklearn.pipeline import make_pipeline
from sklearn.svm import SVC, LinearSVC
from sklearn.tree import DecisionTreeClassifier
import ray
from ray._common.test_utils import wait_for_condition
from ray.util.joblib import register_ray
from ray.util.joblib.ray_backend import RayBackend
def test_register_ray():
register_ray()
assert "ray" in joblib.parallel.BACKENDS
assert not ray.is_initialized()
def test_ray_backend(shutdown_only):
register_ray()
with joblib.parallel_backend("ray"):
assert type(joblib.parallel.get_active_backend()[0]) is RayBackend
def test_svm_single_node(shutdown_only):
digits = load_digits()
param_space = {
"C": np.logspace(-6, 6, 10),
"gamma": np.logspace(-8, 8, 10),
"tol": np.logspace(-4, -1, 3),
"class_weight": [None, "balanced"],
}
class MockParallel(joblib.Parallel):
def _terminate_backend(self):
if self._backend is not None:
# test ObjectRef caching (PR #16879)
assert any(o is digits.data for o, ref in self._backend._pool._registry)
self._backend.terminate()
model = SVC(kernel="rbf")
with mock.patch("sklearn.model_selection._search.Parallel", MockParallel):
search = RandomizedSearchCV(model, param_space, cv=3, n_iter=2, verbose=10)
register_ray()
with joblib.parallel_backend("ray"):
search.fit(digits.data, digits.target)
assert ray.is_initialized()
def test_svm_multiple_nodes(ray_start_cluster_2_nodes):
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"],
}
class MockParallel(joblib.Parallel):
def _terminate_backend(self):
if self._backend is not None:
# test ObjectRef caching (PR #16879)
assert any(o is digits.data for o, ref in self._backend._pool._registry)
self._backend.terminate()
model = SVC(kernel="rbf")
with mock.patch("sklearn.model_selection._search.Parallel", MockParallel):
search = RandomizedSearchCV(model, param_space, cv=5, n_iter=2, verbose=10)
register_ray()
with joblib.parallel_backend("ray"):
search.fit(digits.data, digits.target)
assert ray.is_initialized()
"""This test only makes sure the different sklearn classifiers are supported
and do not fail. It can be improved to check for accuracy similar to
'test_cross_validation' but the classifiers need to be improved (to improve
the accuracy), which results in longer test time.
"""
@pytest.mark.skipif(sys.platform == "win32", reason="Failing on Windows.")
def test_sklearn_benchmarks(ray_start_cluster_2_nodes):
ESTIMATORS = {
"CART": DecisionTreeClassifier(),
"ExtraTrees": ExtraTreesClassifier(n_estimators=10),
"RandomForest": RandomForestClassifier(),
"Nystroem-SVM": make_pipeline(
Nystroem(gamma=0.015, n_components=1000), LinearSVC(C=1)
),
"SampledRBF-SVM": make_pipeline(
RBFSampler(gamma=0.015, n_components=1000), LinearSVC(C=1)
),
"LogisticRegression-SAG": LogisticRegression(solver="sag", tol=1e-1, C=1e4),
"LogisticRegression-SAGA": LogisticRegression(solver="saga", tol=1e-1, C=1e4),
"MultilayerPerceptron": MLPClassifier(
hidden_layer_sizes=(32, 32),
max_iter=100,
alpha=1e-4,
solver="sgd",
learning_rate_init=0.2,
momentum=0.9,
verbose=1,
tol=1e-2,
random_state=1,
),
"MLP-adam": MLPClassifier(
hidden_layer_sizes=(32, 32),
max_iter=100,
alpha=1e-4,
solver="adam",
learning_rate_init=0.001,
verbose=1,
tol=1e-2,
random_state=1,
),
}
# Load dataset.
print("Loading dataset...")
unnormalized_X_train, y_train = pickle.load(
open(os.path.join(os.path.dirname(__file__), "mnist_784_100_samples.pkl"), "rb")
)
# Normalize features.
X_train = unnormalized_X_train / 255
register_ray()
train_time = {}
random_seed = 0
# Use two workers per classifier.
num_jobs = 2
with joblib.parallel_backend("ray"):
for name in sorted(ESTIMATORS.keys()):
print("Training %s ... " % name, end="")
estimator = ESTIMATORS[name]
estimator_params = estimator.get_params()
estimator.set_params(
**{
p: random_seed
for p in estimator_params
if p.endswith("random_state")
}
)
if "n_jobs" in estimator_params:
estimator.set_params(n_jobs=num_jobs)
time_start = time.time()
estimator.fit(X_train, y_train)
train_time[name] = time.time() - time_start
print("training", name, "took", train_time[name], "seconds")
def test_cross_validation(shutdown_only):
register_ray()
iris = load_iris()
clf = SVC(kernel="linear", C=1, random_state=0)
with joblib.parallel_backend("ray", n_jobs=5):
accuracy = cross_val_score(clf, iris.data, iris.target, cv=5)
assert len(accuracy) == 5
for result in accuracy:
assert result > 0.95
def test_ray_remote_args(shutdown_only):
ray.init(num_cpus=4, resources={"custom_resource": 4})
register_ray()
assert ray.available_resources().get("custom_resource", 0) == 4
def check_resource():
assert ray.available_resources().get("custom_resource", 0) < 4
with joblib.parallel_backend(
"ray", ray_remote_args={"resources": {"custom_resource": 1}}
):
joblib.Parallel()(joblib.delayed(check_resource)() for i in range(8))
def test_task_to_actor_assignment(shutdown_only):
ray.init(num_cpus=4)
@ray.remote(num_cpus=0)
class Counter:
def __init__(self):
self._c = 0
def inc(self):
self._c += 1
def get(self) -> int:
return self._c
counter = Counter.remote()
def worker_func(worker_id):
launch_time = time.time()
# Wait for all 4 workers to have started.
ray.get(counter.inc.remote())
wait_for_condition(lambda: ray.get(counter.get.remote()) == 4)
return worker_id, launch_time
output = []
num_workers = 4
register_ray()
with joblib.parallel_backend("ray", n_jobs=-1):
output = joblib.Parallel()(
joblib.delayed(worker_func)(worker_id) for worker_id in range(num_workers)
)
worker_ids = set()
launch_times = []
for worker_id, launch_time in output:
worker_ids.add(worker_id)
launch_times.append(launch_time)
assert len(worker_ids) == num_workers
for i in range(num_workers):
for j in range(i + 1, num_workers):
assert abs(launch_times[i] - launch_times[j]) < 1
if __name__ == "__main__":
sys.exit(pytest.main(["-sv", __file__]))