You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
Without barrier execution mode, the LightGBM driver waits for exactly numTasks training tasks to report before training starts. Two setups left it waiting for the full timeout (20 minutes by default). After that, training failed with a "Connection refused" error that pointed at networking instead of the real problem.
numTasks larger than the input partition count. The non-barrier path used coalesce(numTasks), which can merge partitions but never add them. With 1 input partition and numTasks=2, only 1 task ran while the driver waited for 2. The input is now repartitioned up to numTasks. LightGBMRanker repartitions by its group column, so query groups stay whole even with repartitionByGroupingColumn=false.
numTasks larger than the tasks Spark can run at once. Tasks that started wait for tasks that can't start until they finish. We can't fix that automatically, so the driver now fails with a clear error instead:
The LightGBM driver received network reports from 2 of 3 training tasks, then stopped waiting because no task connected for 10 seconds. Missing partitions: 2. Without barrier execution mode, all numTasks tasks must run at the same time. Check that numTasks is no larger than the number of tasks Spark can run at once ...
The original job failure is attached as a suppressed exception, and the executor-side "could not reach the driver" message now mentions this cause as well.
Addresses part of #2699. Follows #2744, which fixed the same hang for LightGBMRanker under AQE.
What doesn't change
Automatic numTasks (not set): no change. It's already capped at the input partition count, so the new check doesn't run. Reading the partition count of an AQE plan can run upstream shuffle stages early, so the check only runs when numTasks is set explicitly.
Explicit numTasks at or below the partition count: still a coalesce, as before.
Barrier execution mode: unchanged.
Ranker non-barrier grouping path: same single repartition(numTasks, groupCol) as fix: preserve LightGBM ranker grouping under AQE #2744. It no longer reads the input partition count, which it never used in this mode.
No public signature or parameter changes. The new LightGBMMissingTasksException class has a package-private constructor.
Cost
When numTasks is set explicitly, fit now reads the input partition count to decide between coalesce and repartition. Automatic numTasks and barrier mode already do this. For an uncached DataFrame with a shuffle that hasn't run yet under AQE, reading the count runs that shuffle one extra time. On Fabric, an explicit-numTasks fit on an uncached 20M-row aggregate took 10.2s (median) versus 7.9s on master. Automatic numTasks took about 9.5s on both.
This can't be skipped safely: AQE can merge the input below numTasks, and that's the case that hangs. Caching the input avoids the extra pass, and training already reads the input several times. The docs now say this, and the expansion itself is logged at warn level because it adds a shuffle.
Tests
New LightGBMNumTasksSuite (split1) and two new DriverSocketRetrySuite tests:
lightgbm/scalastyle and lightgbm/Test/scalastyle are clean.
Spark 4.0 and 4.1 ports: the commit applied on top of each port branch (only LightGBMRanker.scala conflicted, because fix: preserve LightGBM ranker grouping under AQE #2744 hasn't reached the ports yet). Style is clean, and LightGBMNumTasksSuite and DriverSocketRetrySuite pass (21/21 each, JDK 17, 2 CPUs).
Flakiness: LightGBMNumTasksSuite passed 25 of 25 repeated runs and DriverSocketRetrySuite 8 of 8, on 2 CPUs. After review, the missing-partition test now starts its short 500ms accept timeout only after the first fake worker report is recorded, so a slow agent cannot time out before that report arrives; the suite then passed 5 more runs in a row.
Fabric end-to-end
Ran on Microsoft Fabric (Spark 3.5.5 runtime, 1 executor with 8 task slots, AQE on, dynamic allocation off). Each run loaded this PR's jars (...-227-e51ff84c-SNAPSHOT) and, for the baseline, master's jars (...-226-9d51ad1a-SNAPSHOT). The scenario checked that the classes came from those jars. Non-barrier mode, bulk data transfer:
Case
master
this PR
Classifier, 1 input partition, numTasks=4
fails: "could not reach the driver"
passes, accuracy 1.0
Ranker, 1 input partition, numTasks=4, grouping off
fails with "received network reports from 8 of 9 training tasks ... Missing partitions: 8"
Stock lightgbm-streaming scenario
passes
Streaming data transfer without slotNames fails on this Fabric runtime for both master and this PR, with a native "Feature (Column_) appears more than one time" error during dataset creation. That's a separate, existing problem (see #2242), so the cases above use bulk mode.
Docs
Added a troubleshooting section to the LightGBM overview: "When numTasks is larger than the tasks Spark can run at once". It includes the cost note above.
Note
Open PR #2745 also edits NetworkManager.scala (worker identity diagnostics). The changes are in different methods, but whichever merges second may need a small rebase.
…timeouts
Without barrier execution mode, the LightGBM driver waits for exactly numTasks
training tasks to report. Two cases left it waiting until the timeout, after
which tasks failed with a misleading "Connection refused" error:
- An explicit numTasks larger than the input partition count. coalesce can
only merge partitions, so fewer tasks ran than the driver expected. The
input is now repartitioned up to numTasks. The ranker repartitions by its
grouping column so query groups stay whole. The partition count is only
read for an explicit numTasks, so automatic sizing adds no extra work.
- numTasks larger than the tasks Spark can run at once. This can't be fixed
automatically, so the driver now fails with LightGBMMissingTasksException,
which lists how many tasks reported and which partitions are missing. The
job failure is attached as a suppressed exception.
Adds regression tests that fail on master, and a troubleshooting note in the
LightGBM overview. Addresses part of microsoft#2699.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Hey Rana Singh (@ranadeepsingh) 👋!
Thank you so much for contributing to our repository 🙌.
Someone from SynapseML Team will be reviewing this pull request soon.
We use semantic commit messages to streamline the release process.
Before your pull request can be merged, you should make sure your first commit and PR title start with a semantic prefix.
This helps us to create release messages and credit you for your hard work!
Examples of commit messages with semantic prefixes:
fix: Fix LightGBM crashes with empty partitions
feat: Make HTTP on Spark back-offs configurable
docs: Update Spark Serving usage
build: Add codecov support
perf: improve LightGBM memory usage
refactor: make python code generation rely on classes
style: Remove nulls from CNTKModel
test: Add test coverage for CNTKModel
To test your commit locally, please follow our guild on building from source.
Check out the developer guide for additional guidance on testing your change.
…extra pass
Log the numTasks expansion at warn level, because it adds a shuffle the
user did not ask for, and say how to avoid it. Document that an explicit
numTasks now reads the input partition count, like automatic numTasks and
barrier mode already do, which runs a pending adaptive shuffle once more
on uncached input.
Fabric E2E (runtime Spark 3.5.5, 1 executor x 8 slots, AQE on) measured
explicit numTasks fits at 10.2s median versus 7.9s on master for an
uncached 20M-row aggregate, matching one extra upstream pass.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
A first-attempt connection refusal can also follow an earlier task failure: executeTraining then closes the driver's listener before this task connects. In that case there is no missing-tasks error to find. Keep the earlier-failure case in this diagnostic and direct users to the first failed task or driver error when the missing-tasks explanation is absent.
The missing-partition test set a 500ms accept timeout before the fake
worker connected, so a slow agent could time out with zero reports and
fail the assertion on partition 1. The driver accepts one connection at a
time, so the helper now applies the short timeout on the second accept,
after the first report is recorded, and the test waits for that point.
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
The partition-count cost does not apply to the default non-barrier ranker grouping path: LightGBMRanker.prepareDataframe directly calls repartition(numTasks, groupCol) and skips the count. This note says every explicit-numTasks fit reads the input partition count, which could mislead users about the extra input pass. Qualify it as applying to classifiers, regressors, and rankers with grouping repartition disabled; explain the grouped ranker exception.
Closing in favor of #2745. These changes are now on that PR, on top of Lenin Mookiah (@leninworld)'s attempt identity work, so the #2699 fixes land together and the original contributor keeps credit. The review fixes from this PR carried over, and #2745 also has a newer Fabric end-to-end run and Azure results.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What this fixes
Without barrier execution mode, the LightGBM driver waits for exactly
numTaskstraining tasks to report before training starts. Two setups left it waiting for the fulltimeout(20 minutes by default). After that, training failed with a "Connection refused" error that pointed at networking instead of the real problem.numTaskslarger than the input partition count. The non-barrier path usedcoalesce(numTasks), which can merge partitions but never add them. With 1 input partition andnumTasks=2, only 1 task ran while the driver waited for 2. The input is now repartitioned up tonumTasks.LightGBMRankerrepartitions by its group column, so query groups stay whole even withrepartitionByGroupingColumn=false.numTaskslarger than the tasks Spark can run at once. Tasks that started wait for tasks that can't start until they finish. We can't fix that automatically, so the driver now fails with a clear error instead:The original job failure is attached as a suppressed exception, and the executor-side "could not reach the driver" message now mentions this cause as well.
Addresses part of #2699. Follows #2744, which fixed the same hang for
LightGBMRankerunder AQE.What doesn't change
numTasks(not set): no change. It's already capped at the input partition count, so the new check doesn't run. Reading the partition count of an AQE plan can run upstream shuffle stages early, so the check only runs whennumTasksis set explicitly.numTasksat or below the partition count: still acoalesce, as before.repartition(numTasks, groupCol)as fix: preserve LightGBM ranker grouping under AQE #2744. It no longer reads the input partition count, which it never used in this mode.LightGBMMissingTasksExceptionclass has a package-private constructor.Cost
When
numTasksis set explicitly, fit now reads the input partition count to decide betweencoalesceandrepartition. AutomaticnumTasksand barrier mode already do this. For an uncached DataFrame with a shuffle that hasn't run yet under AQE, reading the count runs that shuffle one extra time. On Fabric, an explicit-numTasksfit on an uncached 20M-row aggregate took 10.2s (median) versus 7.9s on master. AutomaticnumTaskstook about 9.5s on both.This can't be skipped safely: AQE can merge the input below
numTasks, and that's the case that hangs. Caching the input avoids the extra pass, and training already reads the input several times. The docs now say this, and the expansion itself is logged at warn level because it adds a shuffle.Tests
New
LightGBMNumTasksSuite(split1) and two newDriverSocketRetrySuitetests:numTasksabove input partitions givesnumTaskspartitionsnumTasksat or below partitions still coalescesnumTasksadds no shufflenumTasks=2numTasks=2(grouping off)numTasksabove task slots fails with the missing-task errorAll local runs used JDK 11, pinned to 2 CPUs with
taskset -c 0,1to match the CI agents:LightGBMRankerPartitionSuite: 27/27 pass.DriverSocketRetryE2ESuite,NetworkManagerSuite,BarrierNetworkRecoverySuite,WorkerWireFormatSuite,LightGBMValidationDataSuite,TrainUtilsSuite,VerifyLightGBMRankerStream,VerifyLightGBMRankerBulk,VerifyLightGBMRegressorStream): 83/83 pass.lightgbm/scalastyleandlightgbm/Test/scalastyleare clean.LightGBMRanker.scalaconflicted, because fix: preserve LightGBM ranker grouping under AQE #2744 hasn't reached the ports yet). Style is clean, andLightGBMNumTasksSuiteandDriverSocketRetrySuitepass (21/21 each, JDK 17, 2 CPUs).LightGBMNumTasksSuitepassed 25 of 25 repeated runs andDriverSocketRetrySuite8 of 8, on 2 CPUs. After review, the missing-partition test now starts its short 500ms accept timeout only after the first fake worker report is recorded, so a slow agent cannot time out before that report arrives; the suite then passed 5 more runs in a row.Fabric end-to-end
Ran on Microsoft Fabric (Spark 3.5.5 runtime, 1 executor with 8 task slots, AQE on, dynamic allocation off). Each run loaded this PR's jars (
...-227-e51ff84c-SNAPSHOT) and, for the baseline, master's jars (...-226-9d51ad1a-SNAPSHOT). The scenario checked that the classes came from those jars. Non-barrier mode, bulk data transfer:numTasks=4numTasks=4, grouping offnumTasks=4(2+ empty partitions)numTasks=9on 8 slots,timeout=90lightgbm-streamingscenarioStreaming data transfer without
slotNamesfails on this Fabric runtime for both master and this PR, with a native "Feature (Column_) appears more than one time" error during dataset creation. That's a separate, existing problem (see #2242), so the cases above use bulk mode.Docs
Added a troubleshooting section to the LightGBM overview: "When numTasks is larger than the tasks Spark can run at once". It includes the cost note above.
Note
Open PR #2745 also edits
NetworkManager.scala(worker identity diagnostics). The changes are in different methods, but whichever merges second may need a small rebase.