Skip to content

Commit b4667ad

Browse files
author
Donglai Wei
committed
Refactor v2 pipeline boundaries
1 parent afdc0a7 commit b4667ad

61 files changed

Lines changed: 2439 additions & 670 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎.claude/refactor/config_v2.md‎

Lines changed: 124 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,124 @@
1+
# Config V2 Plan
2+
3+
## Goal
4+
5+
Make config strict, typed, and minimal. V2 should reject old field names instead
6+
of translating them, and every runtime option should map to one schema field
7+
with one meaning.
8+
9+
Baseline: `config.md` reports a mature stage-aware Hydra/OmegaConf system. V2
10+
keeps that architecture but removes remaining compatibility tolerance.
11+
12+
## Public API
13+
14+
Keep a small public API under `connectomics.config`:
15+
16+
- `Config`
17+
- stage dataclasses: `DefaultConfig`, `TrainConfig`, `TestConfig`, `TuneConfig`
18+
- section dataclasses for `system`, `model`, `data`, `optimization`,
19+
`monitor`, `inference`, `decoding`, and `evaluation`
20+
- `load_config`
21+
- `save_config`
22+
- `validate_config`
23+
- `resolve_default_profiles`
24+
- `as_plain_dict`
25+
- `cfg_get`
26+
27+
Hardware planning helpers are public only under `connectomics.config.hardware`.
28+
All other helpers are private to `config.pipeline`, `config.schema`, or
29+
`config.hardware`.
30+
31+
## Delete
32+
33+
- Any old `hydra_config` or `hydra_utils` facade if reintroduced.
34+
- Any alias fields kept for old tutorials.
35+
- Any `shared` stage handling.
36+
- Any nested `inference.decoding` or `inference.evaluation` fields.
37+
- Any profile selector accepted outside canonical paths.
38+
- Any runtime fallback that probes a removed field with `getattr`.
39+
- Any scheduler option duplicated both as a top-level field and inside `params`
40+
unless it is genuinely shared by multiple schedulers.
41+
42+
## Move/Rename
43+
44+
- Top-level stage sections should be canonical:
45+
- `default.decoding`
46+
- `test.decoding`
47+
- `tune.decoding`
48+
- `default.evaluation`
49+
- `test.evaluation`
50+
- `tune.evaluation`
51+
- Keep hardware helpers under `config.hardware`.
52+
- Keep schema-only dataclasses under `config.schema`.
53+
- Keep load, merge, profile, and stage-resolution logic under
54+
`config.pipeline`.
55+
56+
## Config Contract
57+
58+
Stage allowlists:
59+
60+
| Stage | Allowed sections |
61+
| --- | --- |
62+
| `default` | `system`, `model`, `data`, `optimization`, `monitor`, `inference`, `decoding`, `evaluation` |
63+
| `train` | `system`, `model`, `data`, `optimization`, `monitor` |
64+
| `test` | `system`, `model`, `data`, `inference`, `decoding`, `evaluation` |
65+
| `tune` | `system`, `model`, `data`, `inference`, `decoding`, `evaluation` |
66+
67+
Mode names:
68+
- use `train`, `test`, `tune`;
69+
- if a future CLI exposes `infer`, `decode`, or `evaluate`, they should resolve
70+
from `test`-compatible sections unless separate dataclasses are needed.
71+
72+
Required v2 fields:
73+
- `inference.decode_after_inference: bool`
74+
- `inference.chunking.output_mode: "decoded" | "raw_prediction"`
75+
- `inference.saved_prediction_path: Optional[str]`
76+
- top-level `decoding` list or decoding stage config
77+
- top-level `evaluation` config
78+
79+
Rejected v1 patterns:
80+
- `inference.decoding`
81+
- `inference.evaluation`
82+
- `shared`
83+
- optimizer step aliases other than `n_steps_per_epoch`
84+
- WandB fields with duplicated `wandb_` prefix
85+
- MONAI arch aliases that duplicate canonical schema fields
86+
87+
## Boundary Rules
88+
89+
- Runtime packages import dataclasses and loader functions, not profile-engine
90+
internals.
91+
- Config code may validate cross-section consistency, but it must not import
92+
training, inference, decoding, or evaluation execution code.
93+
- YAML profile application happens before dataclass conversion.
94+
- Stage resolution happens before execution.
95+
96+
## Implementation Order
97+
98+
1. Add or confirm top-level `decoding` and `evaluation` stage schema.
99+
2. Add `inference.decode_after_inference` and
100+
`inference.chunking.output_mode`.
101+
3. Remove any old aliases and compatibility validators.
102+
4. Update all tutorial YAMLs to v2 fields.
103+
5. Add strict failure tests for removed fields.
104+
6. Run config load tests for every tutorial.
105+
7. Run stage-specific runtime tests.
106+
107+
## Tests
108+
109+
Add or update:
110+
- config rejects `shared`;
111+
- config rejects nested `inference.decoding`;
112+
- config rejects nested `inference.evaluation`;
113+
- config rejects removed optimizer and WandB aliases;
114+
- all tutorials load with v2 schema;
115+
- profile selectors work only at canonical paths;
116+
- stage allowlists reject invalid sections;
117+
- `--debug-config` prints the resolved v2 structure.
118+
119+
## Open Decisions
120+
121+
- Whether `infer`, `decode`, and `evaluate` should become first-class CLI modes
122+
in v2 or remain subcommands/wrappers around `test` stage config.
123+
- Whether scheduler-specific `ReduceLROnPlateau` fields should move fully into
124+
`scheduler.params`.

‎.claude/refactor/data_v2.md‎

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
# Data V2 Plan
2+
3+
## Goal
4+
5+
Keep `connectomics.data` domain-based and remove compatibility leftovers. The
6+
module should own IO, preprocessing, augmentation, label transforms, datasets,
7+
sampling, and splitting with one canonical path for each object.
8+
9+
Baseline: `data.md` reports most data cleanup as complete. V2 mainly deletes
10+
remaining shims and turns the documented structure into the enforced structure.
11+
12+
## Public API
13+
14+
Keep these public subpackages:
15+
16+
- `connectomics.data.io`
17+
- `connectomics.data.augmentation`
18+
- `connectomics.data.processing`
19+
- `connectomics.data.datasets`
20+
21+
Public concepts:
22+
- volume IO: read, save, shape, format detection
23+
- MONAI dictionary transforms for IO, augmentation, and label processing
24+
- pure augmentation operations in `augment.augment_ops`
25+
- deterministic label/target operations in `process`
26+
- dataset classes and sampling/split helpers in `dataset`
27+
28+
Everything else should be private or package-internal.
29+
30+
## Delete
31+
32+
- `connectomics.data.utils` compatibility shim.
33+
- Duplicate copies of `sampling.py` and `split.py` outside `data.datasets`.
34+
- Any `monai_transforms.py` compatibility names.
35+
- Any unused dataset classes kept for legacy code.
36+
- Any transform builder functions that are no longer called.
37+
- Any old exports in `data.__init__` that point to deleted modules.
38+
39+
## Move/Rename
40+
41+
Canonical paths:
42+
43+
| Concept | V2 owner |
44+
| --- | --- |
45+
| Train/val volume split | `connectomics.data.datasets.split` |
46+
| Sample counting | `connectomics.data.datasets.sampling` |
47+
| IO transforms | `connectomics.data.io.transforms` |
48+
| Augmentation transforms | `connectomics.data.augmentation.transforms` |
49+
| Label/process transforms | `connectomics.data.processing.transforms` |
50+
| nnUNet preprocessing | `connectomics.data.processing.nnunet_preprocess` |
51+
| Pure augmentation functions | `connectomics.data.augmentation.augment_ops` |
52+
53+
## Config Contract
54+
55+
Data config should use domain terms:
56+
- `data.input`
57+
- `data.label_transform`
58+
- `data.augmentation`
59+
- `data.dataloader`
60+
- `data.split`
61+
62+
Rules:
63+
- no config objects inside transform classes;
64+
- transform constructors receive plain Python parameters;
65+
- no hidden old-key handling in builders;
66+
- train/test/tune data shorthands must be documented in schema, not handled
67+
through ad hoc fallback logic.
68+
69+
## Boundary Rules
70+
71+
- `data` must not import `training`.
72+
- `data` must not import `inference`.
73+
- `data` must not import `decoding` except through tests or examples.
74+
- Pure functions must not import MONAI.
75+
- MONAI wrappers should be thin adapters around pure functions.
76+
- Dataset code should not know about model architectures or losses.
77+
78+
## Implementation Order
79+
80+
1. Delete `connectomics/data/utils`.
81+
2. Update tests and tutorials to import from `data.datasets`.
82+
3. Remove stale exports from `data.__init__` and subpackage `__init__` files.
83+
4. Search for deleted names across Python and YAML.
84+
5. Tighten transform builders so unsupported config shapes raise early.
85+
6. Add boundary tests for no `data.utils` import path.
86+
7. Run data, training data-factory, and tutorial config tests.
87+
88+
## Tests
89+
90+
Add or update:
91+
- importing `connectomics.data.utils` fails;
92+
- production imports use `data.datasets.split` and `data.datasets.sampling`;
93+
- every `RandomizableTransform` calls `self.randomize()` in `__call__`;
94+
- pure augmentation ops import no MONAI symbols;
95+
- transform builders reject unsupported config types;
96+
- datasets preserve image/label padding semantics;
97+
- nnUNet preprocessing round-trips metadata needed by inference output.
98+
99+
## Open Decisions
100+
101+
- Whether `connectomics.data.__init__` should export only subpackages or also
102+
selected high-use functions like `read_volume` and `save_volume`.
103+
- Whether test/tune `data.batch_size` shorthand should stay or move under
104+
`data.dataloader.batch_size` only.

