fix(ci): serialize heavy shards to fix runner OOM + reshard diffusion (#1454) + Tensors 0.91.2 (#528) - #1485
Conversation
…ffusion by model
The heavy ModelFamily / NN / Diffusion test shards were OOM-killed on the
16 GB CI runners ("The runner has received a shutdown signal"), with tests
passing right up to the kill — a cumulative memory leak, not a hang or a
single oversized model. Confirmed from the #1477 job logs (runner dies 2-6
min in, ~14 GB available at start, fail-fast:false so no cascade).
Two independent contributors, fixed together:
1. Per-shape tensor caches were never cleared between tests. AutoTensorCache
(and the TensorArena persistent pool) pool tensors keyed by SHAPE across
ALL worker threads; over a shard of dozens-to-hundreds of distinct models
they retain multiple GB of weight-sized tensors that GC.Collect cannot
reclaim while the pool holds them. The existing teardowns did compacting
Gen-2 + LOH GC and reset WeightRegistry, but never dropped these pools — so
even MaxParallelThreads=1 + Server GC + GCConserveMemory=9 still accumulated
~160 MB/test net and exhausted RAM after ~80 models. Both
NeuralNetworkModelTestBase and the two diffusion bases now call
TensorCacheSettings.ClearCache() + TensorArena.ClearPersistentPool() +
WeightRegistry.Reset() before their GC pass, returning to a clean baseline
between tests. Clearing reusable caches is safe — it cannot change results
(unlike the reverted weight-streaming attempts, which leaked pool handles).
2. The three ModelFamily Diffusion shards filtered by test-METHOD first letter
(FullyQualifiedName~Tests.A …), which made EACH shard instantiate ALL ~244
diffusion models (just different subsets of each model's methods), so all
three carried the full distinct-model cache footprint. Re-sharded by MODEL
CLASS first letter (~Diffusion.A …) so each shard only loads its letter
range (~1/3 the models) — same three runners, ~3x less cumulative memory.
CI (the real 16 GB runners running the full model sets) is the faithful
validator for the memory envelope; local repro is impractical because the
diffusion training tests run for many minutes serially.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
The latest updates on your projects. Learn more about Vercel for GitHub. 2 Skipped Deployments
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
WalkthroughThis PR fixes issue ooples/AiDotNet#1454 by correcting diffusion model-family shard filters to partition by model-class letter, replaces an ineffective xunit parallelism argument with JSON config rewriting, and bumps AiDotNet package versions. ChangesCI Test Execution Optimization
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related issues
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 2
♻️ Duplicate comments (2)
tests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cs (1)
198-200:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd error handling around cache-clearing calls.
Same issue as in
DiffusionModelTestBase: if any of the cache-clearing methods throw an exception, the subsequentGC.Collect()calls won't run, potentially leaving the test runner in a degraded state.Wrap each clearing call in a try-catch or use a try-finally structure to ensure GC always runs. See the proposed fix in the
DiffusionModelTestBasereview comment for the pattern.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cs` around lines 198 - 200, Wrap each cache-clearing call (AiDotNet.Tensors.TensorCacheSettings.ClearCache, AiDotNet.Tensors.Helpers.TensorArena.ClearPersistentPool, AiDotNet.Tensors.LinearAlgebra.WeightRegistry.Reset) in error-handling so that exceptions do not prevent the subsequent GC.Collect() invocations; follow the DiffusionModelTestBase pattern by surrounding each call with try-catch (logging or swallowing the exception as appropriate) or use a try-finally that guarantees GC.Collect() runs in the finally block, ensuring the test runner isn't left in a degraded state.tests/AiDotNet.Tests/UnitTests/Diffusion/DiffusionUnitTestBase.cs (1)
92-100:⚠️ Potential issue | 🟠 Major | ⚡ Quick winAdd error handling around cache-clearing calls.
Same issue as in the other test base classes: no error handling around the cache-clearing calls. If any of them throw, the GC passes won't run.
The lock correctly protects against concurrent access, but doesn't protect against exceptions from the clearing methods themselves.
Consider wrapping each clearing call in a try-catch to ensure the GC passes run regardless. See the proposed fix in the
DiffusionModelTestBasereview comment for the pattern.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tests/AiDotNet.Tests/UnitTests/Diffusion/DiffusionUnitTestBase.cs` around lines 92 - 100, The three cache-clearing calls (AiDotNet.Tensors.TensorCacheSettings.ClearCache, AiDotNet.Tensors.Helpers.TensorArena.ClearPersistentPool, AiDotNet.Tensors.LinearAlgebra.WeightRegistry.Reset) must be individually wrapped in try-catch blocks so an exception in any one doesn't prevent the subsequent GC/compaction passes; catch Exception (or a more specific exception if known), log the error with the test/logger (or Console) including the method name, and continue execution. Ensure each call has its own try-catch (do not lump them together) so failures are isolated and the GC passes still run.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/AiDotNet.Tests/ModelFamilyTests/Base/DiffusionModelTestBase.cs`:
- Around line 135-145: Wrap the three cache-reset calls
(AiDotNet.Tensors.TensorCacheSettings.ClearCache,
AiDotNet.Tensors.Helpers.TensorArena.ClearPersistentPool,
AiDotNet.Tensors.LinearAlgebra.WeightRegistry.Reset) in a defensive try/finally
(or try/catch and then finally) so any exception from those methods is caught
and logged but does not prevent the subsequent Gen-2/LOH GC compaction from
running; ensure the finally block always invokes the existing GC compaction
sequence so compaction cannot be skipped on failures.
In `@tests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cs`:
- Around line 189-200: The process-wide cache clears
(AiDotNet.Tensors.TensorCacheSettings.ClearCache,
AiDotNet.Tensors.Helpers.TensorArena.ClearPersistentPool,
AiDotNet.Tensors.LinearAlgebra.WeightRegistry.Reset) must be protected by the
same static lock used in sibling bases to avoid races; wrap these three calls
inside lock(_lohCompactionGate) (or declare a private static readonly object
_lohCompactionGate if not present) in DisposeAsync of NeuralNetworkModelTestBase
so concurrent test-class disposals cannot run the clears simultaneously.
---
Duplicate comments:
In `@tests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cs`:
- Around line 198-200: Wrap each cache-clearing call
(AiDotNet.Tensors.TensorCacheSettings.ClearCache,
AiDotNet.Tensors.Helpers.TensorArena.ClearPersistentPool,
AiDotNet.Tensors.LinearAlgebra.WeightRegistry.Reset) in error-handling so that
exceptions do not prevent the subsequent GC.Collect() invocations; follow the
DiffusionModelTestBase pattern by surrounding each call with try-catch (logging
or swallowing the exception as appropriate) or use a try-finally that guarantees
GC.Collect() runs in the finally block, ensuring the test runner isn't left in a
degraded state.
In `@tests/AiDotNet.Tests/UnitTests/Diffusion/DiffusionUnitTestBase.cs`:
- Around line 92-100: The three cache-clearing calls
(AiDotNet.Tensors.TensorCacheSettings.ClearCache,
AiDotNet.Tensors.Helpers.TensorArena.ClearPersistentPool,
AiDotNet.Tensors.LinearAlgebra.WeightRegistry.Reset) must be individually
wrapped in try-catch blocks so an exception in any one doesn't prevent the
subsequent GC/compaction passes; catch Exception (or a more specific exception
if known), log the error with the test/logger (or Console) including the method
name, and continue execution. Ensure each call has its own try-catch (do not
lump them together) so failures are isolated and the GC passes still run.
🪄 Autofix (Beta)
✅ Autofix completed
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: a94e0e8f-aec4-40d2-9d6d-89feeb50a5ba
📒 Files selected for processing (4)
.github/workflows/sonarcloud.ymltests/AiDotNet.Tests/ModelFamilyTests/Base/DiffusionModelTestBase.cstests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cstests/AiDotNet.Tests/UnitTests/Diffusion/DiffusionUnitTestBase.cs
…v/reduction perf) Pulls in AiDotNet.Tensors#528 (allocation-free axis reductions + small-batch channel-parallel conv-backward; ResNet50 fp64 train step 2553ms -> ~1150ms, bit-identical). This is the COMPUTE half of #1463 — together with the memory/OOM cache-clearing + diffusion re-shard already in this PR, #1485 now resolves both halves of the paper-scale CNN/diffusion CI-shard failures, so merging it closes #1463. Native packages bumped in lockstep to 0.91.2. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
|
Note Autofix is a beta feature. Expect some limitations and changes as we gather feedback and continue to improve it. Fixes Applied SuccessfullyFixed 2 file(s) based on 2 unresolved review comments. Files modified:
Commit: The changes have been pushed to the Time taken: |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/AiDotNet.Tests/ModelFamilyTests/Base/DiffusionModelTestBase.cs`:
- Around line 149-153: The catch block in DiffusionModelTestBase that currently
swallows Exception contradicts its comment "log and continue" and risks hiding
teardown failures referenced by NeuralNetworkBase.ResetWeightStreamingForTests;
update the catch to write diagnostic output (e.g.,
System.Diagnostics.Debug.WriteLine or the test logger) including the exception
message/stack and context ("WeightRegistry.Reset failed during test teardown"),
or if swallowing is intentional, change the comment to explicitly state
"silently ignore" and add a TODO referencing WeightRegistry.Reset() so
maintainers know the trade-off; ensure the change targets the catch in
DiffusionModelTestBase and respects the upstream warning about not silently
swallowing WeightRegistry.Reset() failures.
In `@tests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cs`:
- Around line 198-215: NeuralNetworkModelTestBase.DisposeAsync currently calls
AiDotNet.Tensors.TensorCacheSettings.ClearCache,
AiDotNet.Tensors.Helpers.TensorArena.ClearPersistentPool, and
AiDotNet.Tensors.LinearAlgebra.WeightRegistry.Reset inside a lock without
protection, so if any of those throws the
GC.Collect/WaitForPendingFinalizers/GC.Collect calls are skipped; wrap the three
clear/reset calls in a try block and move the GC.Collect();
GC.WaitForPendingFinalizers(); GC.Collect(); into a finally block (keeping the
lock) so gen-2/LOH compaction always runs, mirroring the pattern used in
DiffusionModelTestBase.DisposeAsync.
🪄 Autofix (Beta)
✅ Autofix completed
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: d496a611-03b0-4b99-bcac-97fa6b0884e0
📒 Files selected for processing (2)
tests/AiDotNet.Tests/ModelFamilyTests/Base/DiffusionModelTestBase.cstests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cs
|
Note Autofix is a beta feature. Expect some limitations and changes as we gather feedback and continue to improve it. Fixes Applied SuccessfullyFixed 2 file(s) based on 2 unresolved review comments. Files modified:
Commit: The changes have been pushed to the Time taken: |
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tests/AiDotNet.Tests/ModelFamilyTests/Base/DiffusionModelTestBase.cs`:
- Around line 149-155: The catch message incorrectly blames WeightRegistry.Reset
even though the exception could come from ClearCache(),
TensorArena.ClearPersistentPool(), or WeightRegistry.Reset; update the teardown
to either wrap each call (ClearCache, ClearPersistentPool, Reset) in its own
try/catch so the logged message names the failing method, or include contextual
information about which call threw by catching exceptions immediately around
each call and logging the method name plus exception details (message and stack)
instead of the fixed "WeightRegistry.Reset failed..." text.
In `@tests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cs`:
- Around line 200-220: Add a catch(Exception) between the try and finally in
NeuralNetworkModelTestBase.DisposeAsync around the calls to
AiDotNet.Tensors.TensorCacheSettings.ClearCache(),
AiDotNet.Tensors.Helpers.TensorArena.ClearPersistentPool(), and
AiDotNet.Tensors.LinearAlgebra.WeightRegistry.Reset() so exceptions during
cache/registry clearing are caught, logged (using the same logging used in
DiffusionModelTestBase.DisposeAsync) and swallowed before reaching the finally
block that runs GC.Collect()/GC.WaitForPendingFinalizers(), ensuring teardown
continues even if Reset() or the clear calls fail.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 4d36fc81-87a8-4fe6-b49a-305159b223bd
📒 Files selected for processing (2)
tests/AiDotNet.Tests/ModelFamilyTests/Base/DiffusionModelTestBase.cstests/AiDotNet.Tests/ModelFamilyTests/Base/NeuralNetworkModelTestBase.cs
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In @.github/workflows/sonarcloud.yml:
- Around line 553-565: The current heavy-shard serialization block (guarded by
$serializeShard) only warns when $runnerJson (xunit.runner.json) is missing and
allows the shard to continue; change this so a missing runner config fails the
job: in the block that computes $projDir and $runnerJson and checks Test-Path,
replace the Write-Host warning path with a hard failure (e.g., Write-Error and
exit 1 or throw) so the shard stops rather than running without serialization;
ensure the error message clearly references $runnerJson and that this applies
only inside the $serializeShard branch.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: fb462d69-e1d6-4b97-ada0-941b902efa35
📒 Files selected for processing (1)
.github/workflows/sonarcloud.yml
…tion perf) Pulls in AiDotNet.Tensors#528 (allocation-free axis reductions + small-batch channel-parallel conv-backward; ResNet50 fp64 train step 2553ms -> ~1150ms, bit-identical). This is the COMPUTE half of #1463 — together with the memory/OOM cache-clearing + diffusion re-shard already in this PR, #1485 now resolves both halves of the paper-scale CNN/diffusion CI-shard failures, so merging it closes #1463. Native packages bumped in lockstep to 0.91.2. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Fixed 2 file(s) based on 2 unresolved review comments. Co-authored-by: CodeRabbit <noreply@coderabbit.ai>
Fixed 2 file(s) based on 2 unresolved review comments. Co-authored-by: CodeRabbit <noreply@coderabbit.ai>
A/B measurement (NN ModelFamily, MaxParallelThreads=1, in-process working-set trace after each teardown) shows the cache-clearing does NOT bound memory: teardown: 50 100 200 300 peak no clear: 3328MB 5879MB 10650MB 18067MB 55091MB with clear: 3885MB 10691MB 19294MB 26073MB 55709MB Both climb identically (~94 MB per test-method, linear) to ~55 GB. Clearing AutoTensorCache / TensorArena / WeightRegistry frees the per-shape pools but those are NOT the dominant grower — the leak is a rooted *managed* per-model- instance accumulation (GetTotalMemory climbs too) that survives Dispose + compacting Gen-2 GC + the cache clears. Removing the no-op clearing (and its misleading comments); the real per-instance leak needs a heap-dump root-cause hunt, tracked separately. The diffusion shard re-sharding (#1454) and the Tensors 0.91.2 bump in this PR are unaffected and stand on their own. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The `-- xunit.MaxParallelThreads=1` inline arg is silently ignored by the xunit.runner.visualstudio (VSTest) adapter; only xunit.runner.json in the build output is honored. Heavy shards therefore still ran test classes in parallel (maxParallelThreads = ProcessorCount), each holding a multi-GB model + Adam state for the whole class, peaking ~55 GB on a 16 GB runner and triggering the OOM shutdown signal. Rewrite the built xunit.runner.json to parallelizeTestCollections=false + maxParallelThreads=1 for the heavy shards only, immediately before dotnet test. Locally confirmed: serial trajectory peaks ~4.5 GB (one model in flight) and oscillates ~550 MB with no cumulative leak, vs ~55 GB under the parallel default. Every other shard keeps the parallel JSON default and runs full-speed. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
607e1bb to
67dc004
Compare
…t on missing config
Serialization alone did not fix the heavy-shard runner OOM. CI logs show
`parallel test collections = off` took effect, yet the shard still OOM'd ~70s
into serial execution — so the remaining cause is footprint, not parallelism.
Two contributors: the global Server GC (DOTNET_gcServer=1, chosen for
parallel-collection throughput) reserves a heap segment per core and collects
lazily, and the assembly discovers ~64k test cases. For the now-serial heavy
shards:
- override to Workstation GC (DOTNET_gcServer=0): smaller footprint, eager
collection under memory pressure;
- disable theory pre-enumeration (preEnumerateTheories=false) so discovery
doesn't materialize every theory's data rows up front;
- fail fast (exit 1) when the runner.json rewrite can't be applied instead of
silently falling back to a parallel run that OOMs (addresses CodeRabbit).
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Closes #1454
Part of #1462
Part of #1463
What this PR does (three independent, verified changes)
1. Fix the heavy-shard runner OOM — serialize via
xunit.runner.jsonRoot cause (definitive — heap-dump + A/B confirmed): the heavy shards exhaust the 16 GB runner because xUnit runs many test classes in parallel (
maxParallelThreadsdefaults toProcessorCount), and each big-model class holds a multi-GB model + Adam state for the whole class lifetime. Dozens in flight → ~55 GB peak →runner has received a shutdown signal.The previous mitigation passed
-- xunit.MaxParallelThreads=1as a VSTest inline arg. That arg is silently ignored by thexunit.runner.visualstudioadapter — onlyxunit.runner.jsonin the build output is honored. So the shards we thought were serialized were in fact still fully parallel.Fix (for the heavy shards only):
xunit.runner.jsontoparallelizeTestCollections=false+maxParallelThreads=1immediately beforedotnet test.DOTNET_gcServer=0). The CI run confirmed serialization took effect (parallel test collections = offin the xUnit banner) yet the shard still OOM'd ~70 s into serial execution — so the remaining cause is footprint, not parallelism. The global Server GC (chosen for parallel-collection throughput) reserves a heap segment per core and collects lazily; with the shard now serial it isn't needed, and Workstation GC has a much smaller footprint and collects eagerly under pressure.preEnumerateTheories=false) so discovery doesn't materialize every theory's data rows up front (the assembly discovers ~64k cases).Every other shard keeps the parallel JSON default + Server GC and runs full-speed. (My earlier local A/B showed serial peaking ~4.5 GB, but that ran under Workstation GC and without coverage instrumentation — it did not reproduce CI's Server-GC + coverage footprint, which is why step 1 alone was insufficient. CI is the validator for the combined fix.)
2. Re-shard the diffusion ModelFamily tests by model class — fixes #1454
The three
ModelFamily - Diffusion {A-I,J-R,S-Z}shards filtered by test-method first letter (FullyQualifiedName~Tests.A …), which — becauseTests.M/Tests.Dcollide with theModelFamilyTests/Diffusionnamespace segments — made every shard run the full ~244-model suite (≈3× wasted compute, meaningless per-shard attribution). Switched to model-class sharding (~Diffusion.A …) so each shard only loads its letter range. Same three runners.3. Bump
AiDotNet.Tensors0.91.1 → 0.91.2 (+ Native packages, lockstep) — consumes #528Pulls in AiDotNet.Tensors#528 (allocation-free axis reductions + small-batch channel-parallel conv-backward; ResNet50 fp64 train step 2553 ms → ~1150 ms, ~2.2×, bit-identical). This is the per-step compute half of #1463.
On the earlier "per-instance leak" conclusion (corrected)
My first read of the A/B table below concluded the OOM was a rooted per-instance managed leak, because both arms climbed identically to ~55 GB even at "
MaxParallelThreads=1". That conclusion was wrong: the=1arm never took effect (ignored inline arg — see §1), so both arms ran fully parallel. The 55 GB is unbounded parallelism, not a per-instance leak. Under genuinely serial execution (xunit.runner.jsonrewrite) the trajectory is bounded at ~4.5 GB with no cumulative growth.The per-shape tensor-cache clearing I initially added was still correctly reverted — the A/B (below) showed it made no difference (it was a no-op with misleading comments), and serialization is the real fix.
(Both arms ran parallel — the
MaxParallelThreads=1flag was ignored — which is precisely why neither was bounded.)What remains open
The per-step compute timeouts in #1462/#1463 (e.g.
VGGNetworkTests.MoreData_ShouldNotDegrade/ResNetNetworkTests.MoreData_ShouldNotDegradeat the 120 s per-test cap, paper-scale fp64) are only partially addressed by #528's ~2.2× and may still exceed 120 s. Those stay open and are tracked there — not weakened here. This PR fixes the runner-OOM axis (concurrent model loads) and the #1454 shard mis-filtering, and lands the #528 per-step speedup.🤖 Generated with Claude Code
Summary by CodeRabbit