fix: scale and harden LightGBM validation data transfer

[Rana Singh (ranadeepsingh)]

Summary

Fixes #2294.

Large LightGBM validation sets no longer have to be collected into the Spark
driver heap and broadcast as a complete Array[Row]. Validation partitions are
streamed to an authenticated driver-local spool, represented by a small
versioned descriptor, and streamed from that spool to each native LightGBM
Dataset owner.

This preserves exact validation metrics and early-stopping behavior while
removing the spark.driver.maxResultSize and driver-heap bottleneck reported for
large sparse validation data.

What changed

• Execute validation ingestion with Dataset mapPartitions; no new RDD
  implementation is introduced and validation rows never become Spark task
  results.
• Write each Spark task attempt to a separate spool file. The task waits for a
  driver acknowledgement, returns only
  (partitionId, taskAttemptId, rowCount), and the driver promotes only the
  attempts Spark actually committed.
• Broadcast only a versioned endpoint descriptor. Genuine legacy
  Broadcast[Array[Row]] values remain readable, while unknown future
  descriptor versions fail clearly.
• Authenticate ingest and serving with separate random UUID tokens and bind
  both listeners to the resolved driver host.
• Bound concurrent transfers to eight, reuse fixed 64 KiB buffers, and enforce
  write-inactivity deadlines with one daemon watchdog per JVM. Queued readers
  start their read timeout after admission rather than timing out behind healthy
  transfers.
• Preserve the first asynchronous/training failure and attach later cleanup
  failures as suppressed exceptions.
• Reject null validation indicators inside the existing Catalyst filter action,
  avoiding a separate full-data preflight scan.

File-by-file necessity

File	Why it is required
.github/skills/synapseml-branches/references/branch-spark3p5.md	Records Fabric Runtime 1.3 Python lightgbm 4.3.0 versus JVM lightgbmlib 3.3.510.
.github/skills/synapseml-branches/references/branch-spark4-common.md	Records Runtime 2.0 Python lightgbm 4.6.0, the same JVM artifact, and the Spark 4.0 comparison rule.
.pipelines/release-compat-prerequisites.txt	Replays the three Spark 4.1 Dataset-ownership prerequisites required by this change.
BulkPartitionTask.scala	Ensures validation data is read only by native Dataset owners and keeps the transport comment accurate.
LightGBMBase.scala	Removes driver collection, splits training/validation with DataFrame expressions, and scopes the server/broadcast lifecycle.
NetworkManagerSocketSupport.scala	Preserves primary asynchronous failures and provides bounded write-inactivity cancellation.
StreamingPartitionTask.scala	Consumes the descriptor-backed validation iterator with deterministic native Dataset cleanup.
ValidationDataIngest.scala	Implements acknowledged, task-attempt-aware partition ingestion and successful-attempt promotion.
ValidationDataServer.scala	Owns the authenticated spool ingest/serve lifecycle, bounded transfers, descriptor compatibility, and cleanup.
ValidationDataSpool.scala	Validates and orders the canonical contiguous part-N spool files.
LightGBMValidationDataSuite.scala	Covers public classifier/ranker behavior, persistence, null handling, and no-driver-collection scaling.
ValidationDataIngestRetrySuite.scala	Reproduces a task that sends all rows and then fails before Spark commit; retry rows must win.
ValidationDataServerFlowControlSuite.scala	Proves a ninth reader survives queueing behind the eight-transfer cap while active clients keep progressing.
ValidationDataServerLifecycleSuite.scala	Covers authentication, malformed frames, timeouts, setup failures, stalled clients, and cleanup ordering.
ValidationDataServerSupportSuite.scala	Covers first-failure retention, write watchdog behavior, buffer reuse, and spool validation.
pipeline.yaml	Replays prerequisite renames as delete/add changes with git diff --no-renames.
tools/ci/tests/test_pipeline_yaml.py	Proves the Spark 4.1 prerequisite replay preserves renamed-file semantics.

Every changed file supports production behavior, a required cross-version
prerequisite, documentation requested for branch decisions, or a regression
that directly protects those paths.

Security and credentials

• No Key Vault, tenant, account, certificate, workspace-resolution, or
  INTEGRATION_* behavior changes in this PR.
• The existing credential chain remains unchanged:
  fabrictest-cert-admin-kv -> bami-tenant-adminuser -> SemPy AdminUser01.
• No secret values were added. The final diff contains no secret-like
  assignments, and retained Fabric logs contain token-related diagnostic names
  but no token values.
• Tokens are not logged, dynamic spool names use validated numeric IDs, attempt
  files use atomic CREATE_NEW, and both server paths close active sockets
  before spool deletion.
