{
"cells": [
{
"cell_type": "markdown",
"id": "c1db247a-15ab-41b1-a124-152484a29f29",
"metadata": {
"editable": true,
"slideshow": {
"slide_type": ""
},
"tags": [
"library/xgboost",
"library/optuna",
"library/dask",
"library/dask-kubernetes",
"library/scikit-learn",
"workflow/hpo",
"dataset/nyc-taxi",
"data-storage/gcs",
"data-format/csv",
"platforms/kubeflow",
"platforms/kubernetes"
]
},
"source": [
"# Scaling up Hyperparameter Optimization with Multi-GPU Workload on Kubernetes"
]
},
{
"cell_type": "markdown",
"id": "f7f02171-ed7b-48b4-9d55-32bb1149a3cf",
"metadata": {},
"source": [
"_June, 2024_\n",
"\n",
"Choosing an optimal set of hyperparameters is a daunting task, especially for algorithms like XGBoost that have many hyperparameters to tune. In this notebook, we will speed up hyperparameter optimization by running multiple training jobs in parallel on a Kubernetes cluster. We handle larger data sets by splitting the data into multiple GPU devices."
]
},
{
"cell_type": "markdown",
"id": "a718e21f-5543-4f44-8a68-6ad8e78cb433",
"metadata": {},
"source": [
"## Prerequisites\n",
"Please follow instructions in [Dask Operator: Installation](../../tools/kubernetes/dask-operator) to install the Dask operator on top of a GPU-enabled Kubernetes cluster. (For the purpose of this example, you may ignore other sections of the linked document.\n",
"\n",
"### Optional: Kubeflow\n",
"Kubeflow gives you a nice notebook environment to run this notebook within the k8s cluster. Install Kubeflow by following instructions in [Installing Kubeflow](https://www.kubeflow.org/docs/started/installing-kubeflow/). You may choose any method; we tested this example after installing Kubeflow from manifests."
]
},
{
"cell_type": "markdown",
"id": "7b7f7bb3-5d53-4b8f-8472-bb974c8a597d",
"metadata": {},
"source": [
"## Install extra Python modules\n",
"We'll need a few extra Python modules."
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "27b79db5-bbcd-422c-80a7-af873eb47711",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Collecting dask_kubernetes\n",
" Downloading dask_kubernetes-2024.5.0-py3-none-any.whl.metadata (4.2 kB)\n",
"Collecting optuna\n",
" Downloading optuna-3.6.1-py3-none-any.whl.metadata (17 kB)\n",
"Requirement already satisfied: dask>=2022.08.1 in /opt/conda/lib/python3.11/site-packages (from dask_kubernetes) (2024.1.1)\n",
"Requirement already satisfied: distributed>=2022.08.1 in /opt/conda/lib/python3.11/site-packages (from dask_kubernetes) (2024.1.1)\n",
"Collecting kopf>=1.35.3 (from dask_kubernetes)\n",
" Downloading kopf-1.37.2-py3-none-any.whl.metadata (9.7 kB)\n",
"Collecting kr8s==0.14.* (from dask_kubernetes)\n",
" Downloading kr8s-0.14.4-py3-none-any.whl.metadata (6.7 kB)\n",
"Collecting kubernetes-asyncio>=12.0.1 (from dask_kubernetes)\n",
" Downloading kubernetes_asyncio-29.0.0-py3-none-any.whl.metadata (1.3 kB)\n",
"Collecting kubernetes>=12.0.1 (from dask_kubernetes)\n",
" Downloading kubernetes-29.0.0-py2.py3-none-any.whl.metadata (1.5 kB)\n",
"Collecting pykube-ng>=22.9.0 (from dask_kubernetes)\n",
" Downloading pykube_ng-23.6.0-py3-none-any.whl.metadata (8.0 kB)\n",
"Requirement already satisfied: rich>=12.5.1 in /opt/conda/lib/python3.11/site-packages (from dask_kubernetes) (13.7.1)\n",
"Requirement already satisfied: anyio>=3.7.0 in /opt/conda/lib/python3.11/site-packages (from kr8s==0.14.*->dask_kubernetes) (4.3.0)\n",
"Collecting asyncache>=0.3.1 (from kr8s==0.14.*->dask_kubernetes)\n",
" Downloading asyncache-0.3.1-py3-none-any.whl.metadata (2.0 kB)\n",
"Collecting cryptography>=35 (from kr8s==0.14.*->dask_kubernetes)\n",
" Downloading cryptography-42.0.7-cp39-abi3-manylinux_2_28_x86_64.whl.metadata (5.3 kB)\n",
"Requirement already satisfied: exceptiongroup>=1.2.0 in /opt/conda/lib/python3.11/site-packages (from kr8s==0.14.*->dask_kubernetes) (1.2.0)\n",
"Collecting httpx-ws>=0.5.1 (from kr8s==0.14.*->dask_kubernetes)\n",
" Downloading httpx_ws-0.6.0-py3-none-any.whl.metadata (7.8 kB)\n",
"Requirement already satisfied: httpx>=0.24.1 in /opt/conda/lib/python3.11/site-packages (from kr8s==0.14.*->dask_kubernetes) (0.27.0)\n",
"Collecting python-box>=7.0.1 (from kr8s==0.14.*->dask_kubernetes)\n",
" Downloading python_box-7.1.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (7.8 kB)\n",
"Collecting python-jsonpath>=0.7.1 (from kr8s==0.14.*->dask_kubernetes)\n",
" Downloading python_jsonpath-1.1.1-py3-none-any.whl.metadata (5.3 kB)\n",
"Requirement already satisfied: pyyaml>=6.0 in /opt/conda/lib/python3.11/site-packages (from kr8s==0.14.*->dask_kubernetes) (6.0.1)\n",
"Collecting alembic>=1.5.0 (from optuna)\n",
" Downloading alembic-1.13.1-py3-none-any.whl.metadata (7.4 kB)\n",
"Collecting colorlog (from optuna)\n",
" Downloading colorlog-6.8.2-py3-none-any.whl.metadata (10 kB)\n",
"Requirement already satisfied: numpy in /opt/conda/lib/python3.11/site-packages (from optuna) (1.26.4)\n",
"Requirement already satisfied: packaging>=20.0 in /opt/conda/lib/python3.11/site-packages (from optuna) (24.0)\n",
"Collecting sqlalchemy>=1.3.0 (from optuna)\n",
" Downloading SQLAlchemy-2.0.30-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl.metadata (9.6 kB)\n",
"Requirement already satisfied: tqdm in /opt/conda/lib/python3.11/site-packages (from optuna) (4.66.2)\n",
"Collecting Mako (from alembic>=1.5.0->optuna)\n",
" Downloading Mako-1.3.3-py3-none-any.whl.metadata (2.9 kB)\n",
"Requirement already satisfied: typing-extensions>=4 in /opt/conda/lib/python3.11/site-packages (from alembic>=1.5.0->optuna) (4.11.0)\n",
"Requirement already satisfied: click>=8.1 in /opt/conda/lib/python3.11/site-packages (from dask>=2022.08.1->dask_kubernetes) (8.1.7)\n",
"Requirement already satisfied: cloudpickle>=1.5.0 in /opt/conda/lib/python3.11/site-packages (from dask>=2022.08.1->dask_kubernetes) (3.0.0)\n",
"Requirement already satisfied: fsspec>=2021.09.0 in /opt/conda/lib/python3.11/site-packages (from dask>=2022.08.1->dask_kubernetes) (2024.3.1)\n",
"Requirement already satisfied: partd>=1.2.0 in /opt/conda/lib/python3.11/site-packages (from dask>=2022.08.1->dask_kubernetes) (1.4.1)\n",
"Requirement already satisfied: toolz>=0.10.0 in /opt/conda/lib/python3.11/site-packages (from dask>=2022.08.1->dask_kubernetes) (0.12.1)\n",
"Requirement already satisfied: importlib-metadata>=4.13.0 in /opt/conda/lib/python3.11/site-packages (from dask>=2022.08.1->dask_kubernetes) (7.1.0)\n",
"Requirement already satisfied: jinja2>=2.10.3 in /opt/conda/lib/python3.11/site-packages (from distributed>=2022.08.1->dask_kubernetes) (3.1.3)\n",
"Requirement already satisfied: locket>=1.0.0 in /opt/conda/lib/python3.11/site-packages (from distributed>=2022.08.1->dask_kubernetes) (1.0.0)\n",
"Requirement already satisfied: msgpack>=1.0.0 in /opt/conda/lib/python3.11/site-packages (from distributed>=2022.08.1->dask_kubernetes) (1.0.7)\n",
"Requirement already satisfied: psutil>=5.7.2 in /opt/conda/lib/python3.11/site-packages (from distributed>=2022.08.1->dask_kubernetes) (5.9.8)\n",
"Requirement already satisfied: sortedcontainers>=2.0.5 in /opt/conda/lib/python3.11/site-packages (from distributed>=2022.08.1->dask_kubernetes) (2.4.0)\n",
"Requirement already satisfied: tblib>=1.6.0 in /opt/conda/lib/python3.11/site-packages (from distributed>=2022.08.1->dask_kubernetes) (3.0.0)\n",
"Requirement already satisfied: tornado>=6.0.4 in /opt/conda/lib/python3.11/site-packages (from distributed>=2022.08.1->dask_kubernetes) (6.4)\n",
"Requirement already satisfied: urllib3>=1.24.3 in /opt/conda/lib/python3.11/site-packages (from distributed>=2022.08.1->dask_kubernetes) (1.26.18)\n",
"Requirement already satisfied: zict>=3.0.0 in /opt/conda/lib/python3.11/site-packages (from distributed>=2022.08.1->dask_kubernetes) (3.0.0)\n",
"Requirement already satisfied: python-json-logger in /opt/conda/lib/python3.11/site-packages (from kopf>=1.35.3->dask_kubernetes) (2.0.7)\n",
"Collecting iso8601 (from kopf>=1.35.3->dask_kubernetes)\n",
" Downloading iso8601-2.1.0-py3-none-any.whl.metadata (3.7 kB)\n",
"Requirement already satisfied: aiohttp in /opt/conda/lib/python3.11/site-packages (from kopf>=1.35.3->dask_kubernetes) (3.9.5)\n",
"Requirement already satisfied: certifi>=14.05.14 in /opt/conda/lib/python3.11/site-packages (from kubernetes>=12.0.1->dask_kubernetes) (2024.2.2)\n",
"Requirement already satisfied: six>=1.9.0 in /opt/conda/lib/python3.11/site-packages (from kubernetes>=12.0.1->dask_kubernetes) (1.16.0)\n",
"Requirement already satisfied: python-dateutil>=2.5.3 in /opt/conda/lib/python3.11/site-packages (from kubernetes>=12.0.1->dask_kubernetes) (2.9.0)\n",
"Collecting google-auth>=1.0.1 (from kubernetes>=12.0.1->dask_kubernetes)\n",
" Downloading google_auth-2.29.0-py2.py3-none-any.whl.metadata (4.7 kB)\n",
"Requirement already satisfied: websocket-client!=0.40.0,!=0.41.*,!=0.42.*,>=0.32.0 in /opt/conda/lib/python3.11/site-packages (from kubernetes>=12.0.1->dask_kubernetes) (1.8.0)\n",
"Requirement already satisfied: requests in /opt/conda/lib/python3.11/site-packages (from kubernetes>=12.0.1->dask_kubernetes) (2.31.0)\n",
"Collecting requests-oauthlib (from kubernetes>=12.0.1->dask_kubernetes)\n",
" Downloading requests_oauthlib-2.0.0-py2.py3-none-any.whl.metadata (11 kB)\n",
"Collecting oauthlib>=3.2.2 (from kubernetes>=12.0.1->dask_kubernetes)\n",
" Downloading oauthlib-3.2.2-py3-none-any.whl.metadata (7.5 kB)\n",
"Requirement already satisfied: setuptools>=21.0.0 in /opt/conda/lib/python3.11/site-packages (from kubernetes-asyncio>=12.0.1->dask_kubernetes) (69.5.1)\n",
"Requirement already satisfied: markdown-it-py>=2.2.0 in /opt/conda/lib/python3.11/site-packages (from rich>=12.5.1->dask_kubernetes) (3.0.0)\n",
"Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /opt/conda/lib/python3.11/site-packages (from rich>=12.5.1->dask_kubernetes) (2.17.2)\n",
"Collecting greenlet!=0.4.17 (from sqlalchemy>=1.3.0->optuna)\n",
" Downloading greenlet-3.0.3-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl.metadata (3.8 kB)\n",
"Requirement already satisfied: aiosignal>=1.1.2 in /opt/conda/lib/python3.11/site-packages (from aiohttp->kopf>=1.35.3->dask_kubernetes) (1.3.1)\n",
"Requirement already satisfied: attrs>=17.3.0 in /opt/conda/lib/python3.11/site-packages (from aiohttp->kopf>=1.35.3->dask_kubernetes) (23.2.0)\n",
"Requirement already satisfied: frozenlist>=1.1.1 in /opt/conda/lib/python3.11/site-packages (from aiohttp->kopf>=1.35.3->dask_kubernetes) (1.4.1)\n",
"Requirement already satisfied: multidict<7.0,>=4.5 in /opt/conda/lib/python3.11/site-packages (from aiohttp->kopf>=1.35.3->dask_kubernetes) (6.0.5)\n",
"Requirement already satisfied: yarl<2.0,>=1.0 in /opt/conda/lib/python3.11/site-packages (from aiohttp->kopf>=1.35.3->dask_kubernetes) (1.9.4)\n",
"Requirement already satisfied: idna>=2.8 in /opt/conda/lib/python3.11/site-packages (from anyio>=3.7.0->kr8s==0.14.*->dask_kubernetes) (3.7)\n",
"Requirement already satisfied: sniffio>=1.1 in /opt/conda/lib/python3.11/site-packages (from anyio>=3.7.0->kr8s==0.14.*->dask_kubernetes) (1.3.1)\n",
"Requirement already satisfied: cachetools<6.0.0,>=5.2.0 in /opt/conda/lib/python3.11/site-packages (from asyncache>=0.3.1->kr8s==0.14.*->dask_kubernetes) (5.3.3)\n",
"Requirement already satisfied: cffi>=1.12 in /opt/conda/lib/python3.11/site-packages (from cryptography>=35->kr8s==0.14.*->dask_kubernetes) (1.16.0)\n",
"Collecting pyasn1-modules>=0.2.1 (from google-auth>=1.0.1->kubernetes>=12.0.1->dask_kubernetes)\n",
" Downloading pyasn1_modules-0.4.0-py3-none-any.whl.metadata (3.4 kB)\n",
"Collecting rsa<5,>=3.1.4 (from google-auth>=1.0.1->kubernetes>=12.0.1->dask_kubernetes)\n",
" Downloading rsa-4.9-py3-none-any.whl.metadata (4.2 kB)\n",
"Requirement already satisfied: httpcore==1.* in /opt/conda/lib/python3.11/site-packages (from httpx>=0.24.1->kr8s==0.14.*->dask_kubernetes) (1.0.5)\n",
"Requirement already satisfied: h11<0.15,>=0.13 in /opt/conda/lib/python3.11/site-packages (from httpcore==1.*->httpx>=0.24.1->kr8s==0.14.*->dask_kubernetes) (0.14.0)\n",
"Collecting wsproto (from httpx-ws>=0.5.1->kr8s==0.14.*->dask_kubernetes)\n",
" Downloading wsproto-1.2.0-py3-none-any.whl.metadata (5.6 kB)\n",
"Requirement already satisfied: zipp>=0.5 in /opt/conda/lib/python3.11/site-packages (from importlib-metadata>=4.13.0->dask>=2022.08.1->dask_kubernetes) (3.17.0)\n",
"Requirement already satisfied: MarkupSafe>=2.0 in /opt/conda/lib/python3.11/site-packages (from jinja2>=2.10.3->distributed>=2022.08.1->dask_kubernetes) (2.1.5)\n",
"Requirement already satisfied: mdurl~=0.1 in /opt/conda/lib/python3.11/site-packages (from markdown-it-py>=2.2.0->rich>=12.5.1->dask_kubernetes) (0.1.2)\n",
"Requirement already satisfied: charset-normalizer<4,>=2 in /opt/conda/lib/python3.11/site-packages (from requests->kubernetes>=12.0.1->dask_kubernetes) (3.3.2)\n",
"Requirement already satisfied: pycparser in /opt/conda/lib/python3.11/site-packages (from cffi>=1.12->cryptography>=35->kr8s==0.14.*->dask_kubernetes) (2.22)\n",
"Collecting pyasn1<0.7.0,>=0.4.6 (from pyasn1-modules>=0.2.1->google-auth>=1.0.1->kubernetes>=12.0.1->dask_kubernetes)\n",
" Downloading pyasn1-0.6.0-py2.py3-none-any.whl.metadata (8.3 kB)\n",
"Downloading dask_kubernetes-2024.5.0-py3-none-any.whl (157 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m157.2/157.2 kB\u001b[0m \u001b[31m2.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m00:01\u001b[0m0:01\u001b[0m\n",
"\u001b[?25hDownloading kr8s-0.14.4-py3-none-any.whl (60 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m60.7/60.7 kB\u001b[0m \u001b[31m7.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading optuna-3.6.1-py3-none-any.whl (380 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m380.1/380.1 kB\u001b[0m \u001b[31m8.0 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
"\u001b[?25hDownloading alembic-1.13.1-py3-none-any.whl (233 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m233.4/233.4 kB\u001b[0m \u001b[31m23.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading kopf-1.37.2-py3-none-any.whl (207 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m207.8/207.8 kB\u001b[0m \u001b[31m15.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading kubernetes-29.0.0-py2.py3-none-any.whl (1.6 MB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m1.6/1.6 MB\u001b[0m \u001b[31m27.2 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0ma \u001b[36m0:00:01\u001b[0m\n",
"\u001b[?25hDownloading kubernetes_asyncio-29.0.0-py3-none-any.whl (2.0 MB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m2.0/2.0 MB\u001b[0m \u001b[31m83.5 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading pykube_ng-23.6.0-py3-none-any.whl (26 kB)\n",
"Downloading SQLAlchemy-2.0.30-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (3.2 MB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.2/3.2 MB\u001b[0m \u001b[31m122.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading colorlog-6.8.2-py3-none-any.whl (11 kB)\n",
"Downloading asyncache-0.3.1-py3-none-any.whl (3.7 kB)\n",
"Downloading cryptography-42.0.7-cp39-abi3-manylinux_2_28_x86_64.whl (3.8 MB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m3.8/3.8 MB\u001b[0m \u001b[31m125.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading google_auth-2.29.0-py2.py3-none-any.whl (189 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m189.2/189.2 kB\u001b[0m \u001b[31m29.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading greenlet-3.0.3-cp311-cp311-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl (620 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m620.0/620.0 kB\u001b[0m \u001b[31m61.9 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading httpx_ws-0.6.0-py3-none-any.whl (13 kB)\n",
"Downloading oauthlib-3.2.2-py3-none-any.whl (151 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m151.7/151.7 kB\u001b[0m \u001b[31m24.6 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading python_box-7.1.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl (4.3 MB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m4.3/4.3 MB\u001b[0m \u001b[31m131.1 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading python_jsonpath-1.1.1-py3-none-any.whl (51 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m51.5/51.5 kB\u001b[0m \u001b[31m8.8 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading iso8601-2.1.0-py3-none-any.whl (7.5 kB)\n",
"Downloading Mako-1.3.3-py3-none-any.whl (78 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m78.8/78.8 kB\u001b[0m \u001b[31m12.7 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading requests_oauthlib-2.0.0-py2.py3-none-any.whl (24 kB)\n",
"Downloading pyasn1_modules-0.4.0-py3-none-any.whl (181 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m181.2/181.2 kB\u001b[0m \u001b[31m27.3 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hDownloading rsa-4.9-py3-none-any.whl (34 kB)\n",
"Downloading wsproto-1.2.0-py3-none-any.whl (24 kB)\n",
"Downloading pyasn1-0.6.0-py2.py3-none-any.whl (85 kB)\n",
"\u001b[2K \u001b[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━\u001b[0m \u001b[32m85.3/85.3 kB\u001b[0m \u001b[31m12.4 MB/s\u001b[0m eta \u001b[36m0:00:00\u001b[0m\n",
"\u001b[?25hInstalling collected packages: wsproto, python-jsonpath, python-box, pyasn1, oauthlib, Mako, iso8601, greenlet, colorlog, asyncache, sqlalchemy, rsa, requests-oauthlib, pykube-ng, pyasn1-modules, cryptography, kubernetes-asyncio, kopf, httpx-ws, google-auth, alembic, optuna, kubernetes, kr8s, dask_kubernetes\n",
"Successfully installed Mako-1.3.3 alembic-1.13.1 asyncache-0.3.1 colorlog-6.8.2 cryptography-42.0.7 dask_kubernetes-2024.5.0 google-auth-2.29.0 greenlet-3.0.3 httpx-ws-0.6.0 iso8601-2.1.0 kopf-1.37.2 kr8s-0.14.4 kubernetes-29.0.0 kubernetes-asyncio-29.0.0 oauthlib-3.2.2 optuna-3.6.1 pyasn1-0.6.0 pyasn1-modules-0.4.0 pykube-ng-23.6.0 python-box-7.1.1 python-jsonpath-1.1.1 requests-oauthlib-2.0.0 rsa-4.9 sqlalchemy-2.0.30 wsproto-1.2.0\n"
]
}
],
"source": [
"!pip install dask_kubernetes optuna"
]
},
{
"cell_type": "markdown",
"id": "acc8f524-dc9b-41d7-8faa-3aea23ee1983",
"metadata": {},
"source": [
"## Import Python modules"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "0c8a1ffb-0b03-4d4a-9ab1-0561bf5533d9",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"import threading\n",
"import warnings\n",
"\n",
"import cupy as cp\n",
"import cuspatial\n",
"import dask_cudf\n",
"import optuna\n",
"from cuml.dask.common import utils as dask_utils\n",
"from dask.distributed import Client, wait\n",
"from dask_kubernetes.operator import KubeCluster\n",
"from dask_ml.metrics import mean_squared_error\n",
"from dask_ml.model_selection import KFold\n",
"from xgboost import dask as dxgb"
]
},
{
"cell_type": "markdown",
"id": "b2d61e0b-229b-40c0-889d-b8242e574fc8",
"metadata": {},
"source": [
"## Set up multiple Dask clusters\n",
"\n",
"To run multi-GPU training jobs in parallel, we will create multiple Dask clusters each controlling its share of GPUs. It's best to think of each Dask cluster as a portion of the compute resource of the Kubernetes cluster.\n",
"\n",
"Fill in the following variables:"
]
},
{
"cell_type": "code",
"execution_count": 29,
"id": "d1c22c3c-51b2-4526-b1fa-ac012f616e13",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"n_clusters=2\n",
"n_worker_per_dask_cluster=2\n",
"n_node_per_dask_cluster=3\n"
]
}
],
"source": [
"# Number of nodes in the Kubernetes cluster.\n",
"# Each node is assumed to have a single NVIDIA GPU attached\n",
"n_nodes = 7\n",
"\n",
"# Number of worker nodes to be assigned to each Dask cluster\n",
"n_worker_per_dask_cluster = 2\n",
"\n",
"# Number of nodes to be assigned to each Dask cluster\n",
"# 1 is added since the Dask cluster's scheduler process needs to be mapped to its own node\n",
"n_node_per_dask_cluster = n_worker_per_dask_cluster + 1\n",
"\n",
"# Number of Dask clusters to be created\n",
"# Subtract 1 to account for the notebook Pod (it requires its own node)\n",
"n_clusters = (n_nodes - 1) // n_node_per_dask_cluster\n",
"\n",
"print(f\"{n_clusters=}\")\n",
"if n_clusters == 0:\n",
" raise ValueError(\n",
" \"No cluster can be created. Reduce `n_worker_per_dask_cluster` or create more compute nodes\"\n",
" )\n",
"print(f\"{n_worker_per_dask_cluster=}\")\n",
"print(f\"{n_node_per_dask_cluster=}\")\n",
"\n",
"n_node_active = n_clusters * n_node_per_dask_cluster + 1\n",
"if n_node_active != n_nodes:\n",
" n_idle = n_nodes - n_node_active\n",
" warnings.warn(f\"{n_idle} node(s) will not be used\", stacklevel=2)"
]
},
{
"cell_type": "markdown",
"id": "c0eee823-162f-47e9-be4c-41447b2d7ee9",
"metadata": {},
"source": [
"Once we've determined the number of Dask clusters and their size, we are now ready to launch them:"
]
},
{
"cell_type": "code",
"execution_count": 30,
"id": "8d0b632a-b73d-4351-bb5d-8a1f4ab1a2c4",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# Choose the same RAPIDS image you used for launching the notebook session\n",
"rapids_image = \"\""
]
},
{
"cell_type": "code",
"execution_count": 33,
"id": "62aa9e52-c5b6-487c-8f02-88ea84980cfc",
"metadata": {
"tags": []
},
"outputs": [
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "e24e5095ae78458e804d5f1212372f9a",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Output()"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Launching cluster 0...\n"
]
},
{
"data": {
"text/html": [
"\n"
],
"text/plain": []
},
"metadata": {},
"output_type": "display_data"
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "240e689def1549c1b5dfd87284192e96",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"Output()"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"Launching cluster 1...\n"
]
},
{
"data": {
"text/html": [
"\n"
],
"text/plain": []
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"clusters = []\n",
"for i in range(n_clusters):\n",
" print(f\"Launching cluster {i}...\")\n",
" clusters.append(\n",
" KubeCluster(\n",
" name=f\"rapids-dask{i}\",\n",
" image=rapids_image,\n",
" worker_command=\"dask-cuda-worker\",\n",
" n_workers=2,\n",
" resources={\"limits\": {\"nvidia.com/gpu\": \"1\"}},\n",
" env={\"EXTRA_PIP_PACKAGES\": \"optuna\"},\n",
" )\n",
" )"
]
},
{
"cell_type": "markdown",
"id": "f37fa67f-fa90-432c-bed3-8f2a8a095795",
"metadata": {},
"source": [
"## Set up Hyperparameter Optimization Task with NYC Taxi data\n",
"\n",
"Anaconda has graciously made some of the NYC Taxi dataset available in a public Google Cloud Storage bucket. We'll use our Cluster of GPUs to process it and train a model that predicts the fare amount. We'll use our Dask clusters to process it and train a model that predicts the fare amount."
]
},
{
"cell_type": "code",
"execution_count": 34,
"id": "c84929a5-f13b-4a61-9ed6-aa8060129e17",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"col_dtype = {\n",
" \"VendorID\": \"int32\",\n",
" \"tpep_pickup_datetime\": \"datetime64[ms]\",\n",
" \"tpep_dropoff_datetime\": \"datetime64[ms]\",\n",
" \"passenger_count\": \"int32\",\n",
" \"trip_distance\": \"float32\",\n",
" \"pickup_longitude\": \"float32\",\n",
" \"pickup_latitude\": \"float32\",\n",
" \"RatecodeID\": \"int32\",\n",
" \"store_and_fwd_flag\": \"int32\",\n",
" \"dropoff_longitude\": \"float32\",\n",
" \"dropoff_latitude\": \"float32\",\n",
" \"payment_type\": \"int32\",\n",
" \"fare_amount\": \"float32\",\n",
" \"extra\": \"float32\",\n",
" \"mta_tax\": \"float32\",\n",
" \"tip_amount\": \"float32\",\n",
" \"total_amount\": \"float32\",\n",
" \"tolls_amount\": \"float32\",\n",
" \"improvement_surcharge\": \"float32\",\n",
"}\n",
"\n",
"\n",
"must_haves = {\n",
" \"pickup_datetime\": \"datetime64[ms]\",\n",
" \"dropoff_datetime\": \"datetime64[ms]\",\n",
" \"passenger_count\": \"int32\",\n",
" \"trip_distance\": \"float32\",\n",
" \"pickup_longitude\": \"float32\",\n",
" \"pickup_latitude\": \"float32\",\n",
" \"rate_code\": \"int32\",\n",
" \"dropoff_longitude\": \"float32\",\n",
" \"dropoff_latitude\": \"float32\",\n",
" \"fare_amount\": \"float32\",\n",
"}\n",
"\n",
"\n",
"def compute_haversine_distance(df):\n",
" pickup = cuspatial.GeoSeries.from_points_xy(\n",
" df[[\"pickup_longitude\", \"pickup_latitude\"]].interleave_columns()\n",
" )\n",
" dropoff = cuspatial.GeoSeries.from_points_xy(\n",
" df[[\"dropoff_longitude\", \"dropoff_latitude\"]].interleave_columns()\n",
" )\n",
" df[\"haversine_distance\"] = cuspatial.haversine_distance(pickup, dropoff)\n",
" df[\"haversine_distance\"] = df[\"haversine_distance\"].astype(\"float32\")\n",
" return df\n",
"\n",
"\n",
"def clean(ddf, must_haves):\n",
" # replace the extraneous spaces in column names and lower the font type\n",
" tmp = {col: col.strip().lower() for col in list(ddf.columns)}\n",
" ddf = ddf.rename(columns=tmp)\n",
"\n",
" ddf = ddf.rename(\n",
" columns={\n",
" \"tpep_pickup_datetime\": \"pickup_datetime\",\n",
" \"tpep_dropoff_datetime\": \"dropoff_datetime\",\n",
" \"ratecodeid\": \"rate_code\",\n",
" }\n",
" )\n",
"\n",
" ddf[\"pickup_datetime\"] = ddf[\"pickup_datetime\"].astype(\"datetime64[ms]\")\n",
" ddf[\"dropoff_datetime\"] = ddf[\"dropoff_datetime\"].astype(\"datetime64[ms]\")\n",
"\n",
" for col in ddf.columns:\n",
" if col not in must_haves:\n",
" ddf = ddf.drop(columns=col)\n",
" continue\n",
" if ddf[col].dtype == \"object\":\n",
" # Fixing error: could not convert arg to str\n",
" ddf = ddf.drop(columns=col)\n",
" else:\n",
" # downcast from 64bit to 32bit types\n",
" # Tesla T4 are faster on 32bit ops\n",
" if \"int\" in str(ddf[col].dtype):\n",
" ddf[col] = ddf[col].astype(\"int32\")\n",
" if \"float\" in str(ddf[col].dtype):\n",
" ddf[col] = ddf[col].astype(\"float32\")\n",
" ddf[col] = ddf[col].fillna(-1)\n",
"\n",
" return ddf\n",
"\n",
"\n",
"def prepare_data(client):\n",
" taxi_df = dask_cudf.read_csv(\n",
" \"https://storage.googleapis.com/anaconda-public-data/nyc-taxi/csv/2016/yellow_tripdata_2016-02.csv\",\n",
" dtype=col_dtype,\n",
" )\n",
" taxi_df = taxi_df.map_partitions(clean, must_haves, meta=must_haves)\n",
"\n",
" ## add features\n",
" taxi_df[\"hour\"] = taxi_df[\"pickup_datetime\"].dt.hour.astype(\"int32\")\n",
" taxi_df[\"year\"] = taxi_df[\"pickup_datetime\"].dt.year.astype(\"int32\")\n",
" taxi_df[\"month\"] = taxi_df[\"pickup_datetime\"].dt.month.astype(\"int32\")\n",
" taxi_df[\"day\"] = taxi_df[\"pickup_datetime\"].dt.day.astype(\"int32\")\n",
" taxi_df[\"day_of_week\"] = taxi_df[\"pickup_datetime\"].dt.weekday.astype(\"int32\")\n",
" taxi_df[\"is_weekend\"] = (taxi_df[\"day_of_week\"] >= 5).astype(\"int32\")\n",
"\n",
" # calculate the time difference between dropoff and pickup.\n",
" taxi_df[\"diff\"] = taxi_df[\"dropoff_datetime\"].astype(\"int32\") - taxi_df[\n",
" \"pickup_datetime\"\n",
" ].astype(\"int32\")\n",
" taxi_df[\"diff\"] = (taxi_df[\"diff\"] / 1000).astype(\"int32\")\n",
"\n",
" taxi_df[\"pickup_latitude_r\"] = taxi_df[\"pickup_latitude\"] // 0.01 * 0.01\n",
" taxi_df[\"pickup_longitude_r\"] = taxi_df[\"pickup_longitude\"] // 0.01 * 0.01\n",
" taxi_df[\"dropoff_latitude_r\"] = taxi_df[\"dropoff_latitude\"] // 0.01 * 0.01\n",
" taxi_df[\"dropoff_longitude_r\"] = taxi_df[\"dropoff_longitude\"] // 0.01 * 0.01\n",
"\n",
" taxi_df = taxi_df.drop(\"pickup_datetime\", axis=1)\n",
" taxi_df = taxi_df.drop(\"dropoff_datetime\", axis=1)\n",
"\n",
" taxi_df = taxi_df.map_partitions(compute_haversine_distance)\n",
"\n",
" X = (\n",
" taxi_df.drop([\"fare_amount\"], axis=1)\n",
" .astype(\"float32\")\n",
" .to_dask_array(lengths=True)\n",
" )\n",
" y = taxi_df[\"fare_amount\"].astype(\"float32\").to_dask_array(lengths=True)\n",
"\n",
" X._meta = cp.asarray(X._meta)\n",
" y._meta = cp.asarray(y._meta)\n",
"\n",
" X, y = dask_utils.persist_across_workers(client, [X, y])\n",
" return X, y\n",
"\n",
"\n",
"def train_model(params):\n",
" cluster = get_cluster(threading.get_ident())\n",
"\n",
" default_params = {\n",
" \"objective\": \"reg:squarederror\",\n",
" \"eval_metric\": \"rmse\",\n",
" \"verbosity\": 0,\n",
" \"tree_method\": \"hist\",\n",
" \"device\": \"cuda\",\n",
" }\n",
" params = dict(default_params, **params)\n",
"\n",
" with Client(cluster) as client:\n",
" X, y = prepare_data(client)\n",
" wait([X, y])\n",
"\n",
" scores = []\n",
" kfold = KFold(n_splits=5, shuffle=False)\n",
" for train_index, test_index in kfold.split(X, y):\n",
" dtrain = dxgb.DaskQuantileDMatrix(client, X[train_index, :], y[train_index])\n",
" dtest = dxgb.DaskQuantileDMatrix(client, X[test_index, :], y[test_index])\n",
" model = dxgb.train(\n",
" client,\n",
" params,\n",
" dtrain,\n",
" num_boost_round=10,\n",
" verbose_eval=False,\n",
" )\n",
" y_test_pred = dxgb.predict(client, model, dtest).to_backend(\"cupy\")\n",
" rmse_score = mean_squared_error(y[test_index], y_test_pred, squared=False)\n",
" scores.append(rmse_score)\n",
" return sum(scores) / len(scores)\n",
"\n",
"\n",
"def objective(trial):\n",
" params = {\n",
" \"n_estimators\": trial.suggest_int(\"n_estimators\", 2, 4),\n",
" \"learning_rate\": trial.suggest_float(\"learning_rate\", 0.5, 0.7),\n",
" \"colsample_bytree\": trial.suggest_float(\"colsample_bytree\", 0.5, 1),\n",
" \"colsample_bynode\": trial.suggest_float(\"colsample_bynode\", 0.5, 1),\n",
" \"colsample_bylevel\": trial.suggest_float(\"colsample_bylevel\", 0.5, 1),\n",
" \"reg_lambda\": trial.suggest_float(\"reg_lambda\", 0, 1),\n",
" \"max_depth\": trial.suggest_int(\"max_depth\", 1, 6),\n",
" \"max_leaves\": trial.suggest_int(\"max_leaves\", 0, 2),\n",
" \"max_cat_to_onehot\": trial.suggest_int(\"max_cat_to_onehot\", 1, 10),\n",
" }\n",
" return train_model(params)"
]
},
{
"cell_type": "markdown",
"id": "0c401aa1-2aeb-43d9-955b-4dfd7b495fe9",
"metadata": {},
"source": [
"To kick off multiple training jobs in parallel, we will launch multiple threads, so that each thread controls a Dask cluster.\n",
"One important utility function is `get_cluster`, which returns the Dask cluster that's mapped to a given thread."
]
},
{
"cell_type": "code",
"execution_count": 35,
"id": "97cdeb8a-330e-4d96-92d4-d48c93828d9d",
"metadata": {
"tags": []
},
"outputs": [],
"source": [
"# Map each thread's integer ID to a sequential number (0, 1, 2 ...)\n",
"thread_id_map: dict[int, KubeCluster] = {}\n",
"thread_id_map_lock = threading.Lock()\n",
"\n",
"\n",
"def get_cluster(thread_id: int) -> KubeCluster:\n",
" with thread_id_map_lock:\n",
" try:\n",
" return clusters[thread_id_map[thread_id]]\n",
" except KeyError:\n",
" seq_id = len(thread_id_map)\n",
" thread_id_map[thread_id] = seq_id\n",
" return clusters[seq_id]"
]
},
{
"cell_type": "markdown",
"id": "2e7c923b-f4ea-4f38-b3a5-92dfcd47dfff",
"metadata": {},
"source": [
"Now we are ready to start hyperparameter optimization."
]
},
{
"cell_type": "code",
"execution_count": 36,
"id": "c557d769-0be6-4319-b7f5-8ad52b824961",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[I 2024-05-09 07:53:00,718] A new study created in memory with name: no-name-da830427-bce3-4e42-98e6-c98c0c3da0d7\n"
]
}
],
"source": [
"n_trials = (\n",
" 10 # set to a low number so that the demo finishes quickly. Feel free to adjust\n",
")\n",
"study = optuna.create_study(direction=\"minimize\")"
]
},
{
"cell_type": "code",
"execution_count": 37,
"id": "94ece2d0-b3f7-44c8-9b4e-a2f60fd623b9",
"metadata": {
"tags": []
},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"[I 2024-05-09 07:54:10,229] Trial 1 finished with value: 59.449462890625 and parameters: {'n_estimators': 4, 'learning_rate': 0.6399993857892183, 'colsample_bytree': 0.7020623988319513, 'colsample_bynode': 0.777468318546648, 'colsample_bylevel': 0.7890749134903386, 'reg_lambda': 0.4464953694744921, 'max_depth': 3, 'max_leaves': 0, 'max_cat_to_onehot': 9}. Best is trial 1 with value: 59.449462890625.\n",
"[I 2024-05-09 07:54:19,507] Trial 0 finished with value: 57.77985763549805 and parameters: {'n_estimators': 4, 'learning_rate': 0.674087333032356, 'colsample_bytree': 0.557642421113256, 'colsample_bynode': 0.9719449711676733, 'colsample_bylevel': 0.6984302171973646, 'reg_lambda': 0.7201514298169174, 'max_depth': 4, 'max_leaves': 1, 'max_cat_to_onehot': 4}. Best is trial 0 with value: 57.77985763549805.\n",
"[I 2024-05-09 07:54:59,524] Trial 2 finished with value: 57.77985763549805 and parameters: {'n_estimators': 2, 'learning_rate': 0.6894880267544121, 'colsample_bytree': 0.8171662437182604, 'colsample_bynode': 0.549527686217645, 'colsample_bylevel': 0.890212178266078, 'reg_lambda': 0.5847298606135033, 'max_depth': 2, 'max_leaves': 1, 'max_cat_to_onehot': 5}. Best is trial 0 with value: 57.77985763549805.\n",
"[I 2024-05-09 07:55:22,013] Trial 3 finished with value: 55.01234817504883 and parameters: {'n_estimators': 4, 'learning_rate': 0.6597614733926671, 'colsample_bytree': 0.8437061126308156, 'colsample_bynode': 0.621479934699203, 'colsample_bylevel': 0.8330951489228277, 'reg_lambda': 0.7830102753448884, 'max_depth': 2, 'max_leaves': 2, 'max_cat_to_onehot': 2}. Best is trial 3 with value: 55.01234817504883.\n",
"[I 2024-05-09 07:56:00,678] Trial 4 finished with value: 57.77985763549805 and parameters: {'n_estimators': 4, 'learning_rate': 0.5994587326401378, 'colsample_bytree': 0.9799078215504886, 'colsample_bynode': 0.9766955839079614, 'colsample_bylevel': 0.5088864363378924, 'reg_lambda': 0.18103184809548734, 'max_depth': 3, 'max_leaves': 1, 'max_cat_to_onehot': 4}. Best is trial 3 with value: 55.01234817504883.\n",
"[I 2024-05-09 07:56:11,773] Trial 5 finished with value: 54.936126708984375 and parameters: {'n_estimators': 2, 'learning_rate': 0.5208827661289628, 'colsample_bytree': 0.866258912492528, 'colsample_bynode': 0.6368815844513638, 'colsample_bylevel': 0.9539603435186208, 'reg_lambda': 0.21390618865079458, 'max_depth': 4, 'max_leaves': 2, 'max_cat_to_onehot': 4}. Best is trial 5 with value: 54.936126708984375.\n",
"[I 2024-05-09 07:56:48,737] Trial 6 finished with value: 57.77985763549805 and parameters: {'n_estimators': 2, 'learning_rate': 0.6137888371528442, 'colsample_bytree': 0.9621063205689744, 'colsample_bynode': 0.5306812468481084, 'colsample_bylevel': 0.8527827651989199, 'reg_lambda': 0.3315799968401767, 'max_depth': 6, 'max_leaves': 1, 'max_cat_to_onehot': 9}. Best is trial 5 with value: 54.936126708984375.\n",
"[I 2024-05-09 07:56:59,261] Trial 7 finished with value: 55.204200744628906 and parameters: {'n_estimators': 3, 'learning_rate': 0.6831416027240611, 'colsample_bytree': 0.5311840770388268, 'colsample_bynode': 0.9572535535110238, 'colsample_bylevel': 0.6846894032354778, 'reg_lambda': 0.6091211134408249, 'max_depth': 3, 'max_leaves': 2, 'max_cat_to_onehot': 5}. Best is trial 5 with value: 54.936126708984375.\n",
"[I 2024-05-09 07:57:37,674] Trial 8 finished with value: 54.93584442138672 and parameters: {'n_estimators': 4, 'learning_rate': 0.620742285616388, 'colsample_bytree': 0.7969398985157778, 'colsample_bynode': 0.9049707375663323, 'colsample_bylevel': 0.7209693969245297, 'reg_lambda': 0.6158847054585023, 'max_depth': 1, 'max_leaves': 0, 'max_cat_to_onehot': 10}. Best is trial 8 with value: 54.93584442138672.\n",
"[I 2024-05-09 07:57:50,310] Trial 9 finished with value: 57.76123809814453 and parameters: {'n_estimators': 3, 'learning_rate': 0.5475197727057007, 'colsample_bytree': 0.5381502848057452, 'colsample_bynode': 0.8514705732161596, 'colsample_bylevel': 0.9139277684007088, 'reg_lambda': 0.5117732009332318, 'max_depth': 4, 'max_leaves': 0, 'max_cat_to_onehot': 5}. Best is trial 8 with value: 54.93584442138672.\n"
]
}
],
"source": [
"# With n_jobs parameter, Optuna will launch [n_clusters] threads internally\n",
"# Each thread will deploy a training job to a Dask cluster\n",
"study.optimize(objective, n_trials=n_trials, n_jobs=n_clusters)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac5b3cba-87ba-4470-a166-b6a0815f85e4",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"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.12.8"
}
},
"nbformat": 4,
"nbformat_minor": 5
}