Cooperative cancellation for training and prediction - #66
Draft
hinderling wants to merge 12 commits into
Draft
Conversation
Train, predict, and stack-predict now run on a napari thread_worker so the UI stays responsive. The train/segment/segment-stack buttons flip to 'Cancel' while running and abort the operation when clicked again. Cancellation is cooperative: a CancelToken is threaded through ConvpaintModel.train / segment / predict_probas / get_feature_image and through FeatureExtractor.extract_features*, with checkpoints at the per-scale, per-Z-plane, per-tile, and per-layer (VGG16 forward hook) boundaries. CatBoost / RandomForest fit is uninterruptible C code, so cancel always fires before or after the fit, never during it. Memory-mode annotation state (annot_dict + table) is snapshotted at the start of train and rolled back on CancelledError, so a cancelled train can be retried without hitting 'No new annotations'. The same cancel_token kwarg is available to API users; CancelToken and CancelledError are exposed at the package top level.
napari.qt.threading.thread_worker calls window._register_task_status with cancel_callback=worker.quit, and napari's TaskStatusManager never unregisters tasks — so each worker (and its closure, which captures cp_model with the VGG16 MPS weights) is pinned for the viewer's lifetime. In CI's macOS job, ~16 widget tests each retained a VGG16 on MPS, busting the 7.93 GiB runner cap on the last few tests. Switch to superqt.utils.thread_worker (the same WorkerBase class napari wraps, but without the task-status registration) and recreate the progress bars manually with napari.utils.progress inside each worker body. Verified: repeated widget create/train/destroy on MPS now stays flat at one widget's worth of weights (~573 MB) instead of growing ~593 MB per iteration.
The previous leak fix moved napari.utils.progress construction inside the worker body. In sync mode that still runs on the main thread so tests passed, but an async worker actually runs on a Qt thread-pool thread and Cocoa raises NSInternalInconsistencyException when any QWidget is instantiated off the main thread — so interactive train crashed the napari app immediately on macOS. Reinstate the pattern napari's own thread_worker decorator used: create the progress bar in _begin_worker (which is always invoked on the main thread) and wire it to the worker's finished / yielded signals. This keeps the leak fix intact (no _register_task_status call) while making Qt happy. Also add test_widget_async_worker_completes_without_main_thread_violation, which flips _sync_workers off for one case and drives a full async _on_train round. The rest of the widget tests run sync and would have hidden this regression.
Main branched forward with a widget lazy-init refactor (convpaint_widget moved from eager to QTimer/showEvent-triggered model construction, and several modules now inline their torch/.utils imports) and an independent MPS-cleanup fixture in conftest. Resolutions: - conftest.py: keep both the _sync_workers=True setup (this branch) and the gc+torch.mps.empty_cache autouse fixture (main). They compose — the root-cause fix is still in convpaint_widget (superqt workers, no task-status registration), the fixture is a per-test belt-and-suspenders. - convpaint_widget.py: drop the now-unneeded .utils/ConvpaintModel top-level imports to match main's startup-accelerated pattern; inline-import CancelToken/CancelledError inside the three slot methods that need them. _ActiveOp's cancel_token field annotated as object rather than a forward ref so the dataclass decoration doesn't try to resolve it at import time. - feature_extractors/nnlayers.py: keep check_cancel import, adopt main's import_models() lazy torchvision pattern. - _tests/test_cancellation.py: testing_utils -> testing_data rename, and add widget.ensure_init() to the new async regression test to trigger main's deferred _late_init before we touch cp_model. 11/11 cancellation tests still green after the merge.
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #66 +/- ##
==========================================
- Coverage 63.20% 60.67% -2.54%
==========================================
Files 32 33 +1
Lines 6515 7046 +531
==========================================
+ Hits 4118 4275 +157
- Misses 2397 2771 +374 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
hinderling
marked this pull request as draft
April 22, 2026 15:59
…tures - utils.cancel_scope() installs the token; check_cancel() reads it ambiently. - cancel_token kwarg stays only on public entry points (train, segment, predict_probas, get_feature_image, FeatureExtractor.extract_features) and on _train/_predict used by the widget; all FE-method signatures revert to their pre-cancellation form, so third-party FE subclasses work unmodified and are still cancellable through the base-class loop checkpoints. - Remove Hookmodel._cancel_token instance attribute (hooks use the ambient token; also removes the shared-instance clobbering between threads). - Dask tiled prediction: try/finally cancels pending futures and closes the client on any abort (a cancel previously leaked the whole local cluster); cancel checkpoint added to the gather loop; the token is no longer pickled into dask tasks (it never propagated across processes anyway). - Document cancel_token and CancelledError on the four public methods. - New tests: old-signature custom FE works and cancels; ambient token does not leak out of a cancelled call.
- _on_get_feature_image, _on_get_feature_image_all (kmeans: whole-stack via returned; otherwise per-slice generator with progress) and _on_train_on_project now run on cancellable thread workers like train/predict/predict_all. For project training, the data-collection loop stays on the main thread (the file list drives Qt layer loading); only model.train() runs on the worker. - Op buttons are mutually exclusive while one runs (_other_op_buttons); previously 'Get features' could be clicked mid-train and race the shared FE. - Snapshot annot.data with .copy(): the UI is responsive during training, so the user could paint into the array the worker is reading. - Capture the target slice at click time in _on_predict and _on_get_feature_image; scrolling during the async run no longer writes the result to the wrong slice.
Conflict resolutions and adaptations: - _train gains upstream's sort_features param alongside cancel_token; sort block lives in _train_body. - _parallel_predict_image: upstream's plot_tiles + tile-alignment margins combined with the cancellation try/finally dask cleanup. - Widget renames applied in the worker slots (seg_tag, _approve_annotations_layer_shape, _get_data_dims(data, ndims)). - _on_train_on_project dropped (removed upstream); its cancellable-worker treatment moves to the new shared _train_multiple core, covering both Train-on-selected and Multifile training. - _on_segment_selected_multifile converted to a cancellable generator worker (predict + TIFF writing on the worker; store/tick updates on the main thread via yielded). - _other_op_buttons updated to the new button set (multifile train/preview/ segment-selected, train-on-selected); _on_predict shows Cancel on whichever of its two trigger buttons (Segment / Multifile preview) was clicked.
A cancel click can only take effect at the next cooperative checkpoint — an in-flight CatBoost/RF fit or a single FE forward pass is uninterruptible — so the button now flips to 'Cancelling…' (disabled) the moment cancel is requested, and the progress-bar description follows. When the worker finally finishes, the button is restored and the 'Operation cancelled.' notification is shown as before. Repeat clicks during the grace period are ignored.
- CatBoost fit (the single longest blocking call) is now cancellable on CPU via a per-iteration callback that watches the ambient cancel token; the early-stopped partial fit is discarded by a post-fit checkpoint. GPU fits stay uninterruptible (CatBoost doesn't support callbacks there). - self.classifier is only reassigned after a successful fit — previously a failed or cancelled fit clobbered the previous classifier with an unfitted object (also fixes retrained-state corruption on fit exceptions). - _clf_predict now predicts in 1M-row chunks with a checkpoint between chunks: a whole-plane predict_proba was one multi-second C call on large images; chunking bounds cancel latency at bit-identical results. - Dask gather polls future.result(timeout=1) instead of blocking indefinitely, so cancel is honored while tiles are still computing. - RandomForest fit remains uninterruptible (sklearn has no fit callbacks; warm-start chunking would change the seeded forest). - New test: cancel landing mid-CatBoost-fit aborts, preserves the previous classifier, and the model stays retrainable.
- Progress-bar description while cancelling drops the ellipsis (napari appends ': ' after it, which rendered as 'Cancelling…:'). - The pre-training check with continuous training enabled led with 'Model has not yet been trained', hiding the actual action item; both messages now state that annotations of at least 2 classes are required.
…eshape) Cancelling during prediction on a many-channel image could hang in 'Cancelling' for ~10s: the per-series forward passes are checkpointed, but the per-scale feature rescaling, tensor→numpy transfers and the final concatenation were one unbroken stretch — with e.g. a 20-channel 1024² input that post-processing works on a multi-GB feature image. Rescaling and device transfers now check the token between per-series blocks, and checkpoints precede the big concatenation and prediction reshaping.
The worker object is never read back (superqt keeps running workers referenced for their lifetime); the op record only needs the token, button state and progress bar.
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
2 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.
Summary
thread_worker, keeping the UI responsive. The launching button flips toCancelwhile running; a second click flips it toCancelling…(disabled) until the worker stops at its next checkpoint. Op buttons are mutually exclusive while one is running.CancelToken, carried in a ContextVar (utils.cancel_scope/check_cancel()) rather than through method signatures: only the public entry points (ConvpaintModel.train / segment / predict_probas / get_feature_image,FeatureExtractor.extract_features) take acancel_token=kwarg. Custom FeatureExtractor subclasses need no signature changes and are automatically cancellable through the base-class checkpoints (per-scale, per-Z-plane, per-channel-series, per-tile, per-hooked-layer, per-rescale-block, per-predict-chunk).self.classifieris now only reassigned after a successful fit — previously a failed fit clobbered the previous training with an unfitted object. GPU CatBoost and RandomForest fits remain uninterruptible (no callback support).predict_probaruns in 1M-row chunks (bit-identical results, bounded cancel latency), and feature post-processing (rescaling, tensor→numpy transfer) checks the token between per-series blocks — this mattered in practice: cancelling during a many-channel extraction used to hang ~10 s in post-processing.try/finally— a cancel previously leaked the whole local cluster.annot_dict,self.table) is snapshotted at the start of_trainand rolled back onCancelledError, so a cancelled train can simply be retried.CancelToken/CancelledErrorare exported at the package level, andcancel_token/CancelledErrorare documented on the public methods. API calls without a token are unaffected (every checkpoint no-ops).What changed
utils.py:CancelToken,CancelledError,cancel_scope()(ContextVar),check_cancel().convpaint_model.py:cancel_token=Noneon the public API;_trainwraps the body incancel_scopewith snapshot/rollback of memory-mode state; cancellable CatBoost fit; chunked_clf_predict; dask cleanup.feature_extractor.py,feature_extractors/nnlayers.py: checkpoints in the base-class loops and VGG16 forward hooks read the ambient token — FE subclass signatures are unchanged frommain.convpaint_widget.py: all long-running slots (_on_train,_on_predict,_on_predict_all,_on_get_feature_image[_all],_train_multipleused by Train-on-selected + Multifile training,_on_segment_selected_multifile) wrapped in cancellable workers;_ActiveOpconsolidates running-op state; Cancel/Cancelling… button handling. Tests set_sync_workers=Trueso existing synchronous assertions keep working._tests/test_cancellation.py: 15 tests covering cancel-before-start, cancel-mid-run, cross-thread cancel, cancel mid-CatBoost-fit (classifier preserved, retrainable), memory-mode rollback, a custom FE with pre-cancellation signatures (works unmodified + cancellable), ambient-token leak-out, and a full async widget round-trip on a real Qt worker thread.Test plan
pytest src/napari_convpaint/_tests/— 111 passed, 2 failed: the twocellpose_backboneparams oftest_all_models_train_predict, which fail because cellpose is not installed in the test env (identical onmain, unrelated to this branch)check_cancel()pathmain(Multifile rework, DINOv3, v1.1.0) is merged into the branch; the new Multifile train/segment operations got the same worker + cancellation treatment.