stream crashes on first prune: SharedPool.prune() got an unexpected keyword argument 'survival' #1
Description
stream crashes on its first prune — SharedPool.prune() signature mismatch
stream is unusable end-to-end: it trains correctly for a while, then raises TypeError the first time the pruner runs. read is unaffected (details at the bottom), so this looks isolated to the SharedPool path.
Environment
- mini-AGI @
0cc950b(currentorigin/mainas of 2026-09-21) - macOS 26.6, Apple M4 Pro, 64 GB — training via MPS
- Python 3.14, torch 2.14.0, numpy
Reproduction
python3 train.py --device mps stream \
--mix data_math_char:1.0 \
--weights-dir smoke/weights --out smoke/runs \
--steps 40 --chunk 512 --context 2048 --ctx-start 1024 \
--log-every 10
Runs fine to the end of the step budget (loss 5.74 → 4.72 over 40 steps, val 4.26), then crashes. Reproduced on two clean --weights-dirs; exit code 1 both times. Deterministic — it fires at the first prune, i.e. immediately after the first validation once a grower is armed.
Traceback
File "train.py", line 592, in cmd_stream
gone = model.pool.prune(step, survival=args.grow_min_age)
File "torch/utils/_contextlib.py", line 124, in decorate_context
return func(*args, **kwargs)
TypeError: SharedPool.prune() got an unexpected keyword argument 'survival'
Cause
The repo has two pool implementations whose prune() signatures have diverged, and the callers all use the newer one:
| Location | Signature |
|---|---|
minagi/pool.py:208 — SharedPool.prune |
prune(self, step, min_gate=0.005, min_age=8000, min_delta=5e-4, protect=0) |
minagi/paged.py:900 — PagedPool.prune |
prune(self, step, survival=8600, protect=0) |
Three call sites all pass survival=:
train.py:376—model.pool.prune(step, survival=args.grow_min_age)train.py:592— same (this is the one that fires incmd_stream)train.py:1508—pool.prune(step, survival=survival_steps)(paged pool, works)
SharedPool appears to have kept the older min_gate/min_age/min_delta API while the callers moved to survival. Both pool.py and train.py were last touched by the same commit (0cc950b), so there is no earlier revision that ran clean.
Note this is a plain Python argument mismatch — independent of device or torch version. It would fire identically on CUDA.
Suggested direction
I don't know which side is authoritative here, so this is a question rather than a patch suggestion:
- align
SharedPool.prunewithsurvival(and decide ifmin_gateetc. are still wanted), or - point
streamat the paged pool, ifreadis now the only supported lane andstreamis legacy.
If stream is legacy, it may be worth updating its help text — it currently reads "batch 1, cached, chunked - the way it runs", which is what led me here first.
read is fine (for contrast)
For completeness, the read lane works end-to-end on the same machine: it built a fresh model (64 experts × 3,145,728 params, 2048 hidden, depth 1 — pool 201.33M, resident 32, 838.5 MB on disk), read and trained for its --minutes budget, and reported read 0.06M characters in 1.0m (0.9k char/s) on MPS. No crash. So the defect really does look confined to stream.
Unrelated question: license
There's no LICENSE file and the GitHub API reports license: None. Is one planned?
Asking because I have a small, self-contained patch I'd like to offer — --device currently defaults to cuda if torch.cuda.is_available() else "cpu", so on Apple Silicon it silently falls back to CPU and MPS is only reachable by passing --device mps explicitly. It worked well here, so it seems worth auto-detecting. I'd rather wait for clear terms before opening a PR.
Activity
#1 + License
volotat commentedon Sep 21, 2026
Thanks for the report. I completely forgot about the license. Added it now. MIT, as usual. The fix is in the place and everything should be fine now.