From 51c2deec72ec3e4b1ab3598e1babee7dc03c95c3 Mon Sep 17 00:00:00 2001 From: Michal Harakal Date: Sun, 6 Sep 2026 15:12:42 +0200 Subject: [PATCH] fix(backend-cpu): CoroutineSchedule no longer deadlocks when entered from its own pool The jvm CI leg has timed out at 40 minutes on three of the last four runs since SKEEP-005 (#1262), hanging inside skainet-backend-cpu:jvmTest with 11 GB of memory free. It was not the OOM #1264 assumed but a coroutine deadlock: CoroutineSchedule.forRange ran a region as runBlocking { coroutineScope { launch(Dispatchers.Default) ... } }, and a coroutineScope waits for every child, including the ones the pool never got a thread for. On the 4-vCPU runner, once every Dispatchers.Default worker was inside a region (CoroutineScheduleTest does exactly that), nobody was left to run the children. Never reproduced on a 14-core laptop. A region is now a shared chunk queue: tasks - 1 helpers are dispatched to the dispatcher's executor, the caller runs chunk 0 and then drains the queue itself, and it waits only for chunks a thread has already claimed. A caller can therefore always finish a region alone, whatever the pool is doing; a helper that never ran finds the queue empty and exits. The contract is unchanged: first failure stops unstarted chunks and is rethrown once running ones finish (later failures suppressed), nested regions run inline, writes happen-before the return. Public API unchanged. Tests: a new deterministic pool-exhaustion test (two callers saturate a two-thread pool) times out on the old implementation and passes now; the caller-runs-first-chunk test allows the caller to help with later chunks. Verified with the test JVM pinned to 4 processors (-XX:ActiveProcessorCount=4), which reproduced the CI hang before the fix. Docs: schedules explanation and the SKEEP-005 proposal describe the new mechanism; the build.yml comment on the jvm leg states the real cause. Co-Authored-By: Claude Fable 5.1 --- .github/workflows/build.yml | 18 +-- .../ROOT/pages/explanation/schedules.adoc | 25 ++-- .../005-schedules-structured-concurrency.adoc | 29 ++-- .../ainet/exec/schedule/CoroutineSchedule.kt | 124 +++++++++++++----- .../exec/schedule/CoroutineScheduleTest.kt | 43 +++++- 5 files changed, 177 insertions(+), 62 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 8a4f7150..3c4b85dd 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -40,14 +40,16 @@ jobs: jvmTest :skainet-apps:skainet-plan:test :skainet-backends:benchmarks:jvm-cpu-publish:test - # This leg runs nearly every module's jvmTest sequentially, each forking its own - # 8g-heap Test JVM (build.gradle.kts `maxHeapSize`). ci-gradle.properties already caps - # org.gradle.workers.max=2 repo-wide, but on this leg two 8g forks running at once plus - # the 4g Gradle client JVM (GRADLE_OPTS below) can still exceed the 16 GB runner — the - # same OOM signature described above ("The operation was canceled" with no BUILD - # FAILED), just needing the extra memory pressure SKEEP-005's concurrency tests - # (KernelDispatchConcurrencyTest, CoroutineSchedule, parallel SDPA) added to tip it - # over. Force this leg fully serial so at most one 8g fork runs at a time. + # This leg runs nearly every module's jvmTest, each forking its own 8g-heap Test JVM + # (build.gradle.kts `maxHeapSize`). ci-gradle.properties caps org.gradle.workers.max=2 + # repo-wide; this leg goes fully serial so at most one 8g fork runs next to the 4g + # Gradle client JVM (GRADLE_OPTS below) on the 16 GB runner. + # + # History: the 40-minute reds after SKEEP-005 (#1262) were NOT memory — the runner had + # 11 GB free when the job was killed. `CoroutineSchedule.forRange` blocked a + # Dispatchers.Default worker on children that needed the same 4-thread pool, so + # `skainet-backend-cpu:jvmTest` deadlocked in CoroutineScheduleTest on the 4-vCPU + # runner. Fixed in the schedule itself; the serial setting stays for memory headroom. extraGradleArgs: -Dorg.gradle.workers.max=1 # verifyNpmPins guards the npm-* pins in gradle/libs.versions.toml against # lockfile drift; it belongs on the leg that already has the JS toolchain. diff --git a/docs/modules/ROOT/pages/explanation/schedules.adoc b/docs/modules/ROOT/pages/explanation/schedules.adoc index 9762ea72..f3c68834 100644 --- a/docs/modules/ROOT/pages/explanation/schedules.adoc +++ b/docs/modules/ROOT/pages/explanation/schedules.adoc @@ -50,16 +50,21 @@ public interface Schedule { It is deliberately small, non-suspending and dependency-free: ops are synchronous on every Kotlin target, and `skainet-lang-core` takes no coroutine dependency. The JVM implementation, -`CoroutineSchedule` in `skainet-backend-cpu`, is structured concurrency in the textbook sense — -one region is one `coroutineScope`: - -* the calling thread runs the first chunk itself, so no core idles blocked on the join; -* the other chunks are `launch`ed on the dispatcher (`Dispatchers.Default` by default); -* the scope closes only when every child has finished or been cancelled — the first failure - cancels the siblings and is rethrown to the caller; -* a region reached from inside another region runs inline, so the dispatcher never waits on - itself; `CoroutineSchedule.dedicated(n)` owns its own pool for callers that already live on - `Dispatchers.Default`. +`CoroutineSchedule` in `skainet-backend-cpu`, keeps the structured-concurrency guarantees without +ever letting the caller block on work that still needs a pool thread. One region is one shared +queue of chunks: + +* the calling thread runs the first chunk itself, then keeps claiming chunks until the queue is + empty, so no core idles blocked on the join; +* `tasks - 1` helpers are dispatched to the dispatcher (`Dispatchers.Default` by default) and + claim chunks from the same queue; a helper the pool never started finds the queue empty and + exits without running a body; +* the region returns only when every *claimed* chunk has finished — the first failure stops the + chunks that have not started and is rethrown to the caller once the running ones are done; +* because the caller only waits for chunks that are already running somewhere, a region entered + from a thread of the dispatcher's own pool cannot deadlock, however many pool threads are busy; +* a region reached from inside another region runs inline; `CoroutineSchedule.dedicated(n)` owns + its own pool for callers that want isolation from `Dispatchers.Default`. == The contract a body must keep diff --git a/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc b/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc index a02f239d..af85c2f0 100644 --- a/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc +++ b/docs/modules/skeep/pages/005-schedules-structured-concurrency.adoc @@ -93,13 +93,22 @@ cancelled; all writes happen-before the return. === The JVM implementation (`skainet-backend-cpu`) -`CoroutineSchedule(dispatcher = Dispatchers.Default, parallelism = cores, sink)`: a region is -`runBlocking { coroutineScope { launch(dispatcher) { … } … ; body(first chunk) } }` — the caller -runs the first chunk itself, so no core idles on the join; a nested region (detected with a -thread-local) runs inline so the dispatcher never waits on itself; `dedicated(parallelism)` owns -its own pool for callers that already live on `Dispatchers.Default`. `parallelChunks(outputDim, -schedule)` replaces the old island; the Panama Q4_K/Q5_K SPI kernels gained a schedule-aware -overload; `ScheduleRegion` events report what ran when a sink is attached. +`CoroutineSchedule(dispatcher = Dispatchers.Default, parallelism = cores, sink)`: a region is a +shared queue of `tasks` chunks — `tasks - 1` helpers are dispatched to the dispatcher, the caller +runs the first chunk itself and then drains the queue, and finally waits only for chunks some +thread has already claimed. A caller can therefore always finish a region on its own, which is +what makes it safe to enter from `Dispatchers.Default`. A nested region (detected with a +thread-local) runs inline; `dedicated(parallelism)` owns its own pool for callers that want +isolation from `Dispatchers.Default`. `parallelChunks(outputDim, schedule)` replaces the old +island; the Panama Q4_K/Q5_K SPI kernels gained a schedule-aware overload; `ScheduleRegion` +events report what ran when a sink is attached. + +NOTE: The first merged version ran a region as +`runBlocking { coroutineScope { launch(dispatcher) { … } … } }`. A `coroutineScope` waits for +every child, including the ones the pool never got to start, so once every `Dispatchers.Default` +worker was inside a region the pool had nobody left to run the children and the region hung. +That is exactly the 4-vCPU CI runner's situation; it never reproduced on a 14-core laptop. The +chunk-queue region above replaced it. === First consumer: `scaledDotProductAttention` @@ -169,8 +178,10 @@ bodies. Behaviour change on the JVM only: `DirectCpuExecutionContext()` now runs == Risks -* **Nested `runBlocking` on `Dispatchers.Default`** — mitigated by caller participation, inline - nested regions and `dedicated()`; tested from inside a Default worker. +* **Blocking a `Dispatchers.Default` worker on the pool it belongs to** — the region never + waits for a chunk that still needs a pool thread: the caller drains the queue itself and only + joins chunks already running. Tested from inside a Default worker and, deterministically, from + every thread of a saturated two-thread pool. * **A body that touches the context** — the contract is documented on `forRange`; the parity tests run bodies on foreign threads. * **JIT-dependent lane reductions** — the Panama matmul's `reduceLanes` order changes when the diff --git a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt index 037b7d61..0057c289 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmMain/kotlin/sk/ainet/exec/schedule/CoroutineSchedule.kt @@ -3,30 +3,40 @@ package sk.ainet.exec.schedule import kotlinx.coroutines.CoroutineDispatcher import kotlinx.coroutines.Dispatchers import kotlinx.coroutines.asCoroutineDispatcher -import kotlinx.coroutines.coroutineScope -import kotlinx.coroutines.launch -import kotlinx.coroutines.runBlocking +import kotlinx.coroutines.asExecutor import sk.ainet.context.schedule.Schedule import sk.ainet.lang.memory.ExperimentalMemoryApi import sk.ainet.lang.memory.trace.NoopTraceSink import sk.ainet.lang.memory.trace.TraceClock import sk.ainet.lang.memory.trace.TraceEvent import sk.ainet.lang.memory.trace.TraceSink +import java.util.concurrent.CountDownLatch +import java.util.concurrent.Executor import java.util.concurrent.ExecutorService import java.util.concurrent.Executors import java.util.concurrent.atomic.AtomicInteger +import java.util.concurrent.atomic.AtomicReference /** - * The JVM [Schedule] (SKEEP-005): a region is one `coroutineScope` — structured concurrency — - * whose children run on [dispatcher] while the *calling thread runs the first chunk itself*, so - * no core sits idle blocked on the join. The scope closes only after every child finished or was - * cancelled, which gives the contract its guarantees: the first failure cancels the siblings and - * is rethrown, no task outlives [forRange], and every task's writes happen-before the return. + * The JVM [Schedule] (SKEEP-005): a region is a queue of `tasks` chunks shared between the + * calling thread and `tasks - 1` helpers dispatched to [dispatcher]. The caller runs the first + * chunk itself, then keeps claiming chunks until the queue is empty, and only then waits — and + * only for chunks a thread has already claimed. That is what makes the region safe to enter + * from a thread of the dispatcher's own pool: a caller never blocks on work that still needs a + * pool thread, because it does that work itself. A helper that the pool never got round to + * running finds the queue empty and exits without touching the body. + * + * The contract's guarantees hold as before: the first failure stops every chunk that has not + * started, is rethrown once every running chunk has finished (later failures are attached as + * suppressed), no body outlives [forRange], and every chunk's writes happen-before the return. * * A region reached from inside another region (a body that calls into a parallel op despite the - * contract) runs inline: the dispatcher never waits on itself, and a nested `runBlocking` on - * `Dispatchers.Default` cannot starve the pool. Callers that already run on `Dispatchers.Default` - * and want true isolation use [dedicated]. + * contract) runs inline. Callers that want isolation from `Dispatchers.Default` use [dedicated]. + * + * A previous version ran the region as `runBlocking { coroutineScope { launch(dispatcher) … } }`. + * A `coroutineScope` waits for every child, including the ones the pool never started, so a + * region entered from `Dispatchers.Default` deadlocked as soon as every worker was inside one — + * routinely on a 4-vCPU CI runner, never on a 14-core laptop. */ @OptIn(ExperimentalMemoryApi::class) public open class CoroutineSchedule @JvmOverloads constructor( @@ -42,6 +52,8 @@ public open class CoroutineSchedule @JvmOverloads constructor( final override val name: String = "$label($parallelism)" + private val executor: Executor = dispatcher.asExecutor() + override fun forRange(n: Int, grain: Int, body: (start: Int, end: Int) -> Unit) { val tasks = Schedule.tasksFor(n, grain, parallelism) if (tasks == 0) return @@ -51,36 +63,86 @@ public open class CoroutineSchedule @JvmOverloads constructor( } val chunk = Schedule.chunkFor(n, tasks) val started = if (sink.isEnabled) TraceClock.nowNanos() else 0L + val region = Region(n, chunk, tasks, body) inRegion.set(true) try { - runBlocking { - coroutineScope { - var start = chunk - while (start < n) { - val s = start - val e = minOf(start + chunk, n) - launch(dispatcher) { inRegion(s, e, body) } - start = e - } - body(0, minOf(chunk, n)) - } - } + repeat(tasks - 1) { executor.execute(region) } + region.runChunk(0) + region.drain() } finally { inRegion.set(false) } + region.awaitClaimed() + region.rethrow() if (sink.isEnabled) { val now = TraceClock.nowNanos() sink.emit(TraceEvent.ScheduleRegion(op = "forRange", schedule = name, elements = n, tasks = tasks, durationNanos = now - started, timeNanos = now)) } } - private fun inRegion(start: Int, end: Int, body: (Int, Int) -> Unit) { - val previous = inRegion.get() - inRegion.set(true) - try { - body(start, end) - } finally { - inRegion.set(previous) + /** One `forRange` call: the chunk queue, the completion latch and the first failure. */ + private class Region( + private val n: Int, + private val chunk: Int, + private val tasks: Int, + private val body: (Int, Int) -> Unit, + ) : Runnable { + /** Next chunk to claim; the caller takes chunk 0 before the helpers start. */ + private val next = AtomicInteger(1) + + /** Counts down once per chunk, whether it ran or was skipped after a failure. */ + private val remaining = CountDownLatch(tasks) + + private val failure = AtomicReference(null) + + /** Helper entry point: run on a pool thread, claim chunks until none are left. */ + override fun run() { + val previous = inRegion.get() + inRegion.set(true) + try { + drain() + } finally { + inRegion.set(previous) + } + } + + fun drain() { + while (true) { + val i = next.getAndIncrement() + if (i >= tasks) return + runChunk(i) + } + } + + fun runChunk(i: Int) { + try { + if (failure.get() == null) body(i * chunk, minOf((i + 1) * chunk, n)) + } catch (t: Throwable) { + if (!failure.compareAndSet(null, t)) failure.get()!!.addSuppressed(t) + } finally { + remaining.countDown() + } + } + + /** + * Waits for the chunks other threads have claimed. By the time the caller gets here it + * has drained the queue, so every outstanding chunk is on a thread that is running it. + */ + fun awaitClaimed() { + var interrupted = false + while (true) { + try { + remaining.await() + break + } catch (_: InterruptedException) { + interrupted = true + } + } + if (interrupted) Thread.currentThread().interrupt() + } + + fun rethrow() { + failure.get()?.let { throw it } } } @@ -97,7 +159,7 @@ public open class CoroutineSchedule @JvmOverloads constructor( /** * A schedule with its own daemon pool of `parallelism - 1` workers (the caller is the last - * worker), for code that already runs on `Dispatchers.Default`. Close it when done. + * worker), for code that wants isolation from `Dispatchers.Default`. Close it when done. */ @JvmStatic @JvmOverloads diff --git a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt index 995efbb0..9bcc0e49 100644 --- a/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt +++ b/skainet-backends/skainet-backend-cpu/src/jvmTest/kotlin/sk/ainet/exec/schedule/CoroutineScheduleTest.kt @@ -1,6 +1,7 @@ package sk.ainet.exec.schedule import kotlinx.coroutines.Dispatchers +import kotlinx.coroutines.asCoroutineDispatcher import kotlinx.coroutines.launch import kotlinx.coroutines.runBlocking import sk.ainet.context.schedule.Schedule @@ -8,6 +9,9 @@ import sk.ainet.lang.memory.ExperimentalMemoryApi import sk.ainet.lang.memory.trace.RecordingTraceSink import sk.ainet.lang.memory.trace.TraceEvent import java.util.BitSet +import java.util.concurrent.CyclicBarrier +import java.util.concurrent.Executors +import java.util.concurrent.TimeUnit import java.util.concurrent.atomic.AtomicInteger import kotlin.test.Test import kotlin.test.assertEquals @@ -60,16 +64,20 @@ class CoroutineScheduleTest { } @Test - fun callerThreadRunsTheFirstChunkAndWorkersRunTheRest() { + fun callerThreadRunsTheFirstChunkAndWorkersHelpWithTheRest() { val schedule = CoroutineSchedule(parallelism = 4) val caller = Thread.currentThread() val onCaller = AtomicInteger() val elsewhere = AtomicInteger() - schedule.forRange(4000, grain = 1) { _, _ -> + var first: Thread? = null + schedule.forRange(4000, grain = 1) { s, _ -> + if (s == 0) first = Thread.currentThread() if (Thread.currentThread() === caller) onCaller.incrementAndGet() else elsewhere.incrementAndGet() + Thread.sleep(30) // long enough for the pool to claim the other chunks } - assertEquals(1, onCaller.get(), "the caller runs exactly one chunk itself") - assertEquals(3, elsewhere.get(), "the other chunks run on the dispatcher") + assertSame(caller, first, "the caller runs the first chunk itself") + assertTrue(elsewhere.get() >= 1, "at least one chunk ran on the dispatcher") + assertEquals(4, onCaller.get() + elsewhere.get(), "every chunk ran exactly once") } @Test @@ -119,6 +127,33 @@ class CoroutineScheduleTest { assertEquals(4 * 4096, total.get()) } + /** + * The CI deadlock: every thread of the schedule's own pool enters a region at once, so no + * pool thread is free to run the helpers. The caller must finish the region by itself. With + * the old `runBlocking { coroutineScope { … } }` region this hung forever on a pool of any + * size, because the scope waited for children the pool could never start. + */ + @Test + fun aRegionEnteredFromEveryThreadOfItsOwnPoolStillCompletes() { + val poolSize = 2 + val pool = Executors.newFixedThreadPool(poolSize) + try { + val schedule = CoroutineSchedule(dispatcher = pool.asCoroutineDispatcher(), parallelism = 4) + val allInside = CyclicBarrier(poolSize) + val total = AtomicInteger() + val futures = List(poolSize) { + pool.submit { + allInside.await(10, TimeUnit.SECONDS) + schedule.forRange(4096, grain = 1) { s, e -> total.addAndGet(e - s) } + } + } + futures.forEach { it.get(30, TimeUnit.SECONDS) } + assertEquals(poolSize * 4096, total.get()) + } finally { + pool.shutdownNow() + } + } + @Test fun dedicatedScheduleOwnsItsPoolAndCloses() { CoroutineSchedule.dedicated(parallelism = 3).use { schedule ->