[BUG] LightGBM is failing on a cluster with 2 executors or more #2674

from synapse.ml.lightgbm import LightGBMRegressor
from pyspark.ml.feature import StringIndexer, VectorAssembler

# Generate synthetic data with 100,000 rows and 6 categorical columns
import numpy as np
import pandas as pd

num_rows = 100000
cat7_values = [i for i in range(200)]  # 200 distinct values for cat7

df_pd = pd.DataFrame({
    "cat1": np.random.choice(['A', 'B', 'C'], num_rows),
    "cat2": np.random.choice(['X', 'Y', 'Z'], num_rows),
    "cat3": np.random.choice(['foo', 'bar', 'baz'], num_rows),
    "cat4": np.random.choice(['red', 'green', 'blue'], num_rows),
    "cat5": np.random.choice(['dog', 'cat', 'mouse'], num_rows),
    "cat6": np.random.choice(['apple', 'banana', 'cherry'], num_rows),
    "cat7": np.random.choice(cat7_values, num_rows),  # new high-cardinality column
    "num1": np.random.randn(num_rows),
    "num2": np.random.randn(num_rows),
    "num3": np.random.randn(num_rows),
    "num4": np.random.randn(num_rows),
    "label": np.random.randn(num_rows)
})

df = spark.createDataFrame(df_pd)

# Index categorical columns
indexers = [
    StringIndexer(inputCol=col, outputCol=f"{col}_idx")
    for col in ["cat1", "cat2", "cat3", "cat4", "cat5", "cat6", "cat7"]
]

# Assemble features
assembler = VectorAssembler(
    inputCols=["cat1_idx", "cat2_idx", "cat3_idx", "cat4_idx", "cat5_idx", "cat6_idx", "cat7_idx", "num1", "num2", "num3", "num4"],
    outputCol="features"
)

# LightGBMRegressor
lgbm = LightGBMRegressor(
    featuresCol="features",
    labelCol="label",
    categoricalSlotIndexes=[0, 1, 2, 3, 4, 5, 6],
    useBarrierExecutionMode=True,
    # numTasks=2
)

from pyspark.ml import Pipeline

pipeline = Pipeline(stages=indexers + [assembler, lgbm])
# df = df.repartition(2)
model = pipeline.fit(df)