‎.claude/refactor/decoding_v2.md‎

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
# Decoding V2 Plan
2+
3+
## Goal
4+
5+
Make decoding a standalone stage that consumes raw prediction artifacts or
6+
arrays and writes segmentation artifacts. It should not construct models, load
7+
checkpoints, run sliding-window inference, or compute final evaluation reports.
8+
9+
Baseline: `decoding.md` reports the decoder registry and implementations are
10+
mostly clean. V2 focuses on stage separation and artifact contracts.
11+
12+
## Public API
13+
14+
Keep:
15+
16+
- `Decoder`
17+
- `DecoderOutput`
18+
- `register_decoder`
19+
- `get_decoder`
20+
- `list_decoders`
21+
- `decode_prediction`
22+
- `run_decoding_stage`
23+
- decoder implementations for affinity, watershed, synapse, and ABISS paths
24+
- decoding parameter tuning helpers if they operate only on saved predictions
25+
and labels
26+
27+
## Delete
28+
29+
- Any decoder registration call duplicated outside package initialization.
30+
- Any decode path that assumes model inference is happening in the same call.
31+
- Any final metric reporting that belongs to `evaluation`.
32+
- Any postprocessing module names that conflict with inference output
33+
postprocessing semantics.
34+
- Any print-based operational output.
35+
- Any old config nesting under `inference`.
36+
37+
## Move/Rename
38+
39+
Canonical responsibilities:
40+
41+
| Concept | V2 owner |
42+
| --- | --- |
43+
| Decoder registry | `connectomics.decoding.registry` |
44+
| Decoder base types | `connectomics.decoding.base` |
45+
| Segmentation decoders | `connectomics.decoding.segmentation` |
46+
| Synapse decoder | `connectomics.decoding.synapse` |
47+
| Decoding postprocess | `connectomics.decoding.postprocess` |
48+
| Decoding stage runner | `connectomics.decoding.stage` |
49+
| Decoding tuning | `connectomics.decoding.tuning` or current tuner module |
50+
51+
If `optuna_tuner.py` remains large, split stage-independent parameter search
52+
from CLI/reporting helpers.
53+
54+
## Config Contract
55+
56+
Top-level `decoding` config owns:
57+
- decoder name;
58+
- decoder kwargs;
59+
- postprocessing kwargs;
60+
- output segmentation artifact path;
61+
- optional tuning search space for decoder/postprocess parameters.
62+
63+
It does not own:
64+
- model checkpoint;
65+
- sliding-window parameters;
66+
- raw input image paths, except for metadata or visualization;
67+
- metric selection.
68+
69+
Decode-only mode:
70+
71+
```yaml
72+
inference:
73+
saved_prediction_path: /path/to/raw_prediction.h5
74+
decoding:
75+
- name: decode_affinity_cc
76+
kwargs:
77+
threshold: 0.7
78+
backend: numba
79+
```
80+
81+
V2 may rename this so `decode.input_prediction_path` is canonical. If so,
82+
`inference.saved_prediction_path` should be deleted and decode-only mode should
83+
not depend on the inference section.
84+
85+
## Boundary Rules
86+
87+
- Decoding may import data processing helpers and top-level utils.
88+
- Decoding may import config dataclasses.
89+
- Decoding must not import training or inference managers.
90+
- Decoding must not write metrics reports except optional decoder diagnostics.
91+
- Decoding output artifacts must be accepted by evaluation without special
92+
decoder-specific handling.
93+
94+
## Implementation Order
95+
96+
1. Decide whether prediction input path lives under `decoding` or `inference`.
97+
2. Add `run_decoding_stage` that accepts a prediction artifact path and decoder
98+
config.
99+
3. Update combined test path to call `run_decoding_stage`.
100+
4. Remove decode logic from inference/test orchestration.
101+
5. Ensure all decoders consume canonical artifact arrays and metadata.
102+
6. Move final metric calls out to evaluation.
103+
7. Add decode-only tests.
104+
105+
## Tests
106+
107+
Add or update:
108+
- decode-only mode works without model construction;
109+
- decoding consumes a raw prediction artifact from inference;
110+
- decoded segmentation artifact metadata records decoder name and params;
111+
- all registered decoders are discoverable;
112+
- invalid decoder names raise clear errors;
113+
- decoding does not import training or inference execution modules;
114+
- tuning can run from saved predictions.
115+
116+
## Open Decisions
117+
118+
- Whether decode input path should be `decoding.input_prediction_path` instead
119+
of `inference.saved_prediction_path`.
120+
- Whether decoding postprocessing should remain a subpackage or be folded into
121+
decoder-specific modules where only one decoder uses it.

0 commit comments

Comments
 (0)