Repository navigation
[BUG] LightGBM streaming allocates executor buffers using the global partition count #2702
Description
Summary
PartitionTaskContext.executorPartitionCount returns the number of partitions in the whole training dataset, not the number assigned to the current executor. Streaming initialization passes that value to the native Dataset as its number of external ingestion threads.
This is a concrete SynapseML sizing defect found while continuing #2699. It can be corrected in Scala without changing the native dependency. It is not proof of the cause of the customer's connection failures or the SIGSEGV/SIGBUS reports in #2333.
Affected code
Confirmed on master commit 8c7143875c843c649a817cf3e8ba9c7bee23689c, using Spark 3.5.0, Scala 2.12.17, JDK 11, and lightgbmlib:3.3.510.
- BasePartitionTask.scala, lines 81-98:
threadIndexand row counts useexecutorPartitionIdList, butexecutorPartitionCountusestrainingCtx.partitionCounts.get.length. - ReferenceDatasetUtils.scala, lines 102-124: the incorrect count is passed to
LGBM_DatasetInitStreaming. - Native API contract:
nthreadsmeans the external threads using this Dataset's push APIs. - Source-era SparseBin implementation: streaming initialization resizes the push-buffer array to
omp_max_threads * num_thread.
The native links identify the API contract and a source-era implementation. They do not establish the exact unpublished source commit used to build the Maven .510 binary.
Reproduction
The existing investigation's StreamingLayoutSuite was rerun on September 7, 2026, against unchanged production code at the commit above:
| Global partition count | Executor-local partition IDs | Expected external-thread slots | Actual |
|---|---|---|---|
| 12 | [7, 2, 11] |
3 | 12 |
| 6 | [0, 2, 3], including empty local partitions |
3 | 6 |
| 3 | [0, 1, 2], all partitions local |
3 | 3 |
The first case can be reproduced by putting this test in the LightGBM split1 test package:
package com.microsoft.azure.synapse.ml.lightgbm.split1
import com.microsoft.azure.synapse.ml.core.test.base.TestBase
import com.microsoft.azure.synapse.ml.io.http.SharedSingleton
import com.microsoft.azure.synapse.ml.lightgbm._
import org.apache.spark.ml.linalg.SQLDataTypes
import org.apache.spark.sql.types.{StructField, StructType}
class StreamingLayoutReproSuite extends TestBase {
test("streaming initialization counts executor-local partitions") {
val counts = Array.fill(12)(10L)
val localPartitions = Array(7, 2, 11)
val features = StructField("features", SQLDataTypes.VectorType)
val params = new LightGBMRegressor()
.getTrainParams(counts.length, features, localPartitions.length)
val training = TrainingContext(
batchIndex = 0,
sharedStateSingleton = SharedSingleton(new SharedState(params)),
schema = StructType(Seq(features)),
numCols = 1,
numInitScoreClasses = 0,
trainingParams = params,
networkParams = NetworkParams(12400, "127.0.0.1", 12400, barrierExecutionMode = false),
columnParams = ColumnParams("label", "features", None, None, None),
datasetParams = "",
featureNames = None,
numTasksPerExecutor = localPartitions.length,
validationData = None,
serializedReferenceDataset = None,
partitionCounts = Some(counts))
val ctx = PartitionTaskContext(
trainingCtx = training,
partitionId = 7,
taskId = 7L,
measures = new TaskInstrumentationMeasures(7),
networkTopologyInfo = NetworkTopologyInfo("127.0.0.1:12400", localPartitions, 12400),
shouldExecuteTraining = true,
isEmptyPartition = false,
shouldReturnBooster = true,
shouldCalcValidationDataset = false)
assert(ctx.threadIndex == 1)
assert(ctx.executorRowCount == 30)
assert(ctx.streamingPartitionOffset == 10)
assert(ctx.executorPartitionCount == localPartitions.length)
}
}Run with the repository's JDK 11 SBT setup:
sbt "lightgbm/testOnly com.microsoft.azure.synapse.ml.lightgbm.split1.StreamingLayoutReproSuite"
The last assertion fails with 12 did not equal 3. The recorded baseline also fails the empty-partition case with 6 did not equal 3; the one-executor control passes.
Impact and proposed correction
Each executor should allocate for its own ingestion threads, not for tasks running on other executors. For native sparse bins, the source implementation above creates excess buffer slots and subsequently visits those slots during finalization. The overhead can grow as cluster-wide parallelism increases even when per-executor work does not.
This is over-allocation, not evidence of an out-of-bounds write. No executor RSS measurement, OOM reproduction, or claim that this change alone prevents a native crash is attached to this issue.
Derive the external-thread count from the same executor-local topology used by threadIndex. Preserve slots for empty local partitions so nonempty partitions' thread indexes stay valid. Do not substitute the number of global partitions, the number of native workers, or just the count of nonempty local partitions.
Acceptance criteria
- Count executor-local partition slots consistently with
threadIndex. - Cover noncontiguous and unsorted local IDs, empty local partitions, and the one-executor case.
- Exercise the public streaming estimator across multiple executors, including sparse native bins and validation data.
- Compare native allocation/finalization overhead with the same per-executor workload at different cluster sizes.
- Preserve predictions and parameter/persistence compatibility.
Track delivery under #2699. Keep #2333 open unless its independent corruption reproduction is also resolved.
Activity
github-actions commentedon Sep 7, 2026
Hey Rana Singh (@ranadeepsingh) 👋!
Thank you so much for reporting the issue/feature request 🚨.
Someone from SynapseML Team will be looking to triage this issue soon.
We appreciate your patience.
ranadeepsingh commentedon Sep 7, 2026
Reproduced through real Fabric training, 2026-09-07
This is now reproduced through public
LightGBMRegressor.fit()calls on Fabric, not just a constructedPartitionTaskContext. An opt-in diagnostic records the actual argument immediately beforeLGBM_DatasetInitStreaming.The first valid fit has 8,192 rows, 4,096 per executor. The local partition IDs are non-contiguous. For example, the baseline records:
In the corrected application, the partition groups were assigned to the opposite executor IDs, but each native initialization used its own local count of eight. The correction does not depend on a particular executor numbering or partition-placement order.
Workload and evidence
3.5.5.5.4.20260807.1, source base8c7143875c843c649a817cf3e8ba9c7bee23689c.numThreads=1,maxStreamingOMPThreads=1, and micro-batches of eight rows.total number of machines: 2. This is not inferred from an executor discovery job.The other four cases concern malformed vectors and are reported in #2703. They return models on the baseline and fail before native initialization on the corrected candidate.
What this proves
The wrong native allocation count occurs during actual distributed training, and the executor-local correction changes the argument to the required value without breaking these valid fits.
This is not an RSS or OOM measurement. Sixteen rather than eight external-thread slots does not mean total executor memory doubled. The run does not prove that this defect caused a customer's connection errors, resolves #2333, or fixes the separate categorical native crash.
Exact candidate identity
Both candidates use byte-identical core jars and the unchanged published native package:
The diagnostic source and compiled diagnostic class are identical in the pair. It is off by default and was enabled with
-Dsynapseml.lightgbm.streaming.layoutDiagnostics=true. Loaded class origins and both driver/executor native-library hashes were checked against the supplied artifacts.The first corrected application's initial discovery action observed only one executor and stopped before any fit. It is excluded above. The successful rerun retained the same Scala/native jars and added a bounded executor-readiness check to the Python diagnostic. All three scratch lakehouses were deleted successfully.
The correction and its regression tests remain in an isolated worktree. No dependency pin, published package, or Fabric runtime was changed.
test: preserve group-aligned topology in Fabric ranker control
ranadeepsingh commentedon Sep 15, 2026
The fix is now in draft PR #2707, with final-head Fabric evidence rather than only an uncommitted candidate.
Head:
412bdeade842db7509ee1f8cc693d500a4c87c17.The executor logs contain 12 and 24 distinct native Dataset initialization records, respectively: six successful fits on every executor, with local partition count and external slot count matching in every record. The one-executor control, noncontiguous IDs, empty local slots, thread bounds, and offsets also have local regression coverage.
Both complete applications passed, including repeated dense/sparse regressors, exact model save/load prediction equality, classifier/ranker controls, stable executor membership, and the expected driver/executor native hashes. The native dependency remains the published
lightgbmlib:3.3.510. All scratch lakehouses were cleaned up.This measures the native slot argument, not total RSS, an OOM fix, or the cause of every connection failure. It does not resolve #2697 or #2333.
The PR records 38 passing targeted Scala tests and passing final-head LightGBM Scala CI shards. It remains draft: Spark 4.1 compatibility replay needs missing master prerequisites, other CI legs have certificate/infrastructure failures or are pending, and automated review has not arrived. This issue stays open until the fix is merged; release/Fabric delivery is still separate.
fix: harden LightGBM streaming allocation and reference validation
fix: validate LightGBM streaming input and size executor-local buffers (