From 0df0110857334c4443be71b174a82bcbc2e8ca0b Mon Sep 17 00:00:00 2001 From: Rana Singh Date: Fri, 4 Sep 2026 08:52:30 -0700 Subject: [PATCH] fix: propagate LightGBM iteration failures ## Summary Stop treating native LightGBM iteration exceptions as successful early completion. Log and rethrow the original failure so Spark fails the task instead of constructing a partial model, with a deterministic regression test that preserves the exception and unfinished state. ## Prompting Intent Investigate the value of closed PR #2684 and the linked distributed LightGBM issue family, then implement evidence-backed actionable work in an isolated worktree and validate it locally and on Microsoft Fabric. ## Linked Sources - Closed contribution: https://github.com/microsoft/SynapseML/pull/2684 - Historical multiclass report: https://github.com/microsoft/SynapseML/issues/569 - Historical binary socket report: https://github.com/microsoft/SynapseML/issues/728 - Retry diagnostics fix: https://github.com/microsoft/SynapseML/pull/2612 ## Rationale Normal LightGBM completion is already returned by the native update call. An exception represents failure, not early stopping. Rethrowing the same object preserves the native cause and lets Spark apply task and job failure semantics; swallowing it can return an incomplete model that appears valid. This deliberately does not claim to repair every network or executor failure in the broader issue family. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../synapse/ml/lightgbm/TrainUtils.scala | 6 +- .../ml/lightgbm/split1/TrainUtilsSuite.scala | 58 ++++++++++++++++++- 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/TrainUtils.scala b/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/TrainUtils.scala index dc2e631da46..087af26ba5e 100644 --- a/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/TrainUtils.scala +++ b/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 } } diff --git a/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/TrainUtilsSuite.scala b/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/TrainUtilsSuite.scala index c89f6140af3..208d482ddc2 100644 --- a/lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/TrainUtilsSuite.scala +++ b/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()