[Merge only onto large-scale] Large scale Kmeans changes.

configs/spmd/large_scale/kmeans_narrow_weak.json

@@ -0,0 +1,33 @@
+{
+    "INCLUDE": ["../../common/sklearn.json", "../large_scale/large_scale.json"],
+    "PARAMETERS_SETS": {
+        "spmd kmeans parameters": {
+            "algorithm": {
+                "estimator": "KMeans",
+                "estimator_params": {
+                    "algorithm": "lloyd",
+		    "max_iter": 20,
+		    "n_clusters": 10,
+		    "random_state": 42
+                },
+                "estimator_methods": { "training": "fit", "inference": "" },
+                "sklearnex_context": { "use_raw_input": true }
+            }
+        },
+        "synthetic data": {
+                "data": [
+		    { "source": "make_blobs", "generation_kwargs": { "n_samples": 2000000,  "n_features": 100, "centers": 2000, "cluster_std": 3, "center_box": 100.0}}
+                ]
+        }
+    },
+    "TEMPLATES": {
+        "kmeans": {
+            "SETS": [
+                "synthetic data",
+                "sklearnex spmd implementation",
+                "large scale 2k parameters sample shift",
+                "spmd kmeans parameters"
+            ]
+        }
+    }
+}

configs/spmd/large_scale/kmeans_strong.json

@@ -1,20 +1,21 @@
 {
-    "INCLUDE": ["../../common/sklearn.json", "large_scale.json"],
+    "INCLUDE": ["../../common/sklearn.json", "../large_scale/large_scale.json"],
     "PARAMETERS_SETS": {
         "spmd kmeans parameters": {
             "algorithm": {
                 "estimator": "KMeans",
                 "estimator_params": {
-                    "algorithm": "lloyd"
+                    "algorithm": "lloyd",
+		    "max_iter": 20,
+		    "n_clusters": 100
                 },
-                "estimator_methods": { "training": "fit", "inference": "predict" }
+                "estimator_methods": { "training": "fit", "inference": "predict" },
+                "sklearnex_context": { "use_raw_input": true }
             }
-	},
-	"synthetic data": {
+        },
+        "synthetic data": {
                 "data": [
-                        { "source": "make_blobs", "generation_kwargs": { "n_samples": 5000000,  "n_features": 10, "centers": 10 }, "algorithm": { "n_clusters": 10, "max_iter": 10 } },
-	                { "source": "make_blobs", "generation_kwargs": { "n_samples": 30000,  "n_features": 1000, "centers": 10 }, "algorithm": { "n_clusters": 10, "max_iter": 10 } },
-			{ "source": "make_blobs", "generation_kwargs": { "n_samples": 1000000,  "n_features": 100, "centers": 100 }, "algorithm": { "n_clusters": 100, "max_iter": 100 } }
+                        { "source": "make_blobs", "generation_kwargs": { "n_samples": 25000000,  "n_features": 100, "centers": 100 }}
                 ]
         }
     },
@@ -29,3 +30,4 @@
         }
     }
 }
+

configs/spmd/large_scale/kmeans_wide_weak.json

@@ -0,0 +1,34 @@
+{
+    "INCLUDE": ["../../common/sklearn.json", "../large_scale/large_scale.json"],
+    "PARAMETERS_SETS": {
+        "spmd kmeans parameters": {
+            "algorithm": {
+                "estimator": "KMeans",
+                "estimator_params": {
+                    "algorithm": "lloyd",
+		    "max_iter": 20,
+		    "n_clusters": 10,
+		    "random_state": 42
+                },
+                "estimator_methods": { "training": "fit", "inference": "" },
+                "sklearnex_context": { "use_raw_input": true }
+            }
+        },
+        "synthetic data": {
+                "data": [
+		    { "source": "make_blobs", "generation_kwargs": { "n_samples": 1000000,  "n_features": 1000, "centers": 2000}}
+                ]
+        }
+    },
+    "TEMPLATES": {
+        "kmeans": {
+            "SETS": [
+                "synthetic data",
+                "sklearnex spmd implementation",
+                "large scale 2k parameters",
+                "spmd kmeans parameters"
+            ]
+        }
+    }
+}
+

configs/spmd/large_scale/large_scale.json

@@ -27,6 +27,15 @@
                 "mpi_params": {"n": [1,2,6,12,24,48,96,192,384,768,1536,3072,6144,12288,24576], "ppn": 12, "-hostfile": "", "-cpu-bind=list:0-7,104-111:8-15,112-119:16-23,120-127:24-31,128-135:32-39,136-143:40-47,144-151:52-59,156-163:60-67,164-171:68-75,172-179:76-83,180-187:84-91,188-195:92-99,196-203": "--envall gpu_tile_compact.sh" }
             }
         },
