Repository navigation
Add VanillaSort external sorter - #4832
Open
KohakuKirisame wants to merge 6 commits into
Open
KohakuKirisame wants to merge 6 commits into
KohakuKirisame wants to merge 6 commits into
Conversation
Register VanillaSortSorter and map model selection, local checkpoint paths, CPU/CUDA devices, seed, mixture count, profile and batch sizes to vanillasort.sort(). Keep all detection and clustering code in the external package. Use BaseSorter's recording serialization and metadata flow, preserve multi-segment results and main_channel_id with native sorting-folder storage, and capture package progress in the standard runtime log. Report package installation/version and optional GPU support without loading checkpoints during discovery. Reuse SorterCommonTestSuite with a short synthetic recording and add one direct-versus-wrapper integration test. Check exact spike-vector/unit/property equivalence, result reload, recording metadata, version reporting, output cleanup and standard failure logging. Validation: the new external tests and existing sorter registry/dependency checks pass (9 passed, 1 platform-specific skip). Black, targeted Ruff, compileall and git diff --check pass. The optional VanillaSort package is tested from its local 0.1.0 inference-package branch.
Add VanillaSort to the external sorter list and document source installation, the public package API, repository and paper citation, and a basic run_sorter example. Describe 30 kHz and 2D geometry requirements, four-channel local neighborhoods, optional CPU/CUDA execution, configurable GMM K and local models. Explicitly identify experimental larger-probe behavior, independent segment units and full-segment memory requirements. Validation: documentation changes are limited to the existing sorter list and installation guide; whitespace and diff checks pass. Submission and upstream discussion remain for human contributors under AGENTS.md.
Document that default pretrained models are downloaded lazily from a pinned Hugging Face release and cached separately from the Python package. Update the model parameter description and link the published model repository. Generate small local VanillaDet/HuiduRep checkpoints for the existing common-suite and equivalence tests. Exercise the real inference path and local-model parameter without downloading pretrained weights or adding test cases. Validation: HF_HUB_OFFLINE=1 run of the external sorter, registry and runner dependency tests reports 9 passed and 1 platform-specific skip. A separate smoke test with the actual Hub model and installed lightweight wheel confirms exact direct/wrapper equivalence. Black, targeted Ruff and git diff --check pass.
Replace the repository clone and editable-install commands with pip install vanillasort in the external sorter's installation message and the sorter installation guide. Retain the minimum supported version and the link to platform-specific PyTorch guidance. Validation: Black and git diff --check pass. A targeted source/documentation search confirms both VanillaSort installation instructions now use the PyPI command. This changes instructional text only.
zm711
reviewed
Oct 12, 2026
zm711
left a comment
zm711
left a comment
Member
There was a problem hiding this comment.
Few comments to improve performance of our package overall as we add the features. Let me know if you need more guidance.
| if recording is None: | ||
| raise RuntimeError("Unable to load the recording for VanillaSort") | ||
| # BaseSorter includes this trace in spikeinterface_log.json, including on failure. | ||
| with (sorter_output_folder / "vanillasort.log").open("w", encoding="utf-8") as log, redirect_stdout(log): |
Member
There was a problem hiding this comment.
walk me through this in brief. I haven't used redirect_stdout in the past.
Author
There was a problem hiding this comment.
VanillaSort reports progress through print(). redirect_stdout(log) temporarily sends that output to vanillasort.log, which BaseSorter then includes in spikeinterface_log.json under runtime_trace.
When the block exits, stdout is restored and the file is closed automatically—even if sorting fails. This captures stdout only; BaseSorter records exceptions separately. Since redirection affects the whole process, the wrapper declares threaded execution unsupported.
Address the maintainer review by checking module availability with importlib.util.find_spec and reading the installed distribution version with importlib.metadata.version. Preserve the minimum supported version without importing the optional VanillaSort package. Make use_gpu report the requested device intent without importing Torch on the host. Use the recognized nvidia-optional capability so SpikeInterface's container runner handles hardware availability, while VanillaSort resolves auto in the inference environment. Explain how redirect_stdout captures package progress in the standard runtime trace and restores stdout after success or failure. Add one focused import-guard regression test and extend the existing integration test to verify stdout restoration. Validation: the VanillaSort external suite and existing registry/dependency checks pass with HF_HUB_OFFLINE=1 (10 passed, 1 platform-specific skip) against the published vanillasort 0.1.0 package. A fresh-process check confirms discovery, version reporting and GPU intent leave both vanillasort and torch unloaded. Black, targeted Ruff, compileall and git diff --check pass.
Member
|
I tested this on my laptop with a small GPU, but it seems to get stuck. I only see this output ofter more than one hour. Any idea what could be going wrong? I see my GPU being used... |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add VanillaSort, a spike sorter combining VanillaDet detection, HuiduRep embeddings, relative-amplitude features, GMM clustering and template refinement.
The wrapper calls the installable VanillaSort 0.1.0 public API, keeping algorithm code in its own package. It integrates recording serialization, parameter mapping, version reporting, standard logs and reloadable sorting results. Pretrained weights are downloaded lazily from a pinned Hugging Face revision.
Validation includes SorterCommonTestSuite, offline tiny-model tests, exact direct/wrapper equivalence and pretrained-checkpoint regression checks.
See below:
Repository · Models · Paper