fix: propagate LightGBM iteration failures

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/TrainUtils.scala

@@ -107,10 +107,8 @@ private[lightgbm] object TrainUtils extends Serializable {
       }
     } catch {
       case e: java.lang.Exception =>
-        log.warn("LightGBM reached early termination on one task," +
-          " stopping training on task. This message should rarely occur." +
-          " Inner exception: " + e.toString)
-        state.isFinished = true
+        log.error(s"LightGBM task failed during iteration ${state.iteration}", e)
+        throw e
     }
   }
 

lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/TrainUtilsSuite.scala

@@ -3,7 +3,13 @@
 
 package com.microsoft.azure.synapse.ml.lightgbm.split1
 
-import com.microsoft.azure.synapse.ml.lightgbm.{LightGBMRegressor, NetworkManager, TrainUtils}
+import com.microsoft.azure.synapse.ml.io.http.SharedSingleton
+import com.microsoft.azure.synapse.ml.lightgbm.booster.LightGBMBooster
+import com.microsoft.azure.synapse.ml.lightgbm.{ColumnParams, LightGBMRegressor, NetworkManager, NetworkParams,
+  NetworkTopologyInfo, PartitionTaskContext, PartitionTaskTrainingState, SharedState, TaskInstrumentationMeasures,
+  TrainingContext, TrainUtils}
+import org.apache.spark.ml.linalg.SQLDataTypes
+import org.apache.spark.sql.types.{StructField, StructType}
 import org.scalatest.funsuite.AnyFunSuite
 import org.slf4j.LoggerFactory
 
@@ -43,6 +49,41 @@ class TrainUtilsSuite extends AnyFunSuite {
 
   private val tolerances = Seq(0.0, 0.25, 2.0, 25.0)
 
+  private def newTrainingState(booster: LightGBMBooster): PartitionTaskTrainingState = {
+    val featuresField = StructField("features", SQLDataTypes.VectorType)
+    val trainParams = new LightGBMRegressor().getTrainParams(
+      numTasks = 1,
+      featuresSchema = featuresField,
+      numTasksPerExec = 1)
+    val trainingContext = TrainingContext(
+      batchIndex = 0,
+      sharedStateSingleton = SharedSingleton(new SharedState(trainParams)),
+      schema = StructType(Seq(featuresField)),
+      numCols = 1,
+      numInitScoreClasses = 0,
+      trainingParams = trainParams,
+      networkParams = NetworkParams(12400, "127.0.0.1", 12400, barrierExecutionMode = false),
+      columnParams = ColumnParams("label", "features", None, None, None),
+      datasetParams = "",
+      featureNames = None,
+      numTasksPerExecutor = 1,
+      validationData = None,
+      serializedReferenceDataset = None,
+      partitionCounts = Some(Array(1L)))
+    val taskContext = PartitionTaskContext(
+      trainingCtx = trainingContext,
+      partitionId = 0,
+      taskId = 0L,
+      measures = new TaskInstrumentationMeasures(0),
+      networkTopologyInfo = NetworkTopologyInfo("127.0.0.1:12400", Array(0), 12400),
+      shouldExecuteTraining = true,
+      isEmptyPartition = false,
+      shouldReturnBooster = true,
+      shouldCalcValidationDataset = false)
+
+    PartitionTaskTrainingState(taskContext, booster)
+  }
+
   test("Improvement tolerance is symmetric across metrics and tolerance values") {
     val bestScore = 100.0
     val margin = 0.125
@@ -146,6 +187,21 @@ class TrainUtilsSuite extends AnyFunSuite {
     assert(TrainUtils.shouldStopEarly(iteration = 10, bestIteration = 5, earlyStoppingRound = 5))
   }
 
+  test("A native iteration failure is not reported as completed training") {
+    val nativeFailure = new RuntimeException("injected native iteration failure")
+    val booster = new LightGBMBooster() {
+      override def updateOneIteration(): Boolean = throw nativeFailure
+    }
+    val state = newTrainingState(booster)
+
+    val thrown = intercept[RuntimeException] {
+      TrainUtils.updateOneIteration(state, log)
+    }
+
+    assert(thrown eq nativeFailure)
+    assert(!state.isFinished)
+  }
+
   test("Early stopping parameters accept valid values and reject invalid values") {
     val learner = new LightGBMRegressor()