+        "large scale 2k parameters sample shift": {
+            "data": {
+                "dtype": "float64",
+                "distributed_split": "sample_shift"
+            },
+            "bench": {
+                "mpi_params": {"n": [1,2,6,12], "ppn": 12, "-hostfile": "", "-cpu-bind=list:0-7,104-111:8-15,112-119:16-23,120-127:24-31,128-135:32-39,136-143:40-47,144-151:52-59,156-163:60-67,164-171:68-75,172-179:76-83,180-187:84-91,188-195:92-99,196-203": "--envall gpu_tile_compact.sh" }
+            }
+        },
         "large scale 32 parameters": {
             "data": {
                 "dtype": "float64",

sklbench/benchmarks/sklearn_estimator.py

@@ -191,19 +191,6 @@ def get_subset_metrics_of_estimator(
                     }
                 )
     elif task == "clustering":
-        if hasattr(estimator_instance, "inertia_"):
-            # compute inertia manually using distances to cluster centers
-            # provided by KMeans.transform
-            metrics.update(
-                {
-                    "inertia": float(
-                        np.power(
-                            convert_to_numpy(estimator_instance.transform(x)).min(axis=1),
-                            2,
-                        ).sum()
-                    )
-                }
-            )
         if hasattr(estimator_instance, "predict"):
             y_pred = convert_to_numpy(estimator_instance.predict(x))
             metrics.update(

sklbench/datasets/__init__.py

@@ -67,6 +67,11 @@ def load_data(bench_case: BenchCase) -> Tuple[Dict, Dict]:
             generation_kwargs = get_bench_case_value(
                 bench_case, "data:generation_kwargs", dict()
             )
+            if "center_box" in generation_kwargs:
+                generation_kwargs["center_box"] = (
+                    -1 * generation_kwargs["center_box"],
+                    generation_kwargs["center_box"],
+                )
             return load_sklearn_synthetic_data(
                 function_name=source,
                 input_kwargs=generation_kwargs,

sklbench/datasets/transformer.py

@@ -16,8 +16,10 @@
 
 import os
 
+import math
 import numpy as np
 import pandas as pd
+from mpi4py import MPI
 from scipy.sparse import csr_matrix
 from sklearn.model_selection import train_test_split
 
@@ -113,8 +115,38 @@ def split_and_transform_data(bench_case, data, data_description):
         "KNeighbors" in get_bench_case_value(bench_case, "algorithm:estimator", "")
         and int(get_bench_case_value(bench_case, "bench:mpi_params:n", 1)) > 1
     )
-    if distributed_split == "rank_based" or knn_split_train:
-        from mpi4py import MPI
+
+    if distributed_split == "sample_shift":
+        comm = MPI.COMM_WORLD
+        rank = comm.Get_rank()
+        size = comm.Get_size()
+
+        n_train = len(x_train)
+        n_test = len(x_test)
+
+        train_start = 0
+        train_end = n_train
+        test_start = 0
+        test_end = n_test
+
+        adjust_number = (math.sqrt(rank) * 0.003) + 1
+
+        if "y" in data:
+            x_train, y_train = (
+                x_train[train_start:train_end] * adjust_number,
+                y_train[train_start:train_end],
+            )
+
+            x_test, y_test = (
+                x_test[test_start:test_end] * adjust_number,
+                y_test[test_start:test_end],
+            )
+        else:
+            x_train = x_train[train_start:train_end]
+
+            x_test = x_test[test_start:test_end] * adjust_number
+
+    elif distributed_split == "rank_based" or knn_split_train:
 
         comm = MPI.COMM_WORLD
         rank = comm.Get_rank()
@@ -127,6 +159,7 @@ def split_and_transform_data(bench_case, data, data_description):
         train_end = (1 + rank) * n_train // size
         test_start = rank * n_test // size
         test_end = (1 + rank) * n_test // size
+        x_train_rank = x_train[train_start:train_end]
 
         if "y" in data:
             x_train, y_train = (
@@ -138,7 +171,7 @@ def split_and_transform_data(bench_case, data, data_description):
         else:
             x_train = x_train[train_start:train_end]
             if distributed_split == "rank_based":
-                x_test = x_test[test_start:test_end]
+                x_test = x_test[test_start:test_end] * adjust_number
 
     device = get_bench_case_value(bench_case, "algorithm:device", None)
     common_data_format = get_bench_case_value(bench_case, "data:format", "pandas")

sklbench/utils/common.py

@@ -120,7 +120,7 @@ def flatten_list(input_list: List, ensure_type_homogeneity: bool = False) -> Lis
 
 
 def get_module_members(
-    module_names_chain: Union[List, str]
+    module_names_chain: Union[List, str],
 ) -> Tuple[ModuleContentMap, ModuleContentMap]:
     def get_module_name(module_names_chain: List[str]) -> str:
         name = module_names_chain[0]