From cc7115f6a2d43dc9ad8dd7dd0ee3cc4d028a2af4 Mon Sep 17 00:00:00 2001
From: Antoni Baum <antoni.baum@protonmail.com>
Date: Wed, 13 Jul 2022 19:14:36 +0200
Subject: [PATCH] [Tune/CI] Fix tune-sklearn notebook example (#26470)

Fixes the tune-sklearn notebook example as found in #26410

Signed-off-by: Antoni Baum <antoni.baum@protonmail.com>
---
 doc/source/tune/examples/BUILD                |  1 -
 doc/source/tune/examples/tune-sklearn.ipynb   | 74 +++++++++++++------
 .../requirements/ml/requirements_upstream.txt |  2 +-
 3 files changed, 53 insertions(+), 24 deletions(-)

diff --git a/doc/source/tune/examples/BUILD b/doc/source/tune/examples/BUILD
index 22b20b8dbead9..16ea4105369ad 100644
--- a/doc/source/tune/examples/BUILD
+++ b/doc/source/tune/examples/BUILD
@@ -22,7 +22,6 @@ py_test_run_all_notebooks(
         "nyc_taxi_basic_processing.ipynb", # REGRESSION
         "ocr_example.ipynb", # REGRESSION
         "sigopt_example.ipynb", # REGRESSION
-        "tune-sklearn.ipynb", # REGRESSION
     ],
     data = ["//doc/source/tune/examples:tune_examples"],
     tags = ["exclusive", "team:ml"],
diff --git a/doc/source/tune/examples/tune-sklearn.ipynb b/doc/source/tune/examples/tune-sklearn.ipynb
index 12deacd5154eb..84ef6969bbaac 100644
--- a/doc/source/tune/examples/tune-sklearn.ipynb
+++ b/doc/source/tune/examples/tune-sklearn.ipynb
@@ -82,7 +82,7 @@
     "```{thebe-button} Activate code examples\n",
     "```\n",
     "\n",
-    "To start out, change the import statement to get tune-scikit-learn’s grid search cross validation interface:"
+    "To start out, change the import the modules required for this example:"
    ]
   },
   {
@@ -94,9 +94,13 @@
    "source": [
     "# Keep this here for https://github.com/ray-project/ray/issues/11547\n",
     "from sklearn.model_selection import GridSearchCV\n",
+    "from sklearn.model_selection import train_test_split\n",
+    "from sklearn.linear_model import SGDClassifier\n",
+    "from sklearn.datasets import make_classification\n",
+    "from sklearn import datasets\n",
     "\n",
-    "# Replace above line with:\n",
-    "from ray.tune.sklearn import TuneGridSearchCV"
+    "from ray.tune.sklearn import TuneGridSearchCV\n",
+    "from ray.tune.sklearn import TuneSearchCV"
    ]
   },
   {
@@ -118,12 +122,6 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Other imports\n",
-    "from sklearn.model_selection import train_test_split\n",
-    "from sklearn.linear_model import SGDClassifier\n",
-    "from sklearn.datasets import make_classification\n",
-    "import numpy as np\n",
-    "\n",
     "# Create dataset\n",
     "X, y = make_classification(\n",
     "    n_samples=11000,\n",
@@ -188,6 +186,30 @@
     "Try running this compared to the GridSearchCV equivalent, and see the speedup for yourself!"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "211615e6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# n_jobs = -1 enables use of all cores like Tune does\n",
+    "n_jobs = -1"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c3d64c5c",
+   "metadata": {
+    "tags": ["remove-cell"]
+   },
+   "outputs": [],
+   "source": [
+    "# Avoids loky error in CI\n",
+    "n_jobs = 1"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -195,10 +217,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sklearn.model_selection import GridSearchCV\n",
-    "\n",
-    "# n_jobs=-1 enables use of all cores like Tune does\n",
-    "sklearn_search = GridSearchCV(SGDClassifier(), parameter_grid, n_jobs=-1)\n",
+    "sklearn_search = GridSearchCV(SGDClassifier(), parameter_grid, n_jobs=n_jobs)\n",
     "\n",
     "start = time.time()\n",
     "sklearn_search.fit(x_train, y_train)\n",
@@ -230,12 +249,6 @@
    "outputs": [],
    "source": [
     "# First run `pip install bayesian-optimization`\n",
-    "from ray.tune.sklearn import TuneSearchCV\n",
-    "from sklearn.linear_model import SGDClassifier\n",
-    "from sklearn import datasets\n",
-    "from sklearn.model_selection import train_test_split\n",
-    "import numpy as np\n",
-    "\n",
     "digits = datasets.load_digits()\n",
     "x = digits.data\n",
     "y = digits.target\n",
@@ -288,12 +301,29 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3",
+   "display_name": "Python 3.8.10 ('venv': venv)",
    "language": "python",
    "name": "python3"
   },
-  "orphan": true
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  },
+  "orphan": true,
+  "vscode": {
+   "interpreter": {
+    "hash": "3c0d54d489a08ae47a06eae2fd00ff032d6cddb527c382959b7b2575f6a8167f"
+   }
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 5
-}
\ No newline at end of file
+}
diff --git a/python/requirements/ml/requirements_upstream.txt b/python/requirements/ml/requirements_upstream.txt
index 686335e6c1a62..28ae7e707eb72 100644
--- a/python/requirements/ml/requirements_upstream.txt
+++ b/python/requirements/ml/requirements_upstream.txt
@@ -3,7 +3,7 @@
 # So we separate its own requirements file.
 
 ray_lightning==0.2.0
-tune-sklearn==0.4.1
+tune-sklearn==0.4.3
 xgboost_ray==0.1.9
 lightgbm_ray==0.1.4
 modin>=0.11.0; python_version >= '3.7'