Fix problem with empty partition assigned to validation data (#2059)

microsoft · Aug 31, 2023 · 0836e40 · 0836e40
1 parent c6d5882
commit 0836e40
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/BasePartitionTask.scala b/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/BasePartitionTask.scala
@@ -324,6 +324,9 @@ abstract class BasePartitionTask extends Serializable with Logging {
         s" shouldExecuteTraining: $shouldExecuteTraining, isEmptyPartition: $isEmptyPartition")
 
     val shouldCalcValidationDataset = trainingCtx.sharedState.validationDatasetWorker.getOrElse(-1) == taskId
+    if (trainingCtx.hasValidationData)
+      log.info(s"Validation data found. Task: $taskId, PartId: $partitionId. Main task: $mainExecutorWorkerId" +
+        s" shouldCalcValidationDataset: $shouldCalcValidationDataset, isEmptyPartition: $isEmptyPartition")
 
     PartitionTaskContext(trainingCtx,
                          partitionId,

diff --git a/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/StreamingPartitionTask.scala b/lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/StreamingPartitionTask.scala
@@ -106,7 +106,7 @@ class StreamingPartitionTask extends BasePartitionTask {
     if (!shouldExecuteTraining && !isEmptyPartition) ctx.sharedState().incrementDataPrepDoneSignal(log)
 
     // First dataset to reach here calculates the validation Dataset if needed
-    if (ctx.hasValidationData) {
+    if (ctx.hasValidationData && !isEmptyPartition) {
       ctx.sharedState().linkValidationDatasetWorker()
     }
   }