Skip to content

AGENTS.md

Latest commit

 

History

History
317 lines (259 loc) · 44.5 KB

AGENTS.md

File metadata and controls

317 lines (259 loc) · 44.5 KB
Raw

KonfAI Agent Guide

Canonical reference for humans and AI agents. Read it before changing code. User-facing detail lives in docs/ and examples/.

1. What KonfAI is

KonfAI is a modular, fully-configurable deep-learning framework for medical imaging (Boussot & Dillenseger, 2025, arXiv:2508.09823). A model, its data pipeline, losses/metrics, optimizer/schedulers, augmentations, and the whole train/predict/evaluate workflow are described in YAML and mapped onto Python objects by a reflection engine, without editing code. The same engine drives a fifth workflow, dataset preparation (TRANSFORM): read a dataset, apply a chain, write it back out. The config is a complete, reproducible record of the experiment. KonfAI has produced top-ranking MICCAI-challenge results (SynthRAD, TrackRAD, CURVAS, PANTHER) across segmentation, registration, and synthesis.

Three pillars run through the codebase:

  1. Config-by-reflection. apply_config(path) reads a callable's signature and builds its arguments from the YAML subtree it owns (@config("Key")), recursing into nested @config objects. Resolved defaults are written back to the file, so a run leaves a fully-resolved config on disk. Reading a config mutates it. Each workflow builder reads its file inside strict_config(root), which writes it once, when the build succeeds: a build that fails leaves the file byte-identical, except under --init (initializing()), which writes what resolved before the error. A key nothing bound reads is reported by path (TRANSFORM refuses, the other three warn), so everything a workflow reads from its file must be bound at construction, not lazily.
  2. Lazy, patch-based imaging. A volume is never loaded whole into RAM on a route that can stream: data is read as overlapping patches (optionally streamed) and predictions are reassembled with overlap blending. Mandatory invariant. Three exceptions are declared. TRANSFORM's whole-volume route, taken only when a stage cannot serve a region: the plan names it WHOLE-VOLUME, sizes it against the memory budget (Transform.working_multiple) and refuses the case when it does not fit. And the statistic pass of the TRAIN, PREDICTION and EVALUATION loaders (DatasetManager.warm_stream_statistics): a GLOBAL_STAT stage whose input an earlier stage changes gets its statistic from one whole-volume read per case before the case streams, which the loader says. And PREDICTION's one-pass reader with TTA copies (SamplesDataset._loads_whole), which loads a case whose whole-volume pass fits the memory budget: a sweep would replay the chain for every copy. Never an implicit read_data().
  3. Declarative models. Networks are routed add_module graphs, written as Python classes in konfai/models/, or entirely as a .yml via the YAML model builder.

2. Repository layout

Path Role
konfai/ Core package (config, data, network, metric, workflows, utils)
konfai-apps/ Independent package konfai_apps (app management, HF repos, FastAPI server) with its own pyproject.toml, deps, and CI
konfai-mcp/ Independent package konfai_mcp (FastMCP server exposing KonfAI to LLM agents) with its own pyproject.toml, tests, CI. On main since v1.6.0; published to PyPI by the release workflow.
konfai-studio/ Independent package konfai_studio (FastAPI BFF + built React front: a chatbot web UI over konfai-mcp). Own pyproject.toml; the front (konfai_studio/web/*) is a CI npm build, not in git; wheel-only.
konfai-scout/ Independent package konfai_scout (KonfAI Scout: a vision-language agent that inspects a segmentation, registration or synthesis result visually and reports localised defects with evidence and KonfAI Eval measurements). Own pyproject.toml, tests, published by the release workflow; depends on konfai alone.
apps/ Ready-to-use model app bundles (excluded from the konfai wheel)
examples/ Runnable Segmentation / Synthesis / Registration / Transform workflows, then the published apps (index: examples/README.md; assume CWD = the example dir)
docs/ · tests/ Sphinx site · core test suite (tests/unit, tests/integration)

konfai/models/ has two halves: models/python/<kind>/ (Python Network subclasses) and models/yaml/ (the shipped declarative catalog, 14 models). models/python has no __init__.py on purpose: it is a PEP 420 namespace package, and the wheel ships it via include = ["konfai", "konfai.*"] with namespaces=true; the catalog *.yml ship via package-data. Changing either breaks the wheel silently, so verify with a clean non-editable install, not an editable one.

Core modules worth knowing: utils/config.py (the reflection engine: read before any config change); utils/dataset/ (one module per storage backend SitkFile/H5File/OmeZarrFile/DicomFile/ItkTransformFile + the Attribute geometry sidecar + DataStream region writes); data/data_manager/ (groups, read order, samples, subsets, the per-workflow sources) + data/patching/ (stage, budget, sizer, sweep, blend, accumulate, grid, manager: the lazy patch index, DDP sharding, overlap-blended reassembly, SlabRegionStream) + data/materialize.py (CaseMaterializer: the TRANSFORM materialization engine over the manager's plan/replay API); data/geometry.py (grids/boxes/affine maps in world coordinates: pure numpy, crosses the mp.spawn pickle boundary) + data/sampling.py (coordinate producers + the torch gather; no SimpleITK); network/network/ (ModuleArgsDict/Network/ModelLoader/Measure across base, loaders, measure, network, model: the heart of the model system); metric/measure/ (Criterion = losses + metrics, one module per family); data/transform/ (one module per family, every stage re-exported by the package) + data/augmentation/ (base, spatial, color, intensity, placed); data/reduction.py (Reduction/Mean/Median/Std/Vote/Concat: a bare combine/reduction name resolves in this module) + data/case_reduction.py (the N-cases-to-one engine); trainer.py/predictor/ (output sink, ensemble, loop, workflow)/evaluator.py/transformer.py (the pipelines; transformer.py is the dataset-preparation one); api.py (the workflows as Python callables); export.py (ONNX export, extra export); config_schema.py (the JSON Schema of each workflow config, read off the constructors and published by the docs build: a new dict[str, X] whose keys name components is registered in its chains table, or the schema refuses their arguments); main.py (CLI); utils/{errors,model_builder}.py, utils/runtime/ (environment, logging, distributed).

3. How it fits together

Commands → config files. KonfAI is command-driven; five CLI states map to four YAML files:

Command File Root key Purpose
TRAIN / RESUME Config.yml Trainer: Model + dataset + losses + augmentations + optimizer/schedulers + training params
PREDICTION Prediction.yml Predictor: Load model(s), patch/TTA/ensemble inference, output post-processing
EVALUATION Evaluation.yml Evaluator: Predictions vs ground truth → per-case + aggregate metric JSON
TRANSFORM Transform.yml Transformer: Dataset preparation. Read a dataset, run a chain per case, Write the result. Plan first; an existing output is skipped, so a run resumes

Each run writes a workspace keyed by train_name (name for TRANSFORM): Checkpoints/, Statistics/ (TensorBoard + the resolved-config snapshot), Predictions/, Evaluations/ (metric JSON), Transforms/ (per-rank logs opening with the plan, the config copy, outputs.json, but not the transformed data, which lands wherever each Write: points, in the user's own tree). Dataset/ is the input data directory, not a run output.

Python API. The workflows are also callables: konfai.transform / run_transform(dry_run=True) / evaluate / predict / train (api.py, lazily re-exported from konfai). Same engine, two spellings: a chain is a list of live stage objects or the equivalent dict tree. The contract: a designed refusal raises KonfAIError (only the konfai and konfai-apps CLIs catch it and exit), results come back structured (TransformResult/EvaluationResult), the KONFAI_* environment and the published per-rank memory budget are restored around every call, one workflow per process (lock), and a caller's config FILE is copied to scratch, so the write-back never lands on it (a model YAML named by a relative path is made absolute against the caller's file first, or the working directory for a tree; the scratch is released when the call returns).

Conventions. Arrays are channel-first [C,(Z),Y,X]; geometry/spacing is (x,y,z) (SimpleITK). A chain (transforms:, patch_transforms:, criteria) is a mapping keyed by class name (a repeated class is spelled module-qualified the second time) or a YAML list of one-entry mappings, which binds in order under occurrence keys (Clip, Clip#2, Clip#3); get_module drops the suffix when it resolves the class. Attribute geometry keys are Origin/Spacing/Direction. In Resample, a spacing/shape value ≤ 0 is the keep-this-axis sentinel (that axis keeps the source's own density/extent). The words with two meanings (group, fold, worker, workspace, bundle, Metric_TRAIN) are defined in docs/source/reference/glossary.md.

Network graph. add_module(name, module, in_branch=[...], out_branch=[...], alias=...) wires a string-keyed branch register (branch '0' = input; execution = insertion order). Named module outputs are referenceable in YAML: an outputs_criterions key is a module's dotted path like UNetBlock_0:Head:Softmax (the :/. separators are load-bearing). out_branch:[-1] marks a terminal/deep-supervision head; alias lists are positional and load-bearing for pretrained-weight remapping.

Runtime. Workflows run under run_distributed_app (utils/runtime/distributed.py): it builds the configured DistributedObject, sets the KONFAI_* env vars, forces KONFAI_CONFIG_MODE='Done', and runs one rank per GPU (or per --cpu worker): a single rank runs inline in the calling process (KONFAI_INLINE_SINGLE_RANK, default on), several are spawned, or the job is submitted to SLURM via submitit. Disk/log side effects are gated on global_rank == 0.

For the full config-key catalogue and a concrete end-to-end trace, read the docs/ config guides and examples/.

4. Extending KonfAI

Every extension point is "subclass a base, reference it by classpath in YAML", with no core edits:

  • Model: subclass network.Network, build the graph in __init__ via add_module. Reference classpath: module.MyNet, a local Model:MyNet, a .yml, or default|<Name>.yml for the shipped catalog.
  • Pretrained weights: utils/pretrained.py:transfer_weights_by_execution_order pairs weighted leaves in forward-execution order (no key map). It fills every target tensor or raises: a tensor held by a parent module (torch.nn.MultiheadAttention owns in_proj_weight beside its out_proj child) or by a submodule the forward skips cannot be paired. Unreached source branches (nnU-Net deep-supervision heads) are ignored on purpose.
  • Loss / metric: subclass metric.measure.Criterion; forward returns a Tensor (loss) or a (value, dict) tuple (metric; consumers isinstance-branch). Attach under outputs_criterions/metrics to a named module output. A config that leaves is_loss out takes the criterion's own role, written back: declare default_is_loss = False for a score whose loss form must be asked for (SSIM, then 1 - SSIM with its gradient, read off as_loss) and loss_capable = False for one that has none (PSNR: is_loss: true is refused). Declare batch_mean = True only if the reported value and the loss are each the mean of one value per patch of the batch: a logged mean then weighs each batch by its patch count (validation over N ranks equals one rank); a pooled value (Dice's sums, a norm, a sum reduction) keeps one per batch. Optional-dep criteria import lazily via _require_optional(...) and raise an actionable MeasureError, never a bare top-level import. Declare any_target_grid = True if forward never compares its target voxel to voxel (Mean, Gram): EVALUATION's shape check and geometry warning then leave the metric out. Declare resamples_target = True if it puts a target of another shape on the output's grid itself (Dice, FocalLoss): the shape check leaves it out, and the geometry warning takes it in when both shapes match, since it then compares voxel to voxel.
  • Transform: subclass data.transform.Transform; implement __call__ and transform_shape() (must predict the output spatial shape exactly, because patch planning depends on it). Declare working_multiple (class attribute) if the whole-volume call allocates beyond its input and output, in volumes-worth of the case (Gradient 5.0; Resample 6.5 as a class, plus what the case's own field costs through case_working_multiple(name), which is the override for a cost only the case settles): the plan sizes the whole-volume fallback with it. Declare single_process = True if the stage cannot run inside a daemonic DataLoader worker (KonfAIInference spawns its nested run): a bound stage that declares it turns the loader's workers off. Declare patch_locality() (a LocalityKind: POINTWISE/HALO/ORIENTATION/CROP/GLOBAL_STAT/REGRID/SLAB/WHOLE_VOLUME) or the base default makes it WHOLE_VOLUME; a WHOLE_VOLUME that is a property of the configuration rather than of the stage must carry reason=, which the plan prints; without it the reader has nothing to change. A per-voxel stage that reads a companion volume (a mask) stays POINTWISE and overrides stream_region(name, tensor, context, attribute): both dispatchers (the reader's patches and regions, the writer's slabs) hand it where its region sits; it also overrides plan_region_reads(name, contexts), which a sweep calls once before its first region with the contexts stream_region will be handed, in order, so the stage declares its companion's windows to the dataset holding it (Dataset.plan_region_reads) and the chunk cache evicts them by next use (a hint: values never depend on it; the patch route, in the loader's order, never calls it); SLAB+stream_slab is the write-side-only contract for a side effect needing the output's slabs in order (InferenceStack), and the reader treats it as WHOLE_VOLUME. A GLOBAL_STAT stage declares stat_keys (a subset of Min/Max/Mean/Std, or their PerChannel forms), records (which of them its whole-volume call leaves in the scope under those names; Statistics none, a Clip bound only when saved) and takes_present (it takes a statistic already there over measuring its own, as Normalize and Standardize do): the streamed route seeds each stage right before it runs, from the store or from a whole-volume pass, and takes back what that route would not have left, so the case attribute, the backup a whole-volume load starts from and a Save header never carry a seed the stage did not record. A REGRID stage owns both halves: stream_region_source() (the pull map) and stream_region(). Pair inverse() if apply_inverse; override prepare(konfai_args) only when the stage builds a sub-object from configuration of its own (Reduce → its operator).
  • Augmentation: subclass data.augmentation.DataAugmentation; _state_init (sample params per case index) + _compute (apply lazily; it returns a fresh tensor or a view and never writes into the tensor it is handed: the copies alias the case). Only PlacedMask, Permute and a quarter-turn Rotate may change shape; a draw that reorders axes returns its index remap from _axis_remap, so an Expand copy's header (Spacing, Direction, Origin) follows the axes as the Permute transform's does. Declare _patch_locality like a transform so the copies stream (Rotate quarter turns are ORIENTATION, free angles and Scale REGRID through the affine's own pull box, Noise/CutOUT POINTWISE because their field and box are functions of the voxel's position); a draw that depends on the place overrides _stream_region(name, index, a, tensor, context). A copy's draws are keyed by the case NAME (_drawn_from(seed, name, kind, occurrence)), never by its index: an Expand copy's and a PREDICTION TTA copy's (DataAugmentationsList.draw_seed, from manual_seed); TRAIN's copies redraw every epoch from the seeded global RNG. A draw is also a chain stage: TransformLoader resolves a bare name against data.transform first and data.augmentation second, and reverses that order past an Expand marker (prefer_augmentation), where the chain is the copies' draws: so Flip, Permute and Foreign, which both packages define, are the transform before the marker and the draw after it (the module-qualified spelling forces either). That is how a transforms: block interleaves draws and transforms.
  • Reduction: subclass data.reduction.Reduction; implement __call__(list[Tensor]) -> Tensor over the [1, K, C, *spatial] layout both engines hand over. Two consumers, one vocabulary: the predictor folds one case's copies (ensemble/TTA), data.transform.Reduce folds N cases into one. Declare voxel_local = True only if every output voxel reads the same voxel of each input (a wrong True corrupts a streamed output: the gate checks nothing else), incremental = True if accumulate can fold one at a time, working_multiple = the buffers-worth the operator allocates ON TOP of what it is handed (the plan multiplies it into the peak it sizes regions against). An operator whose route depends on the COUNT overrides working_multiple_for(cases) and the attribute stays its worst case: Median selects the middle by a network of element-wise min/max up to five members (measured 1.5 at two, 1.0 at three, 2.5 at four, 1.5 at five) and sorts the stack past that (4.0), and override output_channels(channels, cases) when the fold changes the channel count (Concat does). Reduce refuses a non-voxel_local operator outright.
  • Imaging format: add a Dataset.AbstractFile backend, register its format token in BACKENDS (utils/dataset/backend.py), register aliases in SUPPORTED_EXTENSIONS, or in SUPPORTED_BACKEND_FORMATS when the token is not a suffix any file carries (:itktransform writes <group>.h5), since only the extensions are probed on disk; import-guard the heavy lib and raise a DatasetManagerError naming the extra rather than degrading in silence. Declare read_errors, what the library raises on an entry it cannot open or decode: Dataset turns exactly those into a CaseReadError naming the case, and PREDICTION and EVALUATION set that case aside where anything else stops them. Override read_granularity(name) when serving a window costs more than the window: a chunked backend decodes whole blocks (zarr, HDF5) and answers its chunk, and a memory-mapped one is served band by band (SitkFile answers (1, 1, Y, X): one step along the outermost axis a window spans, every axis below it whole, because those are the pages the read touches and the kernel counts). A sweep is priced and cut on that grid, so the grain need not be isotropic and need not come from a compressor; a backend that stays silent is priced at what its reads ask for, which is right only when a read costs exactly that. A compressed ITK file (.nii.gz, a zlib MetaImage, a gzip NRRD) is read by region from an uncompressed twin (utils/dataset/decompressed.py), decompressed once per run under a file lock into KONFAI_DECOMPRESSED_DIRECTORY (default ~/.cache/konfai/decompressed): its bounded_region_reads is True and its read_granularity the twin's band before the twin exists. A region that is the whole volume makes no twin (it decodes once either way). A twin is filed under the entry it is read for (the name Dataset resolves: the case, or an Expand copy's entry however deep its pattern), so DatasetManager.release_case(), which a new one-pass reader calls as it leaves a case, removes the twins of that case and of its copies whichever root or stage read them (a Mask group, a Resample field, a Save cache of the copies); run_scope (entered by execute_distributed_object and by each rank, which a cluster task reaches with no launcher around it) removes the run's twins however it ends.

Classpaths: a bare name (e.g. Dice) resolves inside that kind's package; module:Class imports any module: a local file (Loss:MyWrapper) or an installed library (monai.losses:DiceLoss, torch:nn:L1Loss). Model classpaths resolve against konfai.models.python. Use the relative or default| form.

YAML model builder (utils/model_builder.py): builds a Network from a .yml, safe by construction (node types must come from two curated registries, no eval/import injection). The shipped catalog (models/yaml/, 14 models incl. UNet/NestedUNet/ResNet/UNETR/ViT/VNet) now covers the feed-forward subset; custom-forward models (DiffusionGAN/ConvNeXt) stay Python. default|<Name>.yml addresses the flat catalog only; a name with a path separator is refused.

5. Apps (konfai-apps)

A separate package layered on KonfAI's public API (core never imports it). An "app" bundles a config + custom .py + .pt weights, resolved from a Local dir, a HuggingFace repo, or a Remote server; the apps/* bundles are thin CLI wrappers.

⚠️ Trust model. Resolving an app copies and imports its .py files → it runs arbitrary code. It also pip-installs its requirements.txt by default (only missing/mismatched packages; core packages like torch/konfai are never touched; opt out with KONFAI_APPS_INSTALL_REQUIREMENTS=0). Only resolve apps from sources you trust.

5b. MCP server (konfai-mcp)

A third package exposing a FastMCP server. It depends on KonfAI's public API and, for validation, on two core internals: konfai_mcp.runner imports konfai.trainer/konfai.network.network builders and patches the single-process runtime (_apply_single_process_patches), which is why the sibling pins are exact at a release tag (a range only from a working tree). It exposes the server so an LLM agent can inspect a dataset → author YAML → run train/predict/evaluate → monitor jobs → compare runs → iterate. On main since v1.6.0 and published to PyPI by the release workflow. Jobs run in a spawn subprocess (training may init CUDA); check_config(level='instantiate') and check_component run in a spawn subprocess (never in the server process), are side-effect-free (the child reads a scratch copy placed beside the config, so the original is never written, an edit made meanwhile survives, and the parent sweeps a copy a killed child left), and re-import edited workspace code; discovery is via list_components / describe_extension_points / describe_config_schema / check_component. Tests: pip install -e ./konfai-mcp then python -m pytest konfai-mcp/tests (the segmentation E2E needs the imaging extra).

Working on the MCP server, how to validate a change:

  • Synthetic fixtures: pixi run --environment dev python konfai-mcp/tests/make_fixtures.py builds a segmentation dataset, a registration pair with a known translation, a synthesis pair, a 3-level OME-Zarr store, and corrupted/unsupported inputs under the gitignored konfai-mcp/tests/fixtures/ (procedural, no patient data). Reuse these, do not invent ad-hoc data in /tmp.
  • Drive it black-box first, not by tool name. Formulate a real objective ("segment these CT volumes"), then exercise the loop through a fastmcp.Client exactly as test_mcp_server_segmentation_pipeline.py does. A new tool is not "done" because it returns without an exception.
  • Verify outputs, not return codes. After run_* + get_job_status(wait=True), assert the job status=="done" (never trust a green validate_* alone: its default level instantiate runs no train step; only level='train_step' runs one forward+backward), then open the produced files: read_file the config, check Predictions/<name>/Dataset exists, and read the Metric_*.json via get_run_metrics. Confirm the numbers correspond to the requested task.
  • Validating a new tool: (1) its next_actions must be registered tool names; the anti-drift test in test_mcp_server_tool_index.py enforces this for job payloads only, so if the tool emits next_actions in its own payload, assert it in the tool's pytest; (2) if it takes a workspace path, route it through resolve_workspace_relative_path (jail); (3) if it imports/executes app or workspace code, run it in the spawn subprocess (run_api_in_subprocess) and gate it behind allow_untrusted_code where applicable; (4) document per-parameter meaning via Annotated[..., Field(description=...)], not only prose; (5) add a pytest that inspects the output.
  • Adding a workflow kind is one WorkflowSpec entry in konfai-mcp/konfai_mcp/workflows.py plus the two Literal aliases beside it (and a GUIDE/tool description); every other map (config filename, root key, runner command, capabilities class, retry tool) derives from that table, and tests/test_workflow_registry.py pins the derivations to it.
  • Safety invariants to preserve: validation/smoke-tests never execute in the server process; read/write_session_file and the output of run_app and fine_tune_app are path-jailed to the session; fetch_app(path=...) (path) and package_app_from_session (output) write to a host folder the caller names, by design, and fetch_app(path=...) refuses to replace a file already there unless overwrite=True (dataset tools read arbitrary host paths by design, so keep it that way only for the trusted-local deployment, and never widen writes). cancel_job now reaps the whole process group: the job runs os.setsid() and cancel sends the signal via os.killpg, so mp.spawn DDP grandchildren are killed with the middle process (regression test: test_cancel_reaps_the_whole_process_group_including_grandchildren).
  • Regenerate derived docs: after changing a tool's description run python konfai-mcp/scripts/generate_tool_reference.py (the committed skill reference is generated).
  • KonfAI Scout and FidOpt (konfai_mcp/server_scout.py): run_scout_inspection runs a visual inspection (the separate konfai-scout package) as a tracked job and read_scout_report read it back without images; images reach a cloud model only with allow_cloud_images. The fidopt_* tools reach FidOpt through its own CLI and Python (KONFAI_MCP_FIDOPT), launch only KONFAI_MCP_FIDOPT_ALLOWED operations, and reimplement nothing of it. Studio's Scout tab reads the same reports (/api/scout/*).

5c. KonfAI Scout (konfai-scout)

A fifth package, konfai_scout: an agent that looks at ONE result the way a person would at a viewer (coverage montages, then the slices, zooms and overlays it asks for), measures regions through KonfAI Eval, and writes a versioned report of localised defects (measured property / suspected error / confirmed error, severity, confidence, the views and measurements cited). It depends on konfai only and reuses it everywhere it can: konfai.data.geometry for grids and boxes, konfai.utils.dataset for the staged layout (entry_backend, link_entry, the same as konfai-apps) and region reads through every backend, konfai.transform with Resample and TransformPoints for a registration's warped data, konfai.evaluate for every metric, konfai.utils.ITK.field_jacobian_determinant for the Jacobian view. Scout adds the viewer (views.py, planes.py), the tools and prompts the model drives, the loop, the report and the providers (Claude through the API or through Claude Code's Agent SDK, GPT through the API or through the Codex CLI, any OpenAI-compatible local server, a mock for tests). Images leave the machine only under allow_cloud_images. note enforces the epistemic rules mechanically (an error cites a view that shows its box, a figure comes from a cited measurement, a measurement covers the box; an empty measurement confirms nothing), an optional blind second reading re-reads every error from crops alone, and synthetic.py builds clean cases (negative controls) beside the planted ones so invented defects are counted (examples/study_controls.py). konfai-mcp exposes it (run_scout_inspection, §5b) and Studio shows its reports. konfai-scout/README.md and konfai-scout/docs/ hold the user side, the report schema and the studies.

6. Running things

pixi run check                                                    # lint + format-check + core tests + apps tests (run ONCE before finalising); the test tasks pin OMP_NUM_THREADS=1 (8x faster on 24 cores)
pixi run test                                                     # core unit + integration (tests/), ~3 min (pytest-xdist), not an iteration loop
pixi run test-fast                                                # dev loop (~25 s): skips slow oracle + integration tests
pixi run --environment dev typecheck                              # mypy konfai
pixi run --environment dev python -m pip install -e "./konfai-apps[server]" && pixi run --environment dev python -m pytest konfai-apps/tests   # apps suite (separate; [server] for the job-server tests)
pixi run --environment dev python -m pip install -e ./konfai-mcp && pixi run --environment dev python -m pytest konfai-mcp/tests    # mcp suite (separate; the pin is a range from a working tree, exact at a tag)
pixi run --environment dev python -m pip install -e ./konfai-studio && pixi run --environment dev python -m pytest konfai-studio/tests  # studio suite (separate; build the front first or its two front tests skip)
pixi run --environment dev python -m pip install -e "./konfai-scout[dev]" && pixi run --environment dev python -m pytest konfai-scout/tests  # scout suite (separate; no API key: fake clients and a mock model)

The Pixi dev env and a bare pip install .[dev] carry the same dependency list, imaging extras included (the dev extra IS the dev environment). pixi run test does not run the sibling suites (konfai-apps/, konfai-mcp/, konfai-studio/, konfai-scout/); install those packages first (they pull their own runtime deps), exactly as their CI does. Studio's suite also needs its React front on disk (cd konfai-studio/frontend && npm ci && npm run build, emitting konfai_studio/web/): without it the two tests of the served front skip and the run still reports green. Install runtime extras with pip install konfai[<extra>] (itk, hdf5, dicom, omezarr, imaging, tensorboard, lpips, cluster, export, check, …).

6b. Releasing

Versions are tag-derived (setuptools_scm, ^v(?P<version>.*)$) for every package, so never hand-edit a version. Pushing a v* tag runs .github/workflows/publish.yml: test (core + apps + mcp + studio) → build the 10-package matrix (konfai-apps, konfai-mcp, konfai-studio, konfai-scout, and the 5 apps/* bundles) → publish to PyPI via OIDC → build the Docker image once konfai is visible on PyPI. apps/* bundles pin konfai== and konfai-apps==, and konfai-studio pins konfai-mcp== (via its setup.py), all the same version, so the whole matrix releases in lockstep. konfai-studio is the one exception to the pure-Python build: the build job runs npm ci && npm run build first (its React front is git-ignored) and then python -m build --wheel (wheel-only, because the sdist file-finder would drop the built web/). konfai-mcp is wheel-only too: its setup.py copies the repository's examples/ into the wheel, and an sdist holds none to copy. A break may ship in any release, with a refusal naming the new spelling and no deprecation period (docs/source/development.md).

Before tagging: pixi run check green; both sibling suites green; the HF bundles tagged vX.Y.Z too (§7b); pixi run perf-check within its thresholds on a quiet machine (or the regression named in the release notes; the facts of the series, voxel identity and exact differences, also gate every pull request on a GitHub runner, see benchmarks/perf/README.md); and, because the test job only exercises the source tree, confirm the built wheel still ships konfai/models/python/** and konfai/models/yaml/*.yml by installing it non-editable in a clean venv (an editable install hides PEP 420 / package-data breakage).

The changelog is written before the tag, not after. CHANGELOG.md is generated from the commits and the publish workflow renders the same section into the GitHub Release, so the file must already contain the version being tagged:

uvx --from commitizen cz changelog --unreleased-version vX.Y.Z --start-rev v1.5.8   # 1. draft the section
git commit -am "ci: changelog for vX.Y.Z"                                          # 2. commit it
git tag -s vX.Y.Z -m "vX.Y.Z" && git push origin vX.Y.Z                            # 3. sign, tag, push

The draft is not the section: it sees commit subjects only (a squash merge is one line, a subject with no conventional prefix is dropped), and the command regenerates the whole file, edited sections included. Keep the new section, restore the rest from git, then write what a user gets (docs/source/development.md, "Cutting a release").

-s and not -a: releases are signed (v1.7.0 is), and tag.gpgsign already makes -a sign on a machine that has it configured, which is exactly why -s is spelled out. On a machine that does not, -a publishes an unsigned release in silence where -s refuses. Commitizen's own gpg_sign is deliberately NOT set: it would only apply to cz bump, which does not tag here, so it would be a setting that describes an intention nothing acts on.

cz bump computes the right increment (cz bump --dry-run is a good second opinion on major/minor/patch) but does not tag here: with version_provider = "scm" there is no version file to write, so when the changelog is already current it has nothing to commit and stops silently. Tag by hand. --start-rev v1.5.8 is required: Conventional Commits only took hold at v1.5.9, and rendering further back emits empty version headings.

7. Invariants, do NOT break

  • Never load a full volume into RAM on a route that can stream. Use lazy/patch/streaming access (can_stream_patch, read_data_slice). TRANSFORM's WHOLE-VOLUME verdict, the loaders' statistic pass and PREDICTION's budget-fitting TTA case (§1) are the exceptions, each declared where it runs, never a read a stage takes on its own.
  • Channel-first [C,(Z),Y,X]; spacing (x,y,z). Attribute stringifies every value and reparses geometry via np.fromstring, so only flat scalars / 1-D arrays round-trip. Read via __getitem__/get_np_array.
  • KONFAI_config_file + KONFAI_CONFIG_MODE must be set before any Config() (tests must monkeypatch.setenv both); workflows require KONFAI_CONFIG_MODE='Done'. Reading a config rewrites it on disk once the build succeeds; a failed build leaves it as it was (--init excepted).
  • Patch ordering must match between read (disassemble) and write (Accumulator); for PREDICTION/EVALUATION all patches of a case stay on the same DDP rank.
  • outputs_criterions keys equal a module's dotted path; the :/. separators are load-bearing.
  • state_dict load/save does not recurse into nested Networks (each owns its optimizer/state); alias lists are positional.
  • The YAML model builder is the trusted/untrusted boundary: only registry types, and module names contain no ..
  • konfai-apps is a separate package; apps/ is excluded from the konfai wheel. Core never imports konfai_apps: a chain names its stage konfai_apps.transforms:KonfAIInference.
  • The pretrained bridge fills every target tensor or raises; never report a partial load as success.
  • The config write is atomic (temp + os.replace); a reader must never see a truncated config and bind all-defaults.

7b. Contracts with the ecosystem

  • SlicerKonfAI (separate repo) drives konfai-apps by CLI + JSON. It imports, from konfai_apps.app_repository: current_free_vram(devices, remote_server=None), get_app_repository_info, is_app_repo, LocalAppRepositoryFromDirectory, LocalAppRepositoryFromHF, AppRepositoryError, plus the get_parameters() -> {values, constraints} params primitive that drives its advanced-params dialog. Renaming or dropping any of these breaks Slicer silently (it happened once: PR#33 dropped current_free_vram, PR#35 restored it). Grep the Slicer checkout before touching that surface. SlicerKonfAI and SlicerImpactReg also import nine core symbols: konfai.evaluator.Statistics, konfai.utils.dataset.image_to_data and get_infos, konfai.utils.runtime.MinimalLog, and from konfai itself get_available_devices, get_ram, get_vram, check_server, assert_konfai_install, pinned by tests/unit/test_slicer_core_api_contract.py. konfai-apps/tests/test_slicer_api_contract.py pins the konfai_apps names, the methods Slicer calls on an app repository with the arguments it passes, and the attributes it reads. A symbol whose only caller here is such a test is still in use.
  • HF bundles (hf_bundles/*) carry app.json + config + requirements.txt + .pt + custom .py. Their configs use bundle-relative classpaths (ResidualEncoderUNet.yml, model:Unet_TS_CT, Model:RegistrationNet), never KonfAI's internal module paths, which is why the models→models.python move did not break them. Validate a bundle by loading its config on a copy (reading mutates). An unpinned reference resolves to the tag v<konfai-apps version> when the repository has it, else main: tag every bundle repository vX.Y.Z at each release, once its configs are final (git tag vX.Y.Z && git push origin vX.Y.Z in each hf_bundles/* clone).

7c. Security boundaries

Three, and only three, places decide trust. Keep them honest:

  1. YAML model builder: registry types only; the default| catalog name must stay a bare filename.
  2. konfai-apps app resolution: runs the app's .py and pip-installs its requirements.txt by default (KONFAI_APPS_INSTALL_REQUIREMENTS=0 opts out). Protected core packages are matched by PEP 503 canonical name (konfai_apps ≡ konfai-apps); transitive deps are not policed, so say so and don't overclaim.
  3. konfai-mcp: validation/smoke-tests run only in a spawn subprocess (never the server process); read/write_session_file and the output of run_app and fine_tune_app are path-jailed; dataset tools may read arbitrary host paths by design, but writes must never widen (any tool that composes a write target must reject path separators). Two tools take a free host destination, by design: fetch_app(path=...) (path, which refuses to replace an existing file without overwrite=True) and package_app_from_session (output). cancel_job reaps the whole process group. Over SSE/HTTP with no bearer token, the CLI refuses a non-loopback bind (--i-know-this-is-insecure overrides) and the server answers only loopback Host names (DNS rebinding).

7d. Traps that have bitten before

  • An editable install hides packaging breakage (models/python is PEP 420; catalog .yml is package-data).
  • A green validate_* proves little: its default level instantiate runs no train step; only level='train_step' does a real forward+backward.
  • transform_shape() must be exact: patch planning trusts it, a wrong prediction corrupts reassembly.
  • Reading a config mutates it, so build from a scratch copy beside it (what konfai-mcp's validation does), never from the author's file.
  • The mirrors-mypy pre-commit hook cannot type-check this code: it runs in an environment of its own, without torch, and reads every torch attribute as Any | None (199 errors on a tree mypy passes). The hook is a local one running the dev environment's mypy, the same command as pixi run typecheck and the CI job; keep it that way.
  • Adding a workflow kind: in konfai-mcp this is one WorkflowSpec entry + two Literal aliases (drift-tested, see §5b); in core it still touches several maps (main.py _COMMANDS/_INIT_TARGETS, State, api.py). Prefer extending the descriptor tables over scattering new registries.
  • Union coercion in the config binder is declaration-order-sensitive: overlap: 0.25 once bound 0 (lossy int won). Fixed; pinned by test_config.py::test_apply_config_union_keeps_the_value_type_over_lossy_coercion. Any new union-typed config key still needs a test.
  • A model YAML named by a relative path resolves next to the config file that names it (ModelLoader._yaml_path). The Python API copies a config file to scratch before reading it (the write-back must not land on the caller's file), so _config_copy makes such a classpath absolute first; a new entry point that materializes a config elsewhere must do the same. Pinned by test_api.py::test_a_relative_model_yaml_is_anchored_to_the_config_files_directory.
  • A Prediction.yml that preprocesses a group differently from the Config.yml it trained on is silent: the run succeeds and only the output is wrong. The Synthesis example shipped with Standardize(mask: None) in training against Standardize(mask: MASK) plus an extra input Mask in prediction: same checkpoint, 409 HU of MAE instead of 98. Nothing checks it. The tell is an evaluated metric that contradicts the training loss by an order of magnitude (training MAE 0.018 in [-1, 1] ≈ 36 HU); diff the two transform stacks before blaming the model or the epochs.
  • The sibling packages pin konfai==<their own scm version>, which from a working tree is a .dev version no installed core carries. setup.py:_sibling pins the closest release or newer from a tree and the exact version at a tag; keep that helper when adding a package. Pinned by test_packaging.py::test_sibling_pins_resolve_against_the_core_of_this_tree.

Fixed, pinned by a test (do not re-fix; the guide once listed them as open):

  • The one-pass case FIFO: prediction and evaluation hold the case being finished and the next one, whatever the batch size (the training-sized batch_size + 1 kept whole cases loaded, a case per case of RAM growth on a cohort) (test_data_manager.py::test_a_one_pass_source_holds_two_cases_whatever_its_batch_size).
  • Nested-Network optimizer/scheduler state on RESUME: checkpoint_save and Network.load both use the dotted get_networks() key (test_network.py::test_load_restores_nested_network_optimizer_and_counters).
  • Inline augmentation redraws under persistent DataLoader workers: the loader forces persistent_workers=False when inline augmentation is active (test_data_manager.py::test_inline_augmentations_disable_persistent_workers).
  • The train/val split seed: drawn from manual_seed, else the Seed.txt the TRAIN run recorded on RESUME, else a fresh draw recorded for the next RESUME (test_trainer.py::test_resume_reuses_the_recorded_split_seed).

8. Conventions & rules

  • Code: line length 120 (Ruff); type annotations on new public functions; Apache-2.0 SPDX header on every new source file; prefer pathlib.Path; use the error classes in utils/errors.py (do not invent exceptions); import-guard heavy optional deps (SimpleITK/h5py/pydicom/zarr), failing at point-of-use with an install hint, not at import.
  • Commits: Conventional Commits (cz check): type(scope): subject, imperative, < 72 chars. No AI-agent branding (claude/codex/"generated by/with") and no AI co-author trailers.
  • For agents: read before editing; keep diffs small (one logical change per PR, no unrelated reformats); iterate cheaply, verify once: while developing run only the per-module test file you touched (pytest tests/unit/test_<module>.py), then pixi run test-fast (~25 s); reserve the full pixi run check (and the apps suite if you touched konfai-apps) for one final pass before finalising, because the full suite takes ~3 min and re-running it per edit adds up fast; no new runtime dependency without an explicit request + a matching pyproject.toml update in the same commit; update docs and tests/unit/test_config.py when changing config binding; do not skip pre-commit with --no-verify.

Code navigation

Prefer efficient code retrieval over broad file reads.

  • Use CocoIndex Code for conceptual or semantic code search.
  • Use Serena for symbols, references, declarations, and symbolic refactoring.
  • Prefer targeted reads over entire files.
  • Do not repeatedly search or reread code already understood.