fix: scale LightGBM validation data transfer

[Rana Singh (ranadeepsingh)]

What this fixes

Closes #2294.

Large sparse validation sets supplied through validationIndicatorCol were preprocessed with collect() and broadcast as Array[Row]. That required Spark to return the complete validation set to the driver, exposing users to driver result-size and heap failures.

This change removes the large validation-row collection and broadcast while preserving exact LightGBM validation and early-stopping semantics.

Design

• Materialize every preprocessed validation row to a driver-local spool through DataFrame.foreachPartition.
• Preserve every row and partition order; no sampling or validation sharding is introduced.
• Broadcast only a small, versioned, authenticated stream descriptor while retaining legacy Array[Row] descriptor compatibility.
• Stream the complete spool to each task that owns a native LightGBM validation Dataset.
• In bulk single-Dataset mode, skip transfers on helper tasks that do not own the Dataset.
• Accept only the first complete Spark attempt for each partition so retries and speculation cannot overwrite committed data.
• Preserve primary failures while attaching iterator, socket, executor, Dataset, broadcast, and spool cleanup failures as suppressed exceptions.

LightGBM 3.3.5 calculates validation metrics locally, including non-decomposable metrics such as AUC. Each active native Dataset owner therefore still receives the complete validation set; changing that would alter metrics and early-stopping behavior.

Resource bounds

• Driver copy buffers: 64 KiB.
• Concurrent ingest or serving transfers: eight.
• Ingest and serving executors: nine daemon threads each.
• Socket operations and shutdown are time-bounded.
• Temporary attempt files, committed partition files, clients, executors, iterators, broadcasts, and spool directories are cleaned on success and failure.

The semantic tradeoff remains explicit: exact validation requires O(validation data) driver disk, one complete transfer per native Dataset owner, and one complete native validation Dataset per owner.

Compatibility

• No public estimator parameters, JVM signatures, or serialized parameter shapes change.
• TrainingContext.validationData retains its existing type and constructor position.
• Legacy validation broadcasts remain readable.
• Copy, schema, estimator/model persistence, sparse vectors, streaming mode, both bulk Dataset modes, and repeated fits are covered.
• Null validation indicators fail explicitly rather than being silently removed by SQL null filtering.

Rebase and dependency composition

• Rebased onto current master 3fac9b20e81d8e72c877781c746af171f39477ea.
• Pushed head: c0fe6bdd6e71440df7da9d03105f62d7be1cc5aa.
• The original 20 commits range-diff one-for-one after the rebase. The only intentional semantic differences are the two StreamingPartitionTask conflict compositions that retain both PR fix: clean up failed LightGBM streaming datasets #2662's native Dataset ownership cleanup and this PR's closeable streamed-row handling.
• PR fix: clean up failed LightGBM streaming datasets #2662 is merged. Its three rebased commits are declared in .pipelines/release-compat-prerequisites.txt so the Spark 4.1 replay receives those ownership helpers before this patch.

Validation

Spark 3.5 / Scala 2.12 / JDK 11

• lightgbm/compile: passed.
• lightgbm/Test/compile: passed.
• lightgbm/scalastyle: 0 findings.
• lightgbm/Test/scalastyle: 0 findings.
• Full codegen: passed.
• black==22.3.0 --check --extend-exclude docs/: 198 files unchanged.
• git diff --check: passed.
• Focused suites: 38/38 passed across:
  • LightGBMValidationDataSuite
  • ValidationDataServerLifecycleSuite
  • StreamingDatasetLifecycleSuite
  • ReferenceDatasetUtilsSuite
• Exact authentication follow-up: 34/34 passed across ValidationDataServerLifecycleSuite and
  LightGBMValidationDataSuite.
• Final CI-remediation sequence passed 30/30 lifecycle tests, 5/5 bounded-result validation tests,
  and a 1/1 post-suite repeated-fit test, proving the specialized Spark session no longer poisons
  the shared TestBase session.
• Release compatibility policy suite: 64/64 passed, including a scratch-repository regression that
  replays a prerequisite rename and verifies the source path is removed.

The public scaling regression uses 256 rows, 4,096-dimensional sparse vectors with 1,024 deterministic nonzeros, two LightGBM tasks, and spark.driver.maxResultSize=128k. Two tasks retain the legacy driver-result-size failure contract while fitting Azure's two concurrent Spark task slots.

Spark 4.1 / Scala 2.13 / JDK 17

• Replayed PR fix: clean up failed LightGBM streaming datasets #2662's three prerequisite commits in order onto live spark4.1
06897e5b27e28d84ce7ffa33e93d7f756992d0f2.
• Applied final head c0fe6bdd6e71440df7da9d03105f62d7be1cc5aa using the pipeline's
  release-relevant path filtering; the renamed suite exists only at its destination.
• Ran the repository's exact release-compatibility command, root test:compile, successfully on
  JDK 17 in 683 seconds.

Azure CI follow-up

Azure build 231829385 on the old head exposed a real test configuration defect: four LightGBM workers were requested on a two-slot runner, so the first two workers occupied every slot while waiting for workers that could not start. The regression now uses two workers, preserving multi-worker and active/helper coverage without scheduler starvation.

The automated review of 18750a514d3a079409c879a17912234b70216add found that token mismatches were classified as expected disconnects. Head 2878c517eba7ef6919a6dcb9f99995c54a3b7cd3 records authentication failures as terminal for both ingest and serving, preserves the original SecurityException, and covers both shutdown states plus the real serving/close path. The review thread was answered and resolved.

Copilot reviewed final head c0fe6bdd6e71440df7da9d03105f62d7be1cc5aa with zero active or suppressed findings; all review threads are resolved.

Exact-head Azure build 232332516 was a trigger-driven pullRequest build requested for GitHub against merge commit d0d1b44393743f52177ccac173b5d5363ff21a19. It exposed three independent failures:

• UnitTests lightgbm1: the bounded-result provider reused and then stopped the shared SparkContext, leaving later suites with a stopped cached session. Commit 47f6ac1db4e43a2ed0357eeff1e1c32253ec9c54 isolates the specialized session and restores a healthy shared session.
• Release Branch Compatibility Check spark4.1: name-only prerequisite enumeration collapsed a rename to its destination, so the path-filtered patch omitted deletion of the source. Commit c0fe6bdd6e71440df7da9d03105f62d7be1cc5aa enumerates with rename detection disabled and adds a representative replay regression.
• Fabric E2E: certificate token acquisition succeeded, but workspace cleanup received repeated HTTP 500 responses from the Fabric API before any product test ran. This is retained as infrastructure evidence and will be retried by the fresh exact-head build.

Fresh exact-head Azure build 232346646 validated both remediations: UnitTests lightgbm1
passed 143/143 executed tests, and the Spark 4.1 prerequisite replay plus root test:compile
passed. Every code job completed successfully. Fabric E2E again stopped before product tests
when GET https://api.fabric.microsoft.com/v1/workspaces returned HTTP 500; the log reports
zero tests run. Controlled rerun 232352794 ran against the same merge commit
5c9d0a7e47c7fde716c94a2ace420a2e98d8d0ee, whose parents are current master
3fac9b20e81d8e72c877781c746af171f39477ea and final PR head
c0fe6bdd6e71440df7da9d03105f62d7be1cc5aa. Every non-Fabric job passed again; Fabric
workspace discovery returned the same HTTP 500 and again ran zero tests.
Final cooldown build 232363074 produced the same result: all 75 non-parent/non-Fabric checks
passed, while Fabric workspace discovery failed before running any tests.

This is a target-branch infrastructure failure rather than a PR regression. Current-master builds
232301239 and 232305251 have the same sole failed job, the same workspace-listing HTTP 500,
and zero Fabric tests executed.

Managed Fabric proof

The checked-in lightgbm-streaming scenario passed independently on managed Fabric Spark
3.5.5.5.4.20260807.1 with exact head c0fe6bdd6e71440df7da9d03105f62d7be1cc5aa:

• Loaded the supplied core, lightgbmlib 3.3.510, and LightGBM jars; class-source and native
  mappings name those exact jars.
• Recorded SHA-256 hashes 1ae2c732..., f2b1b131..., and f2527db6... respectively.
• Completed two public LightGBMClassifier.fit() and prediction repetitions over 4,000 rows and
  four partitions; each prediction count was 4,000.
• Ran as Fabric application application_1787385228417_0001.
• Downloaded driver/executor logs and deleted the unique scratch lakehouse successfully.

Residual risk

• The implementation is bounded on driver heap but intentionally does not reduce driver disk, network, or native-worker validation size because doing so would change LightGBM semantics.
• Formal readiness remains blocked by the baseline Fabric CI status and the repository's required
  human approval, despite exact-head managed-runtime coverage passing independently.

[azure-pipelines]

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

[github-actions]

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.

[Rana Singh (ranadeepsingh)]

Lifecycle follow-up pushed at 3eaca96. Active ingest/serve sockets are now closed before executor shutdown; spool deletion requires confirmed termination; construction ownership is covered through bind, executor, and start failures; cancellation/accept timeouts remain nonterminal while real I/O failures propagate. JDK 11 compile/test-compile and both scalastyle tasks passed, Black passed, and the two targeted suites passed 13/13. Copilot automated review was requested and polled twice for 10 minutes but no review exists for this head. Azure was intentionally not triggered. Merge #2662 first, then rebase #2664 and preserve both cleanup changes.

[azure-pipelines]

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

[Copilot]

Pull request overview

This PR fixes LightGBM training failures when validationIndicatorCol is used with large/sparse validation sets by eliminating the driver-side collect() + broadcast of all validation rows. Instead, it spools preprocessed validation rows to driver-local disk and streams them to executors/tasks that construct native LightGBM validation Datasets, preserving full validation semantics (each worker still receives the complete validation set).

Changes:

• Add a driver-side ValidationDataServer that ingests per-partition validation rows via sockets into a spool directory, then serves the spool to executors via authenticated streaming.
• Update streaming and bulk partition tasks to consume validation rows via ValidationDataServer.read(...) (iterator + explicit close) and to use the streamed row count.
• Add lifecycle and end-to-end tests covering cleanup, persistence, bulk/stream modes, and null validation indicators.

Show a summary per file

File	Description
lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/ValidationDataServer.scala	New spool + stream server/descriptor format for scalable validation data transfer.
lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/LightGBMBase.scala	Replaces validation collect() broadcast path with server-backed broadcast descriptor and explicit null-indicator rejection.
lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/StreamingPartitionTask.scala	Streams validation rows into the shared streaming validation Dataset and closes the iterator.
lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/BulkPartitionTask.scala	Streams validation rows for bulk-mode aggregation and closes the iterator.
lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/ValidationDataServerLifecycleSuite.scala	New deterministic tests for server lifecycle cleanup and failure handling.
lightgbm/src/test/scala/com/microsoft/azure/synapse/ml/lightgbm/split1/LightGBMValidationDataSuite.scala	New integration tests ensuring validation streaming avoids driver result-size issues and covers persistence/bulk/null-indicator behavior.

Review details

💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Suppressed comments (1)

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/ValidationDataServer.scala:689

• Row lengths come from the network stream; a negative value other than the EndOfStream marker should be rejected explicitly. Without a guard this can throw NegativeArraySizeException (or lead to unexpected allocation behavior) instead of a clear I/O error.

          val length = input.readInt()
          if (length == EndOfStream) {
            close()
            None
          } else {

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

[Rana Singh (ranadeepsingh)]

Suppressed finding from Copilot review 4964793710 is explicitly fixed in db086b5: executor-side validation row lengths now pass through the same protocol validator as ingest, rejecting negative values other than the -1 end marker before allocation. A deterministic malformed executor-stream test covers that exact path. No upper cap was invented because the protocol has no established maximum serialized-row size and a new cap could reject valid sparse rows. Both active threads were also fixed, replied to, and resolved. Local evidence: Scala 2.12 main/test compile; full patch on spark4.0 Scala 2.13 main/test compile; 17/17 lifecycle/integration tests; both scalastyle tasks; Black. Azure was not triggered. #2662 must still merge first, followed by rebasing #2664 while preserving both StreamingPartitionTask cleanup changes.

[Copilot]

Review details

Suppressed comments (1)

Previously missed (1) — in code that hasn't changed since the last review.

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/ValidationDataServer.scala:678

• The ValidationDataServer.read(params) iterator can leak an open socket if authentication/stream setup fails before input is initialized (e.g., socket.getOutputStream, writeUTF, or flush throws). Wrap initialization in a try/catch that closes the socket on failure and only keep input as a field.

      private val socket = connect(params.host, params.port, params.timeoutMillis)
      socket.setKeepAlive(true)
      socket.setTcpNoDelay(true)
      socket.setSoTimeout(params.timeoutMillis)
      private val auth = new DataOutputStream(new BufferedOutputStream(socket.getOutputStream))
      auth.writeUTF(params.token)
      auth.flush()
      private val input = new DataInputStream(new BufferedInputStream(socket.getInputStream))

• 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).

[Rana Singh (ranadeepsingh)]

Addressed the new-head suppressed socket-setup finding in 965f55d. ValidationDataServer.read now closes its connected socket through the existing primary-preserving cleanup helper if socket configuration, authentication output/write/flush, or input-stream construction fails. A deterministic authentication-output failure test proves the original IOException remains primary and the socket is closed. Final validation: 18/18 lifecycle/integration tests; Scala 2.12 main/test compile; the full PR patch applied cleanly and main/test compiled on spark4.0 Scala 2.13/JDK 17; both scalastyle tasks; Black. Azure remains intentionally untriggered. #2662 still merges first, then #2664 rebases while retaining both StreamingPartitionTask cleanup changes.

[Copilot]

Review details

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/ValidationDataServer.scala:478

• timeoutMillis is calculated using IngestPollTimeoutMillis as the seconds→milliseconds conversion factor. This couples the socket timeout semantics to the accept-poll constant, so changing the poll interval would silently change all timeouts. Use an explicit seconds→ms conversion instead (e.g., timeoutSeconds * 1000.0).

      val timeoutMillis = (timeoutSeconds * IngestPollTimeoutMillis).toLong
      socket.setSoTimeout(Math.max(IngestPollTimeoutMillis, timeoutMillis).min(Int.MaxValue).toInt)

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/ValidationDataServer.scala:658

• rowCount converts a Long to Int via Math.toIntExact, which will throw a bare ArithmeticException if validation row count exceeds Int.MaxValue. Since downstream LightGBM APIs require Int row counts anyway, it would be clearer to fail with a targeted IllegalArgumentException explaining the limit.

  def rowCount(data: Broadcast[Array[Row]]): Int = {
    ValidationDataParams.fromBroadcast(data)
      .map(params => Math.toIntExact(params.rowCount))
      .getOrElse(data.value.length)
  }

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

[Rana Singh (ranadeepsingh)]

Addressed both suppressed findings from the 965f55d review in bcae907. Socket timeout conversion now uses an explicit seconds-to-milliseconds constant rather than the accept-poll interval, and validation row counts outside 0..Int.MaxValue fail with a targeted IllegalArgumentException describing the downstream native limit. Deterministic tests cover 2.5 seconds -> 2500 ms and Int overflow. Final Scala 2.12 evidence: main/test compile, both scalastyle tasks, Black, 20/20 lifecycle/integration tests. Azure remains intentionally untriggered; #2662 still merges first, then #2664 rebases preserving both StreamingPartitionTask cleanup changes.

[Copilot]

Review details

Suppressed comments (2)

Previously missed (2) — in code that hasn't changed since the last review.

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/StreamingPartitionTask.scala:165

• rows.close() can throw (socket/input close), which would mask an exception from insertRowsIntoDataset and make failures harder to diagnose. Suppress non-fatal close failures so the primary training/validation error remains primary.

    } finally {
      rows.close()
    }

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/BulkPartitionTask.scala:53

• rows.close() can throw (socket/input close), which would mask an exception from getChunkedColumns / mergeChunksIntoAggregatedArrays. Suppress non-fatal close failures so the primary error remains primary.

      } finally {
        rows.close()
      }

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

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

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

[Copilot]

Review details

• Files reviewed: 9/9 changed files
• Comments generated: 1
• 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)

Previously missed (1) — in code that hasn't changed since the last review.

lightgbm/src/main/scala/com/microsoft/azure/synapse/ml/lightgbm/ValidationDataServer.scala:350

• spoolDirectory.listFiles() can return null on I/O error. Treating that as Array.empty can silently start a server with no partition files (while rowCount is non-zero), leading to confusing downstream failures and potentially incorrect/partial validation reads. Fail fast if the directory listing fails, and validate that the expected number of part-* files were produced.

      val partitionFiles = Option(spoolDirectory.listFiles()).getOrElse(Array.empty)
        .filter(_.getName.startsWith("part-"))
        .sortBy(file => file.getName.stripPrefix("part-").toInt)
      spoolTransferred = true

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

[Rana Singh (ranadeepsingh)]

Addressed the current-head suppressed spool finding in b8f9f6a. Validation ingest now fails before serving if the spool directory cannot be listed or lacks the complete canonical part-0..part-N sequence. Added unreadable, incomplete, and ordered-valid regressions; the 40-test validation selection, compile/test-compile, and main/test scalastyle pass.

[Rana Singh (ranadeepsingh)]

/azp run

[azure-pipelines]

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

[Copilot]

Review details

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

[Rana Singh (ranadeepsingh)]

Current-head readiness update for b8f9f6a2ad7bc34b21c9bce2e7d92d1c111cd45d:

• 0 commits behind current master (57c799b834080681defec358925ef2a7d5d5eb95).
• Exact-head Azure build 232663670 passed all 66 jobs, including Fabric E2E, Databricks CPU/GPU, Python LightGBM, all six LightGBM unit shards, and Spark 4.1 compatibility.
• Current-head Copilot review generated zero findings and all review threads are resolved.

Scott Votaw (@svotaw) Brendan Walsh (@BrendanWalsh) Kirat Pandya (@kiratp) — could one of you review and provide the required approval for this current head?

[Rana Singh (ranadeepsingh)]

Final exact-head managed-Fabric evidence for b8f9f6a2ad7bc34b21c9bce2e7d92d1c111cd45d:

• lightgbm-streaming passed on workspace singhrana_ws with Fabric Spark 3.5.5.5.4.20260807.1.
• The managed application loaded all classes from the supplied jars: synapseml-core_2.12-HEAD-20260824-2222.jar, lightgbmlib-3.3.510.jar, and synapseml-lightgbm_2.12-HEAD-20260824-2222.jar.
• SHA-256 provenance:
  • core: c455e279b1042e1f5b577066254c71db6bc99ba164ad66977c0e2122f337c352
  • native LightGBM: f2b1b13172699832594303ab4c04f3bc8fc2d24737e3e8c11d98d69a88c09272
  • SynapseML LightGBM: a817fc1a6dc144e1add12210b8c30fdbb1e507d6a5ef56b540803957eed180b0
• Runtime topology: 2 block managers, 4 partitions, 8 executor cores; two independent fits each produced 4,000 predictions.
• JUnit: 1 test, 0 failures. Scratch-lakehouse cleanup completed with exit code 0.

This supplements the fully green exact-head Azure build 232663670.