[Tune/CI] Fix tune-sklearn notebook example (ray-project#26470)

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

Signed-off-by: Antoni Baum <[email protected]>

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"],

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,17 +186,38 @@
  "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,
  "id": "bad624d5",
  "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
-}
+}

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'