[LightGBM] Classifier hangs on fit most of the time

### SynapseML version

0.11.1

### System information

- **Language version** (e.g. python 3.8, scala 2.12): 3.9.13 | packaged by conda-forge
- **Spark Version** (e.g. 3.2.3): 3.3.2
- **Spark Platform** (e.g. Synapse, Databricks): Self hosted on Kubernetes

### Describe the problem

We have been attempting to fit a dataset of about 18 million rows and 225 columns with the synapseML LightGBM classifier. This is a binary classification problem. Each feature is a continuous float value.

When we running our hyperparameter tuning (handled by hyperopt), the model occasionally hangs on the fitting stage. By occasionally, we mean most of the time - but sometimes, if we are “lucky”, there is a chance that it will fully train the model and sometimes even run through a parameter search fully. Specifically, it hangs on the `collect` function on line 599 in this code:
https://github.com/microsoft/SynapseML/blob/14213e25147b99611ee56946933208f8d7520560/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/LightGBMBase.scala

We identified this by going into the spark UI and looking at the specific operations it was running.

We have tried using barrier execution mode to no avail, as well as setting a variety of variables. I shared the variety of parameters we have tested for the classifier near the end.

We are running spark through pyspark on a jupyter notebook, and we installed synapseML via the spark packages. Our spark master/workers are hosted on a kubernetes cluster. The data is stored in parquet format on distributed storage.

Note we are also running spark with spot instances; we scale nodes up and down as we need the resources. We have attempted fixing the nodes to just a specific amount of workers however we have not attempted entirely disabling spot instances.

### Code to reproduce issue

Below is how we set up our spark session, only including what may be relevant to the issue:

```

packages = [
'com.microsoft.azure:synapseml_2.12:0.11.1',
]

spark = SparkSession.builder \
.appName("ModelTraining") \
.master(SPARK_MASTER) \
.config("spark.kubernetes.namespace", "spark") \
.config("spark.kubernetes.container.image", "apache/spark-py:v3.4.0") \
.config("spark.kubernetes.authenticate.driver.serviceAccountName", "spark") \
.config("spark.executor.memory", "32g") \
.config("spark.executor.cores", "12") \
.config("spark.driver.memory", "15g") \
.config("spark.driver.maxResultSize", "0") \
.config("spark.driver.bindAddress", "0.0.0.0") \
.config("spark.driver.host", driver_host) \
.config("spark.driver.port", "31137") \
.config("spark.blockManage r.port", "7777") \
.config("spark.storage.replication", "1") \
.config("spark.jars.packages", ",".join(packages)) \
.config("spark.jars.repositories", "https://mmlspark.azureedge.net/maven") \
.config("spark.sql.mapKeyDedupPolicy", "LAST_WIN") \
.config("spark.executor.extraJavaOptions", "-XX:+UseG1GC") \
.config("spark.driver.extraJavaOptions", "-Xss4M -XX:+UseG1GC") \
.getOrCreate()
```

This is how we are preparing our data:

```
# Repartition data
feature_df = feature_df.repartition(72)

# Split the tagged_feature into a train, test split
train, val, test = feature_df.randomSplit([0.8, .1, 0.1], seed=1)

# Find numeric columns for feature assembler
# Uses a helper function from our preprocessing code that simply gets the feature columns
feature_generation = FeatureGeneration()
feat_columns = feature_generation.generate_feat(feature_df)

# delete unused feature df
feature_df.unpersist()

# pull out target feature columns
featurizer = VectorAssembler(inputCols=feat_columns, outputCol="features")

# Create the train and test data based on the features and split data
train_data = featurizer.transform(train)["label", "features"]
val_data = featurizer.transform(val)["label", "features"]
test_data = featurizer.transform(test)["label", "features"]

# force DAG
train_data.cache()
train_data.head()
val_data.cache()
val_data.head()
test_data.cache()
test_data.head()
```

And lastly, tuning and fitting the model. We have tried a wide combination of the settings below. Note this is not exactly how we run this since we actually pass this parameter space to hyperopt, but we left that out for simplicity.

```
num_eval = 10
boostingTypes = ['gbdt']
objectives = ['binary']
binSampleCounts = [10]

param_hyperopt= {
'featuresCol' : hp.choice('featuresCol', ["features"]),
'labelCol' : hp.choice('labelCol', ["label"]),
'isUnbalance' : hp.choice('isUnbalance', [True]),
'executionMode' : hp.choice('executionMode', ['streaming']),
'useBarrierExecutionMode' : hp.choice('useBarrierExecutionMode', [False]),
'boostingType': hp.choice('boostingType', boostingTypes),
'objective': hp.choice('objective', objectives),
'learningRate': hp.loguniform('learningRate', np.log(0.01), np.log(1)),
'maxDepth': scope.int(hp.quniform('maxDepth', 5, 10, 1)),
'numIterations': scope.int(hp.quniform('numIterations', 100, 250, 1)),
'numLeaves': scope.int(hp.quniform('numLeaves', 5, 50, 1)),
# 'numThreads' : hp.choice('numThreads', [70]),
# 'colSampleByTree': hp.uniform('colsample_bytree', 0.6, 1.0),
# 'regLambda': hp.uniform('reg_lambda', 0.0, 1.0),
# 'binSampleCount' : hp.choice('binSampleCount', binSampleCounts),
#'useSingleDatasetMode' : hp.choice('useSingleDatasetMode', [True]),
}

model = LightGBMClassifier(**param_hyperopt)
model = model.fit(train_data) # <- This will hang on the collect of line 599
```

### Other info / logs

If we use the default executor hearbeat interval we will hit that and receive an missed heartbeat interval error. When we change it to a very large interval / network timeout we see a different error, fairly typical spark crash with not much info. The executors do not show errors in their logs.

```
[Stage 22:================> (7 + 17) / 24]

23/05/31 17:12:31 ERROR TaskSchedulerImpl: Lost executor 0 on 10.244.3.14: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 10.0 in stage 22.0 (TID 433) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 13.0 in stage 22.0 (TID 436) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 5.0 in stage 22.0 (TID 427) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 2.0 in stage 22.0 (TID 421) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 15.0 in stage 22.0 (TID 438) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 7.0 in stage 22.0 (TID 429) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 3.0 in stage 22.0 (TID 423) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 11.0 in stage 22.0 (TID 435) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 14.0 in stage 22.0 (TID 437) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 0.0 in stage 22.0 (TID 419) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 9.0 in stage 22.0 (TID 431) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN TaskSetManager: Lost task 4.0 in stage 22.0 (TID 425) (10.244.3.14 executor 0): ExecutorLostFailure (executor 0 exited caused by one of the running tasks) Reason: Remote RPC client disassociated. Likely due to containers exceeding thresholds, or network issues. Check driver logs for WARN messages.
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_49 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_29 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_68 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_21 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_65 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_2 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_3 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_19 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_26 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_20 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_55 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_27 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_70 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_66 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_32 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_59 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_36 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_30 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_64 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_6 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_31 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_24 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_56 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_67 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_61 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_62 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_57 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_46 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_0 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_33 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_58 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_62_0 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_63 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_41 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_15 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_34 !
23/05/31 17:12:31 WARN BlockManagerMasterEndpoint: No more replicas available for rdd_32_28 !
23/05/31 17:12:33 WARN TaskSetManager: Lost task 0.1 in stage 22.0 (TID 445) (10.244.19.3 executor 2): FetchFailed(null, shuffleId=2, mapIndex=-1, mapId=-1, reduceId=0, message=
org.apache.spark.shuffle.MetadataFetchFailedException: Missing an output location for shuffle 2 partition 0
at org.apache.spark.MapOutputTracker$.validateStatus(MapOutputTracker.scala:1705)
at org.apache.spark.MapOutputTracker$.$anonfun$convertMapStatuses$10(MapOutputTracker.scala:1652)
at org.apache.spark.MapOutputTracker$.$anonfun$convertMapStatuses$10$adapted(MapOutputTracker.scala:1651)
at scala.collection.Iterator.foreach(Iterator.scala:943)
at scala.collection.Iterator.foreach$(Iterator.scala:943)
at scala.collection.AbstractIterator.foreach(Iterator.scala:1431)
at org.apache.spark.MapOutputTracker$.convertMapStatuses(MapOutputTracker.scala:1651)
at org.apache.spark.MapOutputTrackerWorker.getMapSizesByExecutorIdImpl(MapOutputTracker.scala:1294)
at org.apache.spark.MapOutputTrackerWorker.getMapSizesByExecutorId(MapOutputTracker.scala:1256)
at org.apache.spark.shuffle.sort.SortShuffleManager.getReader(SortShuffleManager.scala:140)
at org.apache.spark.shuffle.ShuffleManager.getReader(ShuffleManager.scala:63)
at org.apache.spark.shuffle.ShuffleManager.getReader$(ShuffleManager.scala:57)
at org.apache.spark.shuffle.sort.SortShuffleManager.getReader(SortShuffleManager.scala:73)
at org.apache.spark.sql.execution.ShuffledRowRDD.compute(ShuffledRowRDD.scala:208)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.$anonfun$getOrCompute$1(RDD.scala:378)
at org.apache.spark.storage.BlockManager.$anonfun$doPutIterator$1(BlockManager.scala:1523)
at org.apache.spark.storage.BlockManager.org$apache$spark$storage$BlockManager$$doPut(BlockManager.scala:1450)
at org.apache.spark.storage.BlockManager.doPutIterator(BlockManager.scala:1514)
at org.apache.spark.storage.BlockManager.getOrElseUpdate(BlockManager.scala:1337)
at org.apache.spark.rdd.RDD.getOrCompute(RDD.scala:376)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:327)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.CoalescedRDD.$anonfun$compute$1(CoalescedRDD.scala:99)
at scala.collection.Iterator$$anon$11.nextCur(Iterator.scala:486)
at scala.collection.Iterator$$anon$11.hasNext(Iterator.scala:492)
at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:460)
at scala.collection.Iterator$$anon$10.hasNext(Iterator.scala:460)
at com.microsoft.azure.synapse.ml.lightgbm.BasePartitionTask.initialize(BasePartitionTask.scala:177)
at com.microsoft.azure.synapse.ml.lightgbm.BasePartitionTask.mapPartitionTask(BasePartitionTask.scala:132)
at com.microsoft.azure.synapse.ml.lightgbm.LightGBMBase.$anonfun$executePartitionTasks$1(LightGBMBase.scala:591)
at org.apache.spark.sql.execution.MapPartitionsExec.$anonfun$doExecute$3(objects.scala:201)
at org.apache.spark.rdd.RDD.$anonfun$mapPartitionsInternal$2(RDD.scala:890)
at org.apache.spark.rdd.RDD.$anonfun$mapPartitionsInternal$2$adapted(RDD.scala:890)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.rdd.MapPartitionsRDD.compute(MapPartitionsRDD.scala:52)
at org.apache.spark.rdd.RDD.computeOrReadCheckpoint(RDD.scala:365)
at org.apache.spark.rdd.RDD.iterator(RDD.scala:329)
at org.apache.spark.scheduler.ResultTask.runTask(ResultTask.scala:90)
at org.apache.spark.scheduler.Task.run(Task.scala:136)
at org.apache.spark.executor.Executor$TaskRunner.$anonfun$run$3(Executor.scala:548)
at org.apache.spark.util.Utils$.tryWithSafeFinally(Utils.scala:1504)
at org.apache.spark.executor.Executor$TaskRunner.run(Executor.scala:551)
at java.base/java.util.concurrent.ThreadPoolExecutor.runWorker(ThreadPoolExecutor.java:1128)
at java.base/java.util.concurrent.ThreadPoolExecutor$Worker.run(ThreadPoolExecutor.java:628)
at java.base/java.lang.Thread.run(Thread.java:829)
```

### What component(s) does this bug affect?

- [ ] `area/cognitive`: Cognitive project
- [ ] `area/core`: Core project
- [ ] `area/deep-learning`: DeepLearning project
- [X] `area/lightgbm`: Lightgbm project
- [ ] `area/opencv`: Opencv project
- [ ] `area/vw`: VW project
- [ ] `area/website`: Website
- [ ] `area/build`: Project build system
- [ ] `area/notebooks`: Samples under notebooks folder
- [ ] `area/docker`: Docker usage
- [ ] `area/models`: models related issue

### What language(s) does this bug affect?

- [ ] `language/scala`: Scala source code
- [X] `language/python`: Pyspark APIs
- [ ] `language/r`: R APIs
- [ ] `language/csharp`: .NET APIs
- [ ] `language/new`: Proposals for new client languages

### What integration(s) does this bug affect?

- [ ] `integrations/synapse`: Azure Synapse integrations
- [ ] `integrations/azureml`: Azure ML integrations
- [ ] `integrations/databricks`: Databricks integrations