chore: import upstream snapshot with attribution

python/ray/tests/test_joblib.py

@@ -0,0 +1,239 @@
+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__]))