• A dedicated security review found no concrete vulnerability in binding,
  authentication, protocol framing, filesystem paths, cleanup, serialization,
  legacy compatibility, or credential handling.

Fabric LightGBM versions

Fabric baseline	Spark	Python package	JVM/SWIG artifact
Runtime 1.3	3.5	lightgbm 4.3.0	com.microsoft.ml.lightgbm:lightgbmlib:3.3.510
Runtime 2.0	4.1	lightgbm 4.6.0	com.microsoft.ml.lightgbm:lightgbmlib:3.3.510
Spark 4.0 port	No separate managed Fabric runtime	Use Runtime 2.0 as the comparison point	lightgbmlib:3.3.510

Both managed runtimes load a lightgbmlib-3.3.510.jar whose SHA-256 is
f2b1b13172699832594303ab4c04f3bc8fc2d24737e3e8c11d98d69a88c09272,
byte-for-byte identical to Maven Central. The 4.3.0 and 4.6.0 values are Python
package versions, not Maven/JNI versions, so this PR correctly leaves
lightgbmlib unchanged.

Compatibility and limits

• No public estimator parameters, JVM signatures, generated Python API, or
  serialized parameter shapes change.
• TrainingContext.validationData keeps its existing type and constructor
  position.
• Model save/load behavior remains unchanged.
• The fix removes O(validation data) driver heap/result-size pressure. It
  intentionally still requires O(validation data) driver disk, one complete
  network transfer per native Dataset owner, and native-worker memory because
  exact metrics such as AUC require the complete validation set.

Validation

Final target: master at 3498243d0b5534117eb24c0c3cebde3eb1c1a963
Final head: 8b90ff513749fdf5a2c2521e06d1f0e785f71d34

• Branch state: 32 commits ahead and 0 behind current master.
• Spark 3.5 / Scala 2.12 / JDK 11:
  • LightGBM compile, test-compile, main scalastyle, and test scalastyle passed.
  • Focused public-estimator, retry, flow-control, protocol, lifecycle,
    streaming-ownership, and reference-Dataset coverage passed 53/53.
  • The exact-head null-indicator assertion was rerun after making its exception
    supertype portable across Spark versions.
• Spark 4.1 / Scala 2.13 / JDK 17 compatibility replay:
  • All three required prerequisite commits retained.
  • Compile, test-compile, main scalastyle, and test scalastyle passed.
  • The same focused suites passed 53/53 with the required
    java.prefs/java.util.prefs module opening.
• Pipeline/release-replay contracts: 50 passed, 27 platform-specific skips.
• Black 22.3.0: 201 files unchanged.
• Reviews:
  • Dedicated security review: no concrete vulnerabilities.
  • General code review found three scale-path defects (retry-admission timeout,
    unbounded blocked writes, and an extra null scan); all three were fixed and
    now have passing regressions.
  • Automated review covered the complete 17-file change set. Its Boolean-cast
    and tautological-condition findings were fixed in 8b90ff51 and passed
    targeted Spark 3.5 and Spark 4.1 compile, style, and 10/10 tests on each
    baseline.
  • A follow-up automated review covers exact final head 8b90ff51; it reported
    no new or suppressed findings. All 11 review threads are resolved.
• Managed Fabric lightgbm-streaming on exact final jars:
  • MSIT tenant, DONT_DELETE_SynapseML_Build workspace, Fabric Spark
    3.5.5.5.4.20260807.1.
  • Application application_1787794769102_0001.
  • 4,000 rows, four partitions, two repeated fits, prediction counts
    4,000 / 4,000.
  • Runtime class and native-library provenance names the supplied core,
    lightgbmlib-3.3.510, and final-head LightGBM jars.
  • JUnit 1/1 passed; driver/executor logs retained; unique scratch lakehouse
    deleted successfully (cleanupExitCode: 0).
  • Evidence:
    pr-2664-fabric-e2e-8b90ff51/evidence.json in the retained session
    artifacts.
• Current-head GitHub and Azure validation:
  • Azure pull-request build 233012758 completed successfully for
    8b90ff513749fdf5a2c2521e06d1f0e785f71d34; all 76 PR check contexts pass
    with no failed, pending, or missing required checks.
  • The first RTests core attempt was blocked before tests by a transient
    Azure-agent TLS mismatch (msdata.visualstudio.com received an unrelated
    *.azureedge.net certificate). Retrying the failed job in the same build
    executed the workload and published 70/70 passing JUnit tests.
  • Final repository readiness reports zero unresolved review threads, zero
    suppressed final-head findings, and completeness.complete = true.

Repository policy still requires human approval before merge.

[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: 12/12 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: 12/12 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: 12/12 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

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

[Copilot]

Review details

• Files reviewed: 17/17 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).