Skip to content

Instantly share code, notes, and snippets.

@praateekmahajan

praateekmahajan/dask_cuml_kmeans_multi_node.ipynb

Last active October 24, 2025 22:15
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save praateekmahajan/59f41e28ba8d5f72be33c3320cb9604e to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save praateekmahajan/59f41e28ba8d5f72be33c3320cb9604e to your computer and use it in GitHub Desktop.
Download ZIP
Using cuml 25.10+
Display the source blob
Display the rendered blob
Raw
dask_cuml_kmeans_multi_node.ipynb
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"id": "23f333d0-46f4-46bb-bc30-5c21bd94ff02",
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0,1,2,3\""
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "79a9c670-f388-40f1-93ff-67d99d8860c6",
"metadata": {},
"outputs": [],
"source": [
"import cudf\n",
"import cuml\n",
"import dask_cuda\n",
"import dask"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "d1d82bc2-10ba-41e8-a907-0238c9163065",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"cudf 25.10.00\n",
"cuml 25.10.00\n",
"dask_cuda 25.10.00\n",
"dask 2025.9.1\n"
]
}
],
"source": [
"for lib in [cudf, cuml, dask_cuda, dask]:\n",
" print(str(lib.__name__), lib.__version__)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "ebd79e5b-55f7-4193-8685-912e480131fa",
"metadata": {
"scrolled": true
},
"outputs": [],
"source": [
"from sklearn.datasets import make_blobs\n",
"from pathlib import Path\n",
"import numpy as np\n",
"import pandas as pd\n",
"import cudf\n",
"from sklearn.metrics import adjusted_rand_score\n"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "1adf42c4-7bbe-459e-80f2-c88dd5c3e39b",
"metadata": {},
"outputs": [],
"source": [
"N_CLUSTERS = 10\n",
"N_SAMPLES_PER_CLUSTER = 10_000\n",
"EMBEDDING_DIM = 1024\n",
"RANDOM_STATE = 42"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "4fbd2151-faaa-4eec-954c-40216ad7b913",
"metadata": {},
"outputs": [],
"source": [
"def create_clustered_dataset( # noqa: PLR0913\n",
" tmp_path: Path,\n",
" n_clusters: int = N_CLUSTERS,\n",
" n_samples_per_cluster: int = N_SAMPLES_PER_CLUSTER,\n",
" embedding_dim: int = EMBEDDING_DIM,\n",
" random_state: int = RANDOM_STATE,\n",
" file_format: str = \"parquet\",\n",
") -> tuple[Path, np.ndarray]:\n",
" \"\"\"Create a synthetic clustered dataset using sklearn make_blobs.\n",
"\n",
" Args:\n",
" tmp_path: Temporary directory path\n",
" n_clusters: Number of clusters to create\n",
" n_samples_per_cluster: Number of samples per cluster\n",
" embedding_dim: Dimensionality of embeddings\n",
" random_state: Random seed for reproducibility\n",
" file_format: Output file format ('parquet' or 'jsonl')\n",
"\n",
" Returns:\n",
" Tuple of (input_dir_path, true_labels_array)\n",
" \"\"\"\n",
" # Create clustered data using sklearn\n",
" X, y_true = make_blobs( # noqa: N806\n",
" n_samples=n_clusters * n_samples_per_cluster,\n",
" centers=n_clusters,\n",
" n_features=embedding_dim,\n",
" random_state=random_state,\n",
" cluster_std=0.5, # Reduced cluster standard deviation for tighter clusters\n",
" )\n",
"\n",
" # Normalize embeddings (same as KMeans stage will do)\n",
" X_normalized = X / np.linalg.norm(X, axis=1, keepdims=True) # noqa: N806\n",
"\n",
" # Create input directory\n",
" input_dir = tmp_path / \"input\"\n",
" input_dir.mkdir(parents=True, exist_ok=True)\n",
"\n",
" # Create dataframe with embeddings and IDs\n",
" num_files = 20 # Create multiple files to test file partitioning\n",
" samples_per_file = len(X_normalized) // num_files\n",
" rng = np.random.default_rng(random_state)\n",
"\n",
" for file_idx in range(num_files):\n",
" start_idx = file_idx * samples_per_file\n",
" end_idx = (file_idx + 1) * samples_per_file if file_idx < num_files - 1 else len(X_normalized)\n",
" df = pd.DataFrame(\n",
" {\n",
" \"id\": np.arange(start_idx, end_idx),\n",
" \"embeddings\": X_normalized[start_idx:end_idx].tolist(),\n",
" \"true_cluster\": y_true[start_idx:end_idx].tolist(),\n",
" }\n",
" )\n",
" df[\"random_col\"] = rng.integers(0, 100, size=len(df))\n",
"\n",
" if file_format == \"parquet\":\n",
" file_path = input_dir / f\"data_part_{file_idx:02d}.parquet\"\n",
" df.to_parquet(file_path, index=False)\n",
" elif file_format == \"jsonl\":\n",
" file_path = input_dir / f\"data_part_{file_idx:02d}.jsonl\"\n",
" df.to_json(file_path, orient=\"records\", lines=True)\n",
" else:\n",
" msg = f\"Unsupported file format: {file_format}\"\n",
" raise ValueError(msg)\n",
"\n",
" return input_dir, y_true\n",
"\n"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "e94bda8a-7785-4554-8254-7a5306a09a45",
"metadata": {},
"outputs": [],
"source": [
"import shutil\n",
"shutil.rmtree(\"./clustered_data/\", ignore_errors=True)"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "828381ae-b0bc-49bf-aab1-1c4234a1279f",
"metadata": {},
"outputs": [],
"source": [
"cluster_dir, y_true = create_clustered_dataset(Path(\"./clustered_data\"))"
]
},
{
"cell_type": "code",
"execution_count": 9,
"id": "225c7871-ee47-4500-826a-da8a87455a53",
"metadata": {},
"outputs": [],
"source": [
"import cuml\n",
"from typing import Any\n",
"\n",
"import cudf\n",
"import cupy as cp\n",
"import pylibcudf as plc\n",
"\n",
"single_gpu_kmeans = cuml.KMeans(\n",
" n_clusters=N_CLUSTERS,\n",
" init=\"k-means++\",\n",
" max_iter=300,\n",
" tol=1e-4,\n",
" random_state=RANDOM_STATE,\n",
" output_type=\"numpy\", # Use numpy output for easier comparison\n",
")\n",
"\n",
"def get_array_from_df(df: \"cudf.DataFrame\", embedding_col: str) -> \"cp.ndarray\":\n",
" \"\"\"\n",
" Convert a column of lists to a 2D array.\n",
" \"\"\"\n",
" return df[embedding_col].list.leaves.values.reshape(len(df), -1)\n",
"\n",
"def normalize_embeddings_col_in_df(df: \"cudf.DataFrame\", embedding_col: str) -> \"cudf.DataFrame\":\n",
" tensor = torch.Tensor(get_array_from_df(df, embedding_col))\n",
" normalized_tensor = tensor / torch.norm(tensor, dim=1, keepdim=True)\n",
" df[embedding_col] = create_list_series_from_1d_or_2d_ar(cp.asarray(normalized_tensor), index=df.index)\n",
" return df\n",
"\n",
"\n",
"def create_list_series_from_1d_or_2d_ar(ar: Any, index: cudf.Index) -> cudf.Series: # noqa: ANN401\n",
" arr = cp.asarray(ar)\n",
" if len(arr.shape) == 1:\n",
" arr = arr.reshape(-1, 1)\n",
" if not isinstance(index, (cudf.RangeIndex, cudf.Index, cudf.MultiIndex)):\n",
" index = cudf.Index(index)\n",
" return cudf.Series.from_pylibcudf(\n",
" plc.Column.from_cuda_array_interface(arr),\n",
" metadata={\"index\": index},\n",
" )\n"
]
},
{
"cell_type": "markdown",
"id": "fd7a75d2-f473-4b26-8822-a72ff575dc42",
"metadata": {},
"source": [
"## Ground Truth"
]
},
{
"cell_type": "code",
"execution_count": 10,
"id": "7f20e2ec-bf76-4f4a-aef1-91ec266dad4c",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"true_cluster\n",
"1 10000\n",
"2 10000\n",
"7 10000\n",
"9 10000\n",
"8 10000\n",
"3 10000\n",
"5 10000\n",
"6 10000\n",
"0 10000\n",
"4 10000\n",
"Name: count, dtype: int64"
]
},
"execution_count": 10,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"df = cudf.read_parquet(cluster_dir)\n",
"df['true_cluster'].value_counts()"
]
},
{
"cell_type": "markdown",
"id": "1d0fa828-6f58-4bd2-95f7-32c32be9d005",
"metadata": {},
"source": [
"## Single Node"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "be76060d-783c-4464-91e3-f4f1af06b925",
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"<style>#sk-container-id-1 {\n",
" /* Definition of color scheme common for light and dark mode */\n",
" --sklearn-color-text: #000;\n",
" --sklearn-color-text-muted: #666;\n",
" --sklearn-color-line: gray;\n",
" /* Definition of color scheme for unfitted estimators */\n",
" --sklearn-color-unfitted-level-0: #fff5e6;\n",
" --sklearn-color-unfitted-level-1: #f6e4d2;\n",
" --sklearn-color-unfitted-level-2: #ffe0b3;\n",
" --sklearn-color-unfitted-level-3: chocolate;\n",
" /* Definition of color scheme for fitted estimators */\n",
" --sklearn-color-fitted-level-0: #f0f8ff;\n",
" --sklearn-color-fitted-level-1: #d4ebff;\n",
" --sklearn-color-fitted-level-2: #b3dbfd;\n",
" --sklearn-color-fitted-level-3: cornflowerblue;\n",
"\n",
" /* Specific color for light theme */\n",
" --sklearn-color-text-on-default-background: var(--sg-text-color, var(--theme-code-foreground, var(--jp-content-font-color1, black)));\n",
" --sklearn-color-background: var(--sg-background-color, var(--theme-background, var(--jp-layout-color0, white)));\n",
" --sklearn-color-border-box: var(--sg-text-color, var(--theme-code-foreground, var(--jp-content-font-color1, black)));\n",
" --sklearn-color-icon: #696969;\n",
"\n",
" @media (prefers-color-scheme: dark) {\n",
" /* Redefinition of color scheme for dark theme */\n",
" --sklearn-color-text-on-default-background: var(--sg-text-color, var(--theme-code-foreground, var(--jp-content-font-color1, white)));\n",
" --sklearn-color-background: var(--sg-background-color, var(--theme-background, var(--jp-layout-color0, #111)));\n",
" --sklearn-color-border-box: var(--sg-text-color, var(--theme-code-foreground, var(--jp-content-font-color1, white)));\n",
" --sklearn-color-icon: #878787;\n",
" }\n",
"}\n",
"\n",
"#sk-container-id-1 {\n",
" color: var(--sklearn-color-text);\n",
"}\n",
"\n",
"#sk-container-id-1 pre {\n",
" padding: 0;\n",
"}\n",
"\n",
"#sk-container-id-1 input.sk-hidden--visually {\n",
" border: 0;\n",
" clip: rect(1px 1px 1px 1px);\n",
" clip: rect(1px, 1px, 1px, 1px);\n",
" height: 1px;\n",
" margin: -1px;\n",
" overflow: hidden;\n",
" padding: 0;\n",
" position: absolute;\n",
" width: 1px;\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-dashed-wrapped {\n",
" border: 1px dashed var(--sklearn-color-line);\n",
" margin: 0 0.4em 0.5em 0.4em;\n",
" box-sizing: border-box;\n",
" padding-bottom: 0.4em;\n",
" background-color: var(--sklearn-color-background);\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-container {\n",
" /* jupyter's `normalize.less` sets `[hidden] { display: none; }`\n",
" but bootstrap.min.css set `[hidden] { display: none !important; }`\n",
" so we also need the `!important` here to be able to override the\n",
" default hidden behavior on the sphinx rendered scikit-learn.org.\n",
" See: https://github.com/scikit-learn/scikit-learn/issues/21755 */\n",
" display: inline-block !important;\n",
" position: relative;\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-text-repr-fallback {\n",
" display: none;\n",
"}\n",
"\n",
"div.sk-parallel-item,\n",
"div.sk-serial,\n",
"div.sk-item {\n",
" /* draw centered vertical line to link estimators */\n",
" background-image: linear-gradient(var(--sklearn-color-text-on-default-background), var(--sklearn-color-text-on-default-background));\n",
" background-size: 2px 100%;\n",
" background-repeat: no-repeat;\n",
" background-position: center center;\n",
"}\n",
"\n",
"/* Parallel-specific style estimator block */\n",
"\n",
"#sk-container-id-1 div.sk-parallel-item::after {\n",
" content: \"\";\n",
" width: 100%;\n",
" border-bottom: 2px solid var(--sklearn-color-text-on-default-background);\n",
" flex-grow: 1;\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-parallel {\n",
" display: flex;\n",
" align-items: stretch;\n",
" justify-content: center;\n",
" background-color: var(--sklearn-color-background);\n",
" position: relative;\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-parallel-item {\n",
" display: flex;\n",
" flex-direction: column;\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-parallel-item:first-child::after {\n",
" align-self: flex-end;\n",
" width: 50%;\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-parallel-item:last-child::after {\n",
" align-self: flex-start;\n",
" width: 50%;\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-parallel-item:only-child::after {\n",
" width: 0;\n",
"}\n",
"\n",
"/* Serial-specific style estimator block */\n",
"\n",
"#sk-container-id-1 div.sk-serial {\n",
" display: flex;\n",
" flex-direction: column;\n",
" align-items: center;\n",
" background-color: var(--sklearn-color-background);\n",
" padding-right: 1em;\n",
" padding-left: 1em;\n",
"}\n",
"\n",
"\n",
"/* Toggleable style: style used for estimator/Pipeline/ColumnTransformer box that is\n",
"clickable and can be expanded/collapsed.\n",
"- Pipeline and ColumnTransformer use this feature and define the default style\n",
"- Estimators will overwrite some part of the style using the `sk-estimator` class\n",
"*/\n",
"\n",
"/* Pipeline and ColumnTransformer style (default) */\n",
"\n",
"#sk-container-id-1 div.sk-toggleable {\n",
" /* Default theme specific background. It is overwritten whether we have a\n",
" specific estimator or a Pipeline/ColumnTransformer */\n",
" background-color: var(--sklearn-color-background);\n",
"}\n",
"\n",
"/* Toggleable label */\n",
"#sk-container-id-1 label.sk-toggleable__label {\n",
" cursor: pointer;\n",
" display: flex;\n",
" width: 100%;\n",
" margin-bottom: 0;\n",
" padding: 0.5em;\n",
" box-sizing: border-box;\n",
" text-align: center;\n",
" align-items: start;\n",
" justify-content: space-between;\n",
" gap: 0.5em;\n",
"}\n",
"\n",
"#sk-container-id-1 label.sk-toggleable__label .caption {\n",
" font-size: 0.6rem;\n",
" font-weight: lighter;\n",
" color: var(--sklearn-color-text-muted);\n",
"}\n",
"\n",
"#sk-container-id-1 label.sk-toggleable__label-arrow:before {\n",
" /* Arrow on the left of the label */\n",
" content: \"▸\";\n",
" float: left;\n",
" margin-right: 0.25em;\n",
" color: var(--sklearn-color-icon);\n",
"}\n",
"\n",
"#sk-container-id-1 label.sk-toggleable__label-arrow:hover:before {\n",
" color: var(--sklearn-color-text);\n",
"}\n",
"\n",
"/* Toggleable content - dropdown */\n",
"\n",
"#sk-container-id-1 div.sk-toggleable__content {\n",
" display: none;\n",
" text-align: left;\n",
" /* unfitted */\n",
" background-color: var(--sklearn-color-unfitted-level-0);\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-toggleable__content.fitted {\n",
" /* fitted */\n",
" background-color: var(--sklearn-color-fitted-level-0);\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-toggleable__content pre {\n",
" margin: 0.2em;\n",
" border-radius: 0.25em;\n",
" color: var(--sklearn-color-text);\n",
" /* unfitted */\n",
" background-color: var(--sklearn-color-unfitted-level-0);\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-toggleable__content.fitted pre {\n",
" /* unfitted */\n",
" background-color: var(--sklearn-color-fitted-level-0);\n",
"}\n",
"\n",
"#sk-container-id-1 input.sk-toggleable__control:checked~div.sk-toggleable__content {\n",
" /* Expand drop-down */\n",
" display: block;\n",
" width: 100%;\n",
" overflow: visible;\n",
"}\n",
"\n",
"#sk-container-id-1 input.sk-toggleable__control:checked~label.sk-toggleable__label-arrow:before {\n",
" content: \"▾\";\n",
"}\n",
"\n",
"/* Pipeline/ColumnTransformer-specific style */\n",
"\n",
"#sk-container-id-1 div.sk-label input.sk-toggleable__control:checked~label.sk-toggleable__label {\n",
" color: var(--sklearn-color-text);\n",
" background-color: var(--sklearn-color-unfitted-level-2);\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-label.fitted input.sk-toggleable__control:checked~label.sk-toggleable__label {\n",
" background-color: var(--sklearn-color-fitted-level-2);\n",
"}\n",
"\n",
"/* Estimator-specific style */\n",
"\n",
"/* Colorize estimator box */\n",
"#sk-container-id-1 div.sk-estimator input.sk-toggleable__control:checked~label.sk-toggleable__label {\n",
" /* unfitted */\n",
" background-color: var(--sklearn-color-unfitted-level-2);\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-estimator.fitted input.sk-toggleable__control:checked~label.sk-toggleable__label {\n",
" /* fitted */\n",
" background-color: var(--sklearn-color-fitted-level-2);\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-label label.sk-toggleable__label,\n",
"#sk-container-id-1 div.sk-label label {\n",
" /* The background is the default theme color */\n",
" color: var(--sklearn-color-text-on-default-background);\n",
"}\n",
"\n",
"/* On hover, darken the color of the background */\n",
"#sk-container-id-1 div.sk-label:hover label.sk-toggleable__label {\n",
" color: var(--sklearn-color-text);\n",
" background-color: var(--sklearn-color-unfitted-level-2);\n",
"}\n",
"\n",
"/* Label box, darken color on hover, fitted */\n",
"#sk-container-id-1 div.sk-label.fitted:hover label.sk-toggleable__label.fitted {\n",
" color: var(--sklearn-color-text);\n",
" background-color: var(--sklearn-color-fitted-level-2);\n",
"}\n",
"\n",
"/* Estimator label */\n",
"\n",
"#sk-container-id-1 div.sk-label label {\n",
" font-family: monospace;\n",
" font-weight: bold;\n",
" display: inline-block;\n",
" line-height: 1.2em;\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-label-container {\n",
" text-align: center;\n",
"}\n",
"\n",
"/* Estimator-specific */\n",
"#sk-container-id-1 div.sk-estimator {\n",
" font-family: monospace;\n",
" border: 1px dotted var(--sklearn-color-border-box);\n",
" border-radius: 0.25em;\n",
" box-sizing: border-box;\n",
" margin-bottom: 0.5em;\n",
" /* unfitted */\n",
" background-color: var(--sklearn-color-unfitted-level-0);\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-estimator.fitted {\n",
" /* fitted */\n",
" background-color: var(--sklearn-color-fitted-level-0);\n",
"}\n",
"\n",
"/* on hover */\n",
"#sk-container-id-1 div.sk-estimator:hover {\n",
" /* unfitted */\n",
" background-color: var(--sklearn-color-unfitted-level-2);\n",
"}\n",
"\n",
"#sk-container-id-1 div.sk-estimator.fitted:hover {\n",
" /* fitted */\n",
" background-color: var(--sklearn-color-fitted-level-2);\n",
"}\n",
"\n",
"/* Specification for estimator info (e.g. \"i\" and \"?\") */\n",
"\n",
"/* Common style for \"i\" and \"?\" */\n",
"\n",
".sk-estimator-doc-link,\n",
"a:link.sk-estimator-doc-link,\n",
"a:visited.sk-estimator-doc-link {\n",
" float: right;\n",
" font-size: smaller;\n",
" line-height: 1em;\n",
" font-family: monospace;\n",
" background-color: var(--sklearn-color-background);\n",
" border-radius: 1em;\n",
" height: 1em;\n",
" width: 1em;\n",
" text-decoration: none !important;\n",
" margin-left: 0.5em;\n",
" text-align: center;\n",
" /* unfitted */\n",
" border: var(--sklearn-color-unfitted-level-1) 1pt solid;\n",
" color: var(--sklearn-color-unfitted-level-1);\n",
"}\n",
"\n",
".sk-estimator-doc-link.fitted,\n",
"a:link.sk-estimator-doc-link.fitted,\n",
"a:visited.sk-estimator-doc-link.fitted {\n",
" /* fitted */\n",
" border: var(--sklearn-color-fitted-level-1) 1pt solid;\n",
" color: var(--sklearn-color-fitted-level-1);\n",
"}\n",
"\n",
"/* On hover */\n",
"div.sk-estimator:hover .sk-estimator-doc-link:hover,\n",
".sk-estimator-doc-link:hover,\n",
"div.sk-label-container:hover .sk-estimator-doc-link:hover,\n",
".sk-estimator-doc-link:hover {\n",
" /* unfitted */\n",
" background-color: var(--sklearn-color-unfitted-level-3);\n",
" color: var(--sklearn-color-background);\n",
" text-decoration: none;\n",
"}\n",
"\n",
"div.sk-estimator.fitted:hover .sk-estimator-doc-link.fitted:hover,\n",
".sk-estimator-doc-link.fitted:hover,\n",
"div.sk-label-container:hover .sk-estimator-doc-link.fitted:hover,\n",
".sk-estimator-doc-link.fitted:hover {\n",
" /* fitted */\n",
" background-color: var(--sklearn-color-fitted-level-3);\n",
" color: var(--sklearn-color-background);\n",
" text-decoration: none;\n",
"}\n",
"\n",
"/* Span, style for the box shown on hovering the info icon */\n",
".sk-estimator-doc-link span {\n",
" display: none;\n",
" z-index: 9999;\n",
" position: relative;\n",
" font-weight: normal;\n",
" right: .2ex;\n",
" padding: .5ex;\n",
" margin: .5ex;\n",
" width: min-content;\n",
" min-width: 20ex;\n",
" max-width: 50ex;\n",
" color: var(--sklearn-color-text);\n",
" box-shadow: 2pt 2pt 4pt #999;\n",
" /* unfitted */\n",
" background: var(--sklearn-color-unfitted-level-0);\n",
" border: .5pt solid var(--sklearn-color-unfitted-level-3);\n",
"}\n",
"\n",
".sk-estimator-doc-link.fitted span {\n",
" /* fitted */\n",
" background: var(--sklearn-color-fitted-level-0);\n",
" border: var(--sklearn-color-fitted-level-3);\n",
"}\n",
"\n",
".sk-estimator-doc-link:hover span {\n",
" display: block;\n",
"}\n",
"\n",
"/* \"?\"-specific style due to the `<a>` HTML tag */\n",
"\n",
"#sk-container-id-1 a.estimator_doc_link {\n",
" float: right;\n",
" font-size: 1rem;\n",
" line-height: 1em;\n",
" font-family: monospace;\n",
" background-color: var(--sklearn-color-background);\n",
" border-radius: 1rem;\n",
" height: 1rem;\n",
" width: 1rem;\n",
" text-decoration: none;\n",
" /* unfitted */\n",
" color: var(--sklearn-color-unfitted-level-1);\n",
" border: var(--sklearn-color-unfitted-level-1) 1pt solid;\n",
"}\n",
"\n",
"#sk-container-id-1 a.estimator_doc_link.fitted {\n",
" /* fitted */\n",
" border: var(--sklearn-color-fitted-level-1) 1pt solid;\n",
" color: var(--sklearn-color-fitted-level-1);\n",
"}\n",
"\n",
"/* On hover */\n",
"#sk-container-id-1 a.estimator_doc_link:hover {\n",
" /* unfitted */\n",
" background-color: var(--sklearn-color-unfitted-level-3);\n",
" color: var(--sklearn-color-background);\n",
" text-decoration: none;\n",
"}\n",
"\n",
"#sk-container-id-1 a.estimator_doc_link.fitted:hover {\n",
" /* fitted */\n",
" background-color: var(--sklearn-color-fitted-level-3);\n",
"}\n",
"\n",
".estimator-table summary {\n",
" padding: .5rem;\n",
" font-family: monospace;\n",
" cursor: pointer;\n",
"}\n",
"\n",
".estimator-table details[open] {\n",
" padding-left: 0.1rem;\n",
" padding-right: 0.1rem;\n",
" padding-bottom: 0.3rem;\n",
"}\n",
"\n",
".estimator-table .parameters-table {\n",
" margin-left: auto !important;\n",
" margin-right: auto !important;\n",
"}\n",
"\n",
".estimator-table .parameters-table tr:nth-child(odd) {\n",
" background-color: #fff;\n",
"}\n",
"\n",
".estimator-table .parameters-table tr:nth-child(even) {\n",
" background-color: #f6f6f6;\n",
"}\n",
"\n",
".estimator-table .parameters-table tr:hover {\n",
" background-color: #e0e0e0;\n",
"}\n",
"\n",
".estimator-table table td {\n",
" border: 1px solid rgba(106, 105, 104, 0.232);\n",
"}\n",
"\n",
".user-set td {\n",
" color:rgb(255, 94, 0);\n",
" text-align: left;\n",
"}\n",
"\n",
".user-set td.value pre {\n",
" color:rgb(255, 94, 0) !important;\n",
" background-color: transparent !important;\n",
"}\n",
"\n",
".default td {\n",
" color: black;\n",
" text-align: left;\n",
"}\n",
"\n",
".user-set td i,\n",
".default td i {\n",
" color: black;\n",
"}\n",
"\n",
".copy-paste-icon {\n",
" background-image: url(data:image/svg+xml;base64,PHN2ZyB4bWxucz0iaHR0cDovL3d3dy53My5vcmcvMjAwMC9zdmciIHZpZXdCb3g9IjAgMCA0NDggNTEyIj48IS0tIUZvbnQgQXdlc29tZSBGcmVlIDYuNy4yIGJ5IEBmb250YXdlc29tZSAtIGh0dHBzOi8vZm9udGF3ZXNvbWUuY29tIExpY2Vuc2UgLSBodHRwczovL2ZvbnRhd2Vzb21lLmNvbS9saWNlbnNlL2ZyZWUgQ29weXJpZ2h0IDIwMjUgRm9udGljb25zLCBJbmMuLS0+PHBhdGggZD0iTTIwOCAwTDMzMi4xIDBjMTIuNyAwIDI0LjkgNS4xIDMzLjkgMTQuMWw2Ny45IDY3LjljOSA5IDE0LjEgMjEuMiAxNC4xIDMzLjlMNDQ4IDMzNmMwIDI2LjUtMjEuNSA0OC00OCA0OGwtMTkyIDBjLTI2LjUgMC00OC0yMS41LTQ4LTQ4bDAtMjg4YzAtMjYuNSAyMS41LTQ4IDQ4LTQ4ek00OCAxMjhsODAgMCAwIDY0LTY0IDAgMCAyNTYgMTkyIDAgMC0zMiA2NCAwIDAgNDhjMCAyNi41LTIxLjUgNDgtNDggNDhMNDggNTEyYy0yNi41IDAtNDgtMjEuNS00OC00OEwwIDE3NmMwLTI2LjUgMjEuNS00OCA0OC00OHoiLz48L3N2Zz4=);\n",
" background-repeat: no-repeat;\n",
" background-size: 14px 14px;\n",
" background-position: 0;\n",
" display: inline-block;\n",
" width: 14px;\n",
" height: 14px;\n",
" cursor: pointer;\n",
"}\n",
"</style><body><div id=\"sk-container-id-1\" class=\"sk-top-container\"><div class=\"sk-text-repr-fallback\"><pre>KMeans()</pre><b>In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook. <br />On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.</b></div><div class=\"sk-container\" hidden><div class=\"sk-item\"><div class=\"sk-estimator fitted sk-toggleable\"><input class=\"sk-toggleable__control sk-hidden--visually\" id=\"sk-estimator-id-1\" type=\"checkbox\" checked><label for=\"sk-estimator-id-1\" class=\"sk-toggleable__label fitted sk-toggleable__label-arrow\"><div><div>KMeans</div></div><div><span class=\"sk-estimator-doc-link fitted\">i<span>Fitted</span></span></div></label><div class=\"sk-toggleable__content fitted\" data-param-prefix=\"\"><pre>KMeans()</pre></div></div></div></div></div><script>function copyToClipboard(text, element) {\n",
" // Get the parameter prefix from the closest toggleable content\n",
" const toggleableContent = element.closest('.sk-toggleable__content');\n",
" const paramPrefix = toggleableContent ? toggleableContent.dataset.paramPrefix : '';\n",
" const fullParamName = paramPrefix ? `${paramPrefix}${text}` : text;\n",
"\n",
" const originalStyle = element.style;\n",
" const computedStyle = window.getComputedStyle(element);\n",
" const originalWidth = computedStyle.width;\n",
" const originalHTML = element.innerHTML.replace('Copied!', '');\n",
"\n",
" navigator.clipboard.writeText(fullParamName)\n",
" .then(() => {\n",
" element.style.width = originalWidth;\n",
" element.style.color = 'green';\n",
" element.innerHTML = \"Copied!\";\n",
"\n",
" setTimeout(() => {\n",
" element.innerHTML = originalHTML;\n",
" element.style = originalStyle;\n",
" }, 2000);\n",
" })\n",
" .catch(err => {\n",
" console.error('Failed to copy:', err);\n",
" element.style.color = 'red';\n",
" element.innerHTML = \"Failed!\";\n",
" setTimeout(() => {\n",
" element.innerHTML = originalHTML;\n",
" element.style = originalStyle;\n",
" }, 2000);\n",
" });\n",
" return false;\n",
"}\n",
"\n",
"document.querySelectorAll('.fa-regular.fa-copy').forEach(function(element) {\n",
" const toggleableContent = element.closest('.sk-toggleable__content');\n",
" const paramPrefix = toggleableContent ? toggleableContent.dataset.paramPrefix : '';\n",
" const paramName = element.parentElement.nextElementSibling.textContent.trim();\n",
" const fullParamName = paramPrefix ? `${paramPrefix}${paramName}` : paramName;\n",
"\n",
" element.setAttribute('title', fullParamName);\n",
"});\n",
"</script></body>"
],
"text/plain": [
"KMeans()"
]
},
"execution_count": 11,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"embeddings = get_array_from_df(df, \"embeddings\")\n",
"single_gpu_kmeans.fit(embeddings)"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "63f4f736-c2c7-42e3-9d92-d4586013bf7d",
"metadata": {},
"outputs": [],
"source": [
"df[\"single_gpu_centroid\"] = single_gpu_kmeans.predict(embeddings)"
]
},
{
"cell_type": "code",
"execution_count": 13,
"id": "8e28e1f6-f3ad-4ede-b22f-c528917bf845",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"single_gpu_centroid\n",
"3 10000\n",
"4 10000\n",
"2 10000\n",
"9 10000\n",
"0 10000\n",
"5 10000\n",
"6 10000\n",
"7 10000\n",
"1 10000\n",
"8 10000\n",
"Name: count, dtype: int64"
]
},
"execution_count": 13,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"df['single_gpu_centroid'].value_counts()"
]
},
{
"cell_type": "code",
"execution_count": 14,
"id": "f638e739-0d79-416d-b207-0ef29dad2252",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"1.0"
]
},
"execution_count": 14,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"adjusted_rand_score(\n",
" df[\"single_gpu_centroid\"].to_pandas().to_numpy(), \n",
" df['true_cluster'].to_pandas().to_numpy()\n",
")"
]
},
{
"cell_type": "markdown",
"id": "18b52a31-b7d4-4bf6-b39f-15e215c0c943",
"metadata": {},
"source": [
"## Multi Node Dask"
]
},
{
"cell_type": "code",
"execution_count": 15,
"id": "7f61e070-0645-430e-b7f4-f66dc7fb3b86",
"metadata": {},
"outputs": [],
"source": [
"from dask.distributed import Client, LocalCluster\n",
"from dask_cuda import LocalCUDACluster\n",
"import dask_cudf\n",
"\n",
"# Create a CUDA cluster with 4 workers (one per GPU)\n",
"cluster = LocalCUDACluster(\n",
" n_workers=4,\n",
" threads_per_worker=1,\n",
" device_memory_limit='40GB' # Adjust based on your GPU memory\n",
")\n",
"client = Client(cluster)\n"
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "28b95a53-68e0-420b-9eff-9b70ede4fa4f",
"metadata": {},
"outputs": [],
"source": [
"import cupy as cp\n",
"from cuml.dask.cluster import KMeans as daskKMeans\n",
"\n",
"ddf = dask_cudf.read_parquet(cluster_dir).optimize()\n",
"\n",
"cupy_normalized_darr = ddf.map_partitions(\n",
" get_array_from_df, \"embeddings\", meta=cp.ndarray([1, 1])\n",
")\n",
"cupy_normalized_darr.compute_chunk_sizes()\n",
"\n",
"\n",
"# Initialize distributed KMeans\n",
"kmeans = daskKMeans(\n",
" n_clusters=N_CLUSTERS,\n",
" init=\"k-means++\",\n",
" max_iter=300,\n",
" tol=1e-4,\n",
" random_state=RANDOM_STATE,\n",
")\n"
]
},
{
"cell_type": "code",
"execution_count": 17,
"id": "404e6b93-1599-4252-96b7-55d7111fb45b",
"metadata": {},
"outputs": [],
"source": [
"kmeans.fit(cupy_normalized_darr)\n",
"ddf[\"nearest_cent\"] = kmeans.predict(cupy_normalized_darr).astype(np.int32)"
]
},
{
"cell_type": "code",
"execution_count": 18,
"id": "2d9f34ec-21d5-43df-b0ce-c35355162b33",
"metadata": {},
"outputs": [],
"source": [
"multi_node_cluster = ddf[\"nearest_cent\"].compute()"
]
},
{
"cell_type": "code",
"execution_count": 19,
"id": "640f79d3-009e-4671-8703-c053414325d1",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"nearest_cent\n",
"3 20000\n",
"2 20000\n",
"8 20000\n",
"1 10000\n",
"5 9999\n",
"6 7941\n",
"0 5906\n",
"7 4094\n",
"4 2059\n",
"9 1\n",
"Name: count, dtype: int64"
]
},
"execution_count": 19,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"multi_node_cluster.value_counts()"
]
},
{
"cell_type": "code",
"execution_count": 20,
"id": "4a1699a3-36b1-4146-b064-2546f81a6ac3",
"metadata": {
"scrolled": true
},
"outputs": [
{
"data": {
"text/plain": [
"0.6925127175700695"
]
},
"execution_count": 20,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"adjusted_rand_score(\n",
" multi_node_cluster.to_pandas().to_numpy(), \n",
" df['true_cluster'].to_pandas().to_numpy()\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 21,
"id": "8c52d5b3-1627-47b0-b0d6-d38b0c2f8ca1",
"metadata": {},
"outputs": [],
"source": [
"client.shutdown()"
]
},
{
"cell_type": "markdown",
"id": "092cda92-dc63-437b-a556-f3fc0f4a623d",
"metadata": {},
"source": [
"## Ray KMeans"
]
},
{
"cell_type": "code",
"execution_count": 22,
"id": "802fc19b-1291-402f-a20b-aa3e8a401787",
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"2025-10-24 15:12:28,244\tINFO worker.py:2004 -- Started a local Ray instance. View the dashboard at \u001b[1m\u001b[32mhttp://127.0.0.1:8265 \u001b[39m\u001b[22m\n",
"/tmp/ipykernel_1864803/1854117350.py:2: FutureWarning: Tip: In future versions of Ray, Ray will no longer override accelerator visible devices env var if num_gpus=0 or num_gpus=None (default). To enable this behavior and turn off this error message, set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0\n",
" ray.init(num_gpus=4, _temp_dir=\"/raid/praateekm/ray_tmp_1\")\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "167cab40a93b480ab0f294507b4dce38",
"version_major": 2,
"version_minor": 0
},
"text/html": [
"<div class=\"lm-Widget p-Widget lm-Panel p-Panel jp-Cell-outputWrapper\">\n",
" <div style=\"margin-left: 50px;display: flex;flex-direction: row;align-items: center\">\n",
" <div class=\"jp-RenderedHTMLCommon\" style=\"display: flex; flex-direction: row;\">\n",
" <svg viewBox=\"0 0 567 224\" fill=\"none\" xmlns=\"http://www.w3.org/2000/svg\" style=\"height: 3em;\">\n",
" <g clip-path=\"url(#clip0_4338_178347)\">\n",
" <path d=\"M341.29 165.561H355.29L330.13 129.051C345.63 123.991 354.21 112.051 354.21 94.2307C354.21 71.3707 338.72 58.1807 311.88 58.1807H271V165.561H283.27V131.661H311.8C314.25 131.661 316.71 131.501 319.01 131.351L341.25 165.561H341.29ZM283.29 119.851V70.0007H311.82C331.3 70.0007 342.34 78.2907 342.34 94.5507C342.34 111.271 331.34 119.861 311.82 119.861L283.29 119.851ZM451.4 138.411L463.4 165.561H476.74L428.74 58.1807H416L367.83 165.561H380.83L392.83 138.411H451.4ZM446.19 126.601H398L422 72.1407L446.24 126.601H446.19ZM526.11 128.741L566.91 58.1807H554.35L519.99 114.181L485.17 58.1807H472.44L514.01 129.181V165.541H526.13V128.741H526.11Z\" fill=\"var(--jp-ui-font-color0)\"/>\n",
" <path d=\"M82.35 104.44C84.0187 97.8827 87.8248 92.0678 93.1671 87.9146C98.5094 83.7614 105.083 81.5067 111.85 81.5067C118.617 81.5067 125.191 83.7614 130.533 87.9146C135.875 92.0678 139.681 97.8827 141.35 104.44H163.75C164.476 101.562 165.622 98.8057 167.15 96.2605L127.45 56.5605C121.071 60.3522 113.526 61.6823 106.235 60.3005C98.9443 58.9187 92.4094 54.9203 87.8602 49.0574C83.3109 43.1946 81.0609 35.8714 81.5332 28.4656C82.0056 21.0599 85.1679 14.0819 90.4252 8.8446C95.6824 3.60726 102.672 0.471508 110.08 0.0272655C117.487 -0.416977 124.802 1.86091 130.647 6.4324C136.493 11.0039 140.467 17.5539 141.821 24.8501C143.175 32.1463 141.816 39.6859 138 46.0505L177.69 85.7505C182.31 82.9877 187.58 81.4995 192.962 81.4375C198.345 81.3755 203.648 82.742 208.33 85.3976C213.012 88.0532 216.907 91.9029 219.616 96.5544C222.326 101.206 223.753 106.492 223.753 111.875C223.753 117.258 222.326 122.545 219.616 127.197C216.907 131.848 213.012 135.698 208.33 138.353C203.648 141.009 198.345 142.375 192.962 142.313C187.58 142.251 182.31 140.763 177.69 138L138 177.7C141.808 184.071 143.155 191.614 141.79 198.91C140.424 206.205 136.44 212.75 130.585 217.313C124.731 221.875 117.412 224.141 110.004 223.683C102.596 223.226 95.6103 220.077 90.3621 214.828C85.1139 209.58 81.9647 202.595 81.5072 195.187C81.0497 187.779 83.3154 180.459 87.878 174.605C92.4405 168.751 98.9853 164.766 106.281 163.401C113.576 162.035 121.119 163.383 127.49 167.19L167.19 127.49C165.664 124.941 164.518 122.182 163.79 119.3H141.39C139.721 125.858 135.915 131.673 130.573 135.826C125.231 139.98 118.657 142.234 111.89 142.234C105.123 142.234 98.5494 139.98 93.2071 135.826C87.8648 131.673 84.0587 125.858 82.39 119.3H60C58.1878 126.495 53.8086 132.78 47.6863 136.971C41.5641 141.163 34.1211 142.972 26.7579 142.059C19.3947 141.146 12.6191 137.574 7.70605 132.014C2.79302 126.454 0.0813599 119.29 0.0813599 111.87C0.0813599 104.451 2.79302 97.2871 7.70605 91.7272C12.6191 86.1673 19.3947 82.5947 26.7579 81.6817C34.1211 80.7686 41.5641 82.5781 47.6863 86.7696C53.8086 90.9611 58.1878 97.2456 60 104.44H82.35ZM100.86 204.32C103.407 206.868 106.759 208.453 110.345 208.806C113.93 209.159 117.527 208.258 120.522 206.256C123.517 204.254 125.725 201.276 126.771 197.828C127.816 194.38 127.633 190.677 126.253 187.349C124.874 184.021 122.383 181.274 119.205 179.577C116.027 177.88 112.359 177.337 108.826 178.042C105.293 178.746 102.113 180.654 99.8291 183.44C97.5451 186.226 96.2979 189.718 96.3 193.32C96.2985 195.364 96.7006 197.388 97.4831 199.275C98.2656 201.163 99.4132 202.877 100.86 204.32ZM204.32 122.88C206.868 120.333 208.453 116.981 208.806 113.396C209.159 109.811 208.258 106.214 206.256 103.219C204.254 100.223 201.275 98.0151 197.827 96.97C194.38 95.9249 190.676 96.1077 187.348 97.4873C184.02 98.8669 181.274 101.358 179.577 104.536C177.879 107.714 177.337 111.382 178.041 114.915C178.746 118.448 180.653 121.627 183.439 123.911C186.226 126.195 189.717 127.443 193.32 127.44C195.364 127.443 197.388 127.042 199.275 126.259C201.163 125.476 202.878 124.328 204.32 122.88ZM122.88 19.4205C120.333 16.8729 116.981 15.2876 113.395 14.9347C109.81 14.5817 106.213 15.483 103.218 17.4849C100.223 19.4868 98.0146 22.4654 96.9696 25.9131C95.9245 29.3608 96.1073 33.0642 97.4869 36.3922C98.8665 39.7202 101.358 42.4668 104.535 44.1639C107.713 45.861 111.381 46.4036 114.914 45.6992C118.447 44.9949 121.627 43.0871 123.911 40.301C126.195 37.515 127.442 34.0231 127.44 30.4205C127.44 28.3772 127.038 26.3539 126.255 24.4664C125.473 22.5788 124.326 20.8642 122.88 19.4205ZM19.42 100.86C16.8725 103.408 15.2872 106.76 14.9342 110.345C14.5813 113.93 15.4826 117.527 17.4844 120.522C19.4863 123.518 22.4649 125.726 25.9127 126.771C29.3604 127.816 33.0638 127.633 36.3918 126.254C39.7198 124.874 42.4664 122.383 44.1635 119.205C45.8606 116.027 46.4032 112.359 45.6988 108.826C44.9944 105.293 43.0866 102.114 40.3006 99.8296C37.5145 97.5455 34.0227 96.2983 30.42 96.3005C26.2938 96.3018 22.337 97.9421 19.42 100.86ZM100.86 100.86C98.3125 103.408 96.7272 106.76 96.3742 110.345C96.0213 113.93 96.9226 117.527 98.9244 120.522C100.926 123.518 103.905 125.726 107.353 126.771C110.8 127.816 114.504 127.633 117.832 126.254C121.16 124.874 123.906 122.383 125.604 119.205C127.301 116.027 127.843 112.359 127.139 108.826C126.434 105.293 124.527 102.114 121.741 99.8296C118.955 97.5455 115.463 96.2983 111.86 96.3005C109.817 96.299 107.793 96.701 105.905 97.4835C104.018 98.2661 102.303 99.4136 100.86 100.86Z\" fill=\"#00AEEF\"/>\n",
" </g>\n",
" <defs>\n",
" <clipPath id=\"clip0_4338_178347\">\n",
" <rect width=\"566.93\" height=\"223.75\" fill=\"white\"/>\n",
" </clipPath>\n",
" </defs>\n",
" </svg>\n",
"</div>\n",
"\n",
" <table class=\"jp-RenderedHTMLCommon\" style=\"border-collapse: collapse;color: var(--jp-ui-font-color1);font-size: var(--jp-ui-font-size1);\">\n",
" <tr>\n",
" <td style=\"text-align: left\"><b>Python version:</b></td>\n",
" <td style=\"text-align: left\"><b>3.12.11</b></td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"text-align: left\"><b>Ray version:</b></td>\n",
" <td style=\"text-align: left\"><b>2.50.1</b></td>\n",
" </tr>\n",
" <tr>\n",
" <td style=\"text-align: left\"><b>Dashboard:</b></td>\n",
" <td style=\"text-align: left\"><b><a href=\"http://127.0.0.1:8265\" target=\"_blank\">http://127.0.0.1:8265</a></b></td>\n",
"</tr>\n",
"\n",
"</table>\n",
"\n",
" </div>\n",
"</div>\n"
],
"text/plain": [
"RayContext(dashboard_url='127.0.0.1:8265', python_version='3.12.11', ray_version='2.50.1', ray_commit='7cf6817996f5304b5c808453a997fc1570dcde25')"
]
},
"execution_count": 22,
"metadata": {},
"output_type": "execute_result"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\u001b[36m(RAFTKMeansActor pid=1867442)\u001b[0m Initialized RAFTActor-1 (rank=1, is_root=False)\n",
"\u001b[36m(RAFTKMeansActor pid=1867426)\u001b[0m Initialized RAFTActor-0 (rank=0, is_root=True)\n",
"\u001b[36m(RAFTKMeansActor pid=1867426)\u001b[0m RAFTActor-0: RAFT setup complete\n"
]
}
],
"source": [
"import ray\n",
"ray.init(num_gpus=4, _temp_dir=\"/raid/praateekm/ray_tmp_1\")\n",
"# ray.shutdown()"
]
},
{
"cell_type": "code",
"execution_count": 23,
"id": "c0032782-1f92-4ce6-9198-b91aab68d087",
"metadata": {},
"outputs": [],
"source": [
"import uuid\n",
"\n",
"from raft_dask.common.nccl import nccl\n",
"class Comms:\n",
" valid_nccl_placements = \"ray-actor\"\n",
" def __init__(\n",
" self,\n",
" nccl_root_location: str = \"ray-actor\",\n",
" ) -> None:\n",
" self.nccl_root_location = nccl_root_location.lower()\n",
" if self.nccl_root_location not in Comms.valid_nccl_placements:\n",
" msg = f\"nccl_root_location must be one of: {Comms.valid_nccl_placements}\"\n",
" raise ValueError(msg)\n",
" self.sessionId = uuid.uuid4().bytes\n",
" \n",
" def create_nccl_uniqueid(self) -> None:\n",
" self.uniqueId = nccl.get_unique_id()\n",
" \n",
" def init(self) -> None:\n",
" self.create_nccl_uniqueid()\n"
]
},
{
"cell_type": "code",
"execution_count": 24,
"id": "803e752a-0d75-4714-9d4b-f18ad5881cf9",
"metadata": {},
"outputs": [],
"source": [
"@ray.remote(num_gpus=1)\n",
"class RAFTKMeansActor:\n",
" \"\"\"RAFT-enabled KMeans actor - coordinates with other actors via NCCL\"\"\"\n",
" \n",
" def __init__(self, rank: int, world_size: int, n_clusters: int, random_state: int, actor_name_prefix: str = \"RAFT\"):\n",
" \"\"\"Initialize RAFT actor with NCCL communication\"\"\"\n",
" self.rank = rank\n",
" self.world_size = world_size\n",
" self.n_clusters = n_clusters\n",
" self.random_state = random_state\n",
" self.is_root = (rank == 0)\n",
" self.name = f\"{actor_name_prefix}Actor-{rank}\"\n",
" self.actor_name_prefix = actor_name_prefix\n",
" \n",
" self.cb = Comms(nccl_root_location=\"ray-actor\")\n",
" self.cb.init()\n",
" self.unique_id = self.cb.uniqueId\n",
" self.root_unique_id = self.unique_id if self.is_root else None\n",
" print(f\"Initialized {self.name} (rank={rank}, is_root={self.is_root})\")\n",
" \n",
" def get_name(self):\n",
" return self.name\n",
" \n",
" def broadcast_root_unique_id(self):\n",
" \"\"\"Broadcast the root unique ID to all actors (matches raft_adapter.py lines 92-108)\"\"\"\n",
" if self.is_root:\n",
" actor_handles = [\n",
" ray.get_actor(name=f\"{self.actor_name_prefix}Actor-{i}\", namespace=None)\n",
" for i in range(1, self.world_size)\n",
" ]\n",
" futures = [actor.set_root_unique_id.remote(self.root_unique_id) for actor in actor_handles]\n",
" \n",
" # Block until all futures complete\n",
" ray.get(futures)\n",
" else:\n",
" raise RuntimeError(\"This method should only be called by the root\")\n",
" \n",
" def set_root_unique_id(self, root_unique_id: int):\n",
" \"\"\"Set the root unique ID (matches raft_adapter.py lines 110-120)\"\"\"\n",
" if self.root_unique_id is None:\n",
" self.root_unique_id = root_unique_id\n",
"\n",
" def setup_raft(self):\n",
" \"\"\"Setup RAFT handle with NCCL (matches raft_adapter.py lines 122-157)\"\"\"\n",
" if self.root_unique_id is None:\n",
" raise RuntimeError(\n",
" \"The unique ID of root is not set. Make sure `broadcast_root_unique_id` \"\n",
" \"runs on the root before calling this method.\"\n",
" )\n",
" \n",
" from raft_dask.common.nccl import nccl\n",
" from pylibraft.common.handle import Handle\n",
" from raft_dask.common.comms_utils import inject_comms_on_handle_coll_only\n",
" \n",
" # NCCL + RAFT handle wiring\n",
" self._nccl = nccl()\n",
" self._nccl.init(self.world_size, self.root_unique_id, self.rank)\n",
" \n",
" self._raft_handle = Handle(n_streams=0)\n",
" inject_comms_on_handle_coll_only(\n",
" self._raft_handle, self._nccl, self.world_size, self.rank, verbose=True\n",
" )\n",
" \n",
" # IMPORTANT: pass the RAFT handle for distributed KMeans\n",
" self.kmeans = cuml.KMeans(\n",
" handle=self._raft_handle,\n",
" n_clusters=self.n_clusters,\n",
" init=\"k-means++\",\n",
" max_iter=300,\n",
" tol=1e-4,\n",
" random_state=self.random_state,\n",
" output_type=\"cupy\",\n",
" )\n",
" print(f\"{self.name}: RAFT setup complete\")\n",
" \n",
" def fit(self, data_files: list[str], embedding_col: str):\n",
" \"\"\"Fit distributed KMeans using RAFT\"\"\"\n",
" df = cudf.read_parquet(data_files)\n",
" embeddings = get_array_from_df(df, embedding_col)\n",
" \n",
" self.kmeans._fit(\n",
" embeddings,\n",
" sample_weight=None,\n",
" convert_dtype=False,\n",
" multigpu=True\n",
" )\n",
" # Return numpy labels to avoid GPU object shuffles to driver\n",
" return self.kmeans.predict(embeddings, convert_dtype=False).get()\n",
"\n",
"def raft_multi_gpu_kmeans(data_files: list[str], n_gpus: int = 4, n_clusters: int = N_CLUSTERS, random_state: int = RANDOM_STATE):\n",
" files_per_gpu = np.array_split(data_files, n_gpus)\n",
" \n",
" # Create RAFT actors (one per GPU) - MUST use named actors for broadcast pattern\n",
" actors = [\n",
" RAFTKMeansActor.options(name=f\"RAFTActor-{i}\").remote(\n",
" rank=i,\n",
" world_size=n_gpus,\n",
" n_clusters=n_clusters,\n",
" random_state=random_state,\n",
" actor_name_prefix=\"RAFT\" # Must match the name option\n",
" )\n",
" for i in range(n_gpus)\n",
" ]\n",
" \n",
" print(\"Setting up RAFT communication...\")\n",
" # Get the root actor (index 0) and broadcast root unique ID\n",
" # (matches executor.py lines 194-196)\n",
" root_actor = actors[0]\n",
" ray.get(root_actor.broadcast_root_unique_id.remote())\n",
" \n",
" # Setup all actors (including root) - matches executor.py lines 198-200\n",
" setup_futures = [actor.setup_raft.remote() for actor in actors]\n",
" ray.get(setup_futures)\n",
" print(\"RAFT setup complete!\")\n",
" \n",
" print(\"\\nTraining distributed KMeans across all GPUs...\")\n",
" fit_futures = [\n",
" actors[i].fit.remote(files_per_gpu[i].tolist(), 'embeddings')\n",
" for i in range(n_gpus)\n",
" ]\n",
" return ray.get(fit_futures)"
]
},
{
"cell_type": "code",
"execution_count": 25,
"id": "d4099079-0243-4c78-861b-a6b604fc704c",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Setting up RAFT communication...\n",
"RAFT setup complete!\n",
"\n",
"Training distributed KMeans across all GPUs...\n"
]
}
],
"source": [
"ray_multi_node_centroids = raft_multi_gpu_kmeans(\n",
" list(map(lambda x : str(x.as_posix()), (Path(\"./clustered_data/input/\").glob(\"*.parquet\")))),\n",
" n_gpus=4\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 26,
"id": "2ebd4c2b-7406-486f-bc72-50d92f251d9d",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"1 20000\n",
"2 20000\n",
"5 10000\n",
"3 10000\n",
"8 10000\n",
"4 10000\n",
"9 10000\n",
"0 4293\n",
"6 3577\n",
"7 2130\n",
"Name: count, dtype: int64"
]
},
"execution_count": 26,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"pd.Series(np.concatenate(ray_multi_node_centroids)).value_counts()"
]
},
{
"cell_type": "code",
"execution_count": 27,
"id": "240a8c7d-95d3-419f-bc28-437a4af59c07",
"metadata": {},
"outputs": [
{
"data": {
"text/plain": [
"0.0018910334551606482"
]
},
"execution_count": 27,
"metadata": {},
"output_type": "execute_result"
}
],
"source": [
"adjusted_rand_score(\n",
" pd.Series(np.concatenate(ray_multi_node_centroids)).to_numpy(), \n",
" df['true_cluster'].to_pandas().to_numpy()\n",
")"
]
}
],
"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.11"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment