fix: preserve the real cause when a LightGBM task retry cannot rejoin the network

[Rana Singh (ranadeepsingh)]

Problem

Distributed LightGBM training exchanges network topology with the driver exactly once per training round, then closes driverServerSocket. Any Spark task retried after that point reconnects to an endpoint that is no longer listening, so every retry fails with a bare java.net.ConnectException: Connection refused.

Because Spark only reports the most recent attempt, the failure that actually caused the retry is never surfaced:

org.apache.spark.SparkException: Job aborted due to stage failure: Task 4 in stage 103.0 failed 4 times,
most recent failure: Lost task 4.3 in stage 103.0 (TID 1954) (vm-... executor 2):
java.net.ConnectException: Connection refused

This makes such failures undiagnosable from the job output alone. It is reproduced end to end in DriverSocketRetryE2ESuite, which fails one task after it completes the driver handshake and produces exactly the signature above on unpatched code.

Barrier execution mode is affected too, and more seriously: a barrier stage restarts in its entirety, which is the only recovery a fixed LightGBM network can support — but the restarted stage also found the port closed, so barrier mode could not actually recover from anything.

Changes

Stop masking the real cause — a terminal ConnectException in getGlobalNetworkInfo is translated into an actionable error. On a retry it states that the topology exchange is already closed, that a partial task retry can never rejoin the network, and that the real cause is in the logs of the first failed attempt of that partition. On a first attempt it points at driver reachability instead.

Make barrier stage restarts work — the driver now keeps serving topology rounds in barrier mode until training is done. Task messages carry the Spark stage attempt number so the driver can tell attempts apart: a newer attempt discards the previous attempt's partial topology, and a straggler from a superseded attempt is rejected. This matters because a topology containing a dead host:port makes every surviving task hang trying to connect to it. Message parsing accepts both the old and new formats. A barrier round completes only after the current attempt has both its finished marker and every unique partition report; retried reports replace the prior socket instead of inflating the count.

Preserve the public message API — stage-attempt metadata is carried by a private wire message while TaskMessageInfo keeps its public five-field constructor, extractor, product shape, and string representation.

Fix the socket and thread lifecycle

• The driver network thread closes its sockets in a finally, so the server socket is no longer leaked when a round throws or times out.
• closeConnections is idempotent and synchronized, closes each socket independently, and closes the accepted-but-not-yet-registered socket plus the load-only helper sockets that were leaked on every fit().
• LightGBMBase.executeTraining closes the network manager in a finally. Previously a failure in the partition tasks left the network thread parked in accept() holding the driver port for up to the configured timeout (20 minutes by default).
• The barrier finished-marker connection is closed rather than leaked.
• Cleanup warnings retain the throwable and stack trace.

Algorithm and performance

No change to the training algorithm, the LightGBM parameters, or the data path. The new error handling only runs on a path that was already failing. Per-report work is a constant-time partition-map update; topology strings are built once per completed round. getGlobalNetworkInfo runs in initialize, before preparePartitionData, so a failing retry costs ~800ms and does not redo data preparation.

Tests

Suite	Covers
DriverSocketRetryE2ESuite	Full Spark job: a task fails after the handshake; asserts the reported error explains the cascade instead of Connection refused
DriverSocketRetrySuite	Refused retry, duplicate partition replacement, public message compatibility, helper socket release, idempotent close, unexpected close, and accept/shutdown races
BarrierNetworkRecoverySuite	Restarted stage isolation, stale markers, both marker/report orderings, and duplicate reports before or after completion signaling

Validated locally: 68 tests across 8 suites pass, including VerifyLightGBMClassifierStreamBasic (barrier mode), VerifyLightGBMRegressorBulk, VerifyLightGBMRegressorStream and VerifyLightGBMRankerBulk. Latest focused validation: 19/19 socket lifecycle and barrier recovery tests. scalastyle and Test/scalastyle are clean.

Related

Distinct from #2595, which fixes a native executor-to-executor listen port race. That failure surfaces via LightGBMUtils.validate as a generic Exception, not ConnectException, so the two are complementary rather than overlapping.

Update: barrier round completion fix

Review of this branch found a hang introduced by the earlier round-completion change, now fixed in this PR.

roundComplete required reportedTaskCount == numTasks. numTasks is the requested task count, not the number of tasks a stage actually runs: determineNumTasks never clamps a user-set numTasks, and barrier-mode prepareDataframe only repartitions downwards. So setNumTasks(16) on an 8-partition DataFrame ran 8 tasks, the equality never held, and the driver waited on Spark's barrier timeout (365 days by default) while every executor blocked on a read with no timeout.

Partition 0 now reports the stage's real task count (BarrierTaskContext.getTaskInfos().length) in the finished marker, and the driver completes the round once that many unique partitions have reported. Completing on the marker alone -- the pre-PR behavior -- has its own race: reports still sitting in the accept backlog get dropped from the topology.

Compatibility: the count is appended to the existing marker, so a sender that omits it still parses and the marker alone completes the round (previous behavior, never a hang).

The wire protocol moved to WorkerMessage.scala; NetworkManager.scala was at 797 of the 800-line scalastyle limit.

Validation: 19 tests across 3 lightgbm suites pass on Java 11; the new A barrier stage smaller than numTasks completes instead of hanging test was confirmed to fail against the unfixed condition.

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[Codecov Comments Bot (codecov-commenter)]

Codecov Report

❌ Patch coverage is 93.77778% with 14 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.99%. Comparing base (7088a6a) to head (686e86e).

Files with missing lines	Patch %	Lines
...oft/azure/synapse/ml/lightgbm/NetworkManager.scala	91.72%	12 Missing ⚠️
...apse/ml/lightgbm/NetworkManagerSocketSupport.scala	98.14%	1 Missing ⚠️
...soft/azure/synapse/ml/lightgbm/WorkerMessage.scala	95.45%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #2612      +/-   ##
==========================================
- Coverage   86.81%   82.99%   -3.83%     
==========================================
  Files         336      338       +2     
  Lines       18275    18384     +109     
  Branches     1673     1697      +24     
==========================================
- Hits        15866    15258     -608     
- Misses       2409     3126     +717     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[Copilot]

Pull request overview

Warning

Copilot couldn't run its full agentic review because it didn't start before the timeout. Make sure your repository has a runner available, or add a copilot-code-review.yml file specifying one with the runs-on attribute. See the docs for more details.

Improves diagnosability and recovery of distributed LightGBM training when Spark retries tasks after the driver’s topology exchange socket has closed, with special handling to enable barrier-stage restarts.

Changes:

• Adds clearer error translation for topology handshake ConnectException so retries don’t mask the original failure.
• Extends the driver topology server lifecycle for barrier mode by serving multiple topology rounds and tracking Spark stage attempts.
• Fixes socket/thread lifecycle cleanup, and adds targeted unit + E2E coverage plus documentation guidance.

Show a summary per file

File	Description
lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/NetworkManager.scala	Adds retry-aware ConnectException messaging, stage-attempt-aware handshake protocol, barrier restart support, and more robust socket lifecycle management.
lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/LightGBMBase.scala	Ensures driver network resources are released in a finally block when training fails.
lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/DriverSocketRetrySuite.scala	Adds unit tests for refused retries, helper socket cleanup, idempotent close, and early-failure cleanup.
lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/DriverSocketRetryE2ESuite.scala	Adds an E2E Spark test that verifies retries surface an actionable message rather than only “Connection refused”.
lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/BarrierNetworkRecoverySuite.scala	Adds tests validating barrier stage restart gets a fresh topology and stragglers are excluded.
docs/Explore Algorithms/LightGBM/Overview.md	Documents why non-barrier retries can’t rejoin the network, and how to diagnose “Connection refused”.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 3
• Review effort level: Lite

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines will not run the associated pipelines, because the pull request was updated after the run command was issued. Review the pull request again and issue a new run command.

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[Copilot]

Review details

Suppressed comments (1)

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/NetworkManager.scala:517

• handleNextWorkerConnection() assumes reader.readLine() always returns a non-null message. If a worker connects and closes without sending (or the connection is otherwise truncated), messageStr will be null and parseWorkerProtocolMessage will throw a NullPointerException, failing the driver network thread with a non-actionable error. Handle null explicitly by closing the socket and treating the connection as a non-report.

        val reader = new BufferedReader(new InputStreamReader(socket.getInputStream))
        val messageStr = reader.readLine()
        log.info(s"received worker message string: $messageStr")
        processWorkerConnection(socket, NetworkManager.parseWorkerProtocolMessage(messageStr))
      } catch {

• Files reviewed: 6/6 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[Copilot]

Review details

Suppressed comments (1)

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/NetworkManager.scala:762

• handleNextWorkerConnection assumes the worker will always send a line of text. If the accepted socket is closed before a message is written (e.g., shutdown racing with read, or a peer that connects and closes), reader.readLine() returns null and parseWorkerProtocolMessage will throw a NullPointerException. This can make shutdown/recovery paths fail with an unrelated NPE instead of a controlled SocketException (which barrier-mode already treats as expected when shutting down).

        val reader = new BufferedReader(new InputStreamReader(socket.getInputStream))
        val messageStr = reader.readLine()
        log.info(s"received worker message string: $messageStr")
        processWorkerConnection(socket, NetworkManager.parseWorkerProtocolMessage(messageStr))
      } catch {

• Files reviewed: 6/6 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Rana Singh (ranadeepsingh)]

/azp run

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[Copilot]

Review details

• Files reviewed: 8/8 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).

[Copilot]

Review details

Suppressed comments (1)

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/WorkerMessage.scala:48

• WorkerMessage.parse throws a generic Exception for unexpected task-message field counts. This makes it harder to distinguish parse/protocol errors from other failures higher up the stack (and loses the expected-vs-actual context). Using an IOException (or IllegalArgumentException) with a more actionable message would improve diagnosability without changing control flow.

      if (components.length != TaskMessageFieldCount && components.length != TaskMessageFieldCountWithStageAttempt) {
        throw new Exception(s"Unexpected message: $message")
      }

• Files reviewed: 8/8 changed files
• Comments generated: 0 new
• Review effort level: Lite

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

Azure Pipelines:
Successfully started running 1 pipeline(s).