fix: correctness, IO-safety, and packaging hardening in core and apps

[vboussot]

📚 Stacked PR — bottom of the stack. Base: main. #17 → #18 → #19 build on this; review/merge bottom-up.

Summary

First wave of the correctness/security audit: behaviour-fixing changes across the config reflection engine, the lazy patch pipeline, metrics/losses, the network/model graph, the training/resume path, the evaluator, and the konfai-apps server and client, plus packaging hygiene. Several of these are silent-corruption or hang-inducing bugs — config write-back dropping dict defaults, overlap-blending seams on volume edges, unequal-length DDP shards hanging NCCL, an unsafe pickle loader on remote checkpoints, and konfai-apps leaking into the konfai wheel. Each fix ships with a focused regression test.

Stacked PR — this is the bottom of the stack (#16 → main); #17–#19 build on it. Review/merge order matters.

What changed

Config reflection round-trip (konfai/utils/config.py)

• Persist dict[str, primitive] defaults on write-back instead of collapsing them to None/{} (which silently dropped them on the next run).
• Escape/unescape . and % in dict-key path components (_escape_key_component/_unescape_key_component) so keys containing dots survive the dotted-path split used to recurse into nested configs.

Lazy patch pipeline (konfai/data/patching.py, transform.py, augmentation.py, data_manager.py, konfai/utils/dataset.py)

• Overlap-blended reassembly divides by an accumulated per-voxel weight map in Accumulator, and PathCombine.set_patch_config short-circuits overlap <= 0 to unit weights — fixes darkened patch borders/corners on volume edges without whole-image padding.
• Crop.transform_shape treats shape as the already channel-stripped spatial shape, so the predicted output shape matches __call__ exactly (patch planning depends on this).
• Label maps resample with mode="nearest" under EulerTransform/Elastix (integer dtype → nearest) instead of blending class ids; Translate params are normalized to grid coordinates (value * 2 / (size - 1)) and Rotate quarter-turns are randomized per case/axis; augmentation state is drawn once per case and shared across destination groups (reset_state gating in reset_augmentation).
• get_infos and the per-rank reader in dataset.py reverse SimpleITK.GetSize() for every dimensionality, not only 3-D (2-D/4-D were transposed).
• HDF5: a per-file threading.RLock (_get_h5_file_lock) serialises the open/use/close sequence so cache workers cannot race/truncate; is_dataset_exist-path reads open read-only (Dataset.File(..., True, ...)).
• Streaming trailing transforms persist inversion state (e.g. TensorCast source dtype) back onto the case attribute (_stream_attributes_persisted), matching the non-streamed path.
• Padding inverse-origin correction indexes the correct axis (origin[dim]); ResampleToResolution/ResampleToShape now raise TransformError (were constructed and discarded); unsupported ITK transform types raise DatasetManagerError instead of falling through (elif chains).
• Subset selection is deterministic (sorted); shuffling moved to Data under an explicit subset.shuffle guard.

DDP training shards (konfai/data/data_manager.py)

• Contiguous per-rank shards are trimmed to a common min_len (DistributedSampler drop_last semantics) so every rank runs the same number of backward all-reduces per epoch — prevents the NCCL hang under static_graph=True. world_size == 1 keeps every sample.

Metrics & losses (konfai/metric/measure.py, schedulers.py)

• Dice excludes background (label 0) from auto-labels and normalizes by the count of present labels; SSIM computes over channels via channel_axis, returns a tensor, and tolerates a missing mask; Variance guards single-channel input; PerceptualLoss passes targets with *targets unpacking (was comparing against a list).
• Measure.Loss.get_loss aligns the current losses with their trailing weights so a scheduler drives the gradient; dict-payload metrics (per-label Dice/TRE) are nan-mean-summarized for logging in Measure.Loss.add.
• accepts_init dispatch reads the flag off the criterion module, not the CriterionsAttr value (was always False, silently skipping graph-rewiring criteria like KLDivergence).
• PolyLRScheduler.step respects last_epoch when set.

Network & models (konfai/network/network.py, models/)

• BatchNorm gamma initialised around 1.0 (pix2pix convention), not 0.0 (which stalled early training).
• named_forward resets its per-module inner-match set (tmp = []) so a sibling's output is not silently dropped.
• Mixed precision: Measure holds the GradScaler and scales the backward pass; the scaler is wired from Network when built.
• UNetBlock attention routes the gated skip to a free branch (out_branch=[2], skip_branch) so it captures the Multiply output, not an internal half-resolution projection.
• Experimental models fail fast with an actionable NotImplementedError: DDPM (broken time-embedding wiring) and 3-D VoxelMorph (2-D-hardcoded warping); dead CycleGanDiscriminator.initialized removed.

Training / resume / checkpoints / EMA (konfai/trainer.py, network.py, konfai/main.py)

• BEST-checkpoint scan and _load go through safe_torch_load; https:// model URLs are kept as raw strings (no Path() // collapse) and never fall back to the unsafe unpickler.
• Early-stopping is decided on rank 0 and broadcast to all ranks (_broadcast_stop via synchronize_data) so the loop is left together; the LR-decayed-to-zero clean stop is preserved.
• Checkpoint filenames disambiguate on collision; loss=None is stored as inf; EMA n_averaged is saved/restored (Model_EMA_n_averaged); the EMA blend uses the standard convention (decay * avg + (1 - decay) * model).
• New --lr override for RESUME: threaded build_train(lr=...) → Trainer.set_lr → Network.load(override_lr=...), resetting optimizer lr/initial_lr and scheduler base_lrs/last_epoch/_last_lr. None resumes the checkpoint LR and continues the schedule.

Evaluator (konfai/evaluator.py)

• Aggregate JSON is sanitized (_to_serializable → null for non-finite, allow_nan=False); count no longer becomes NaN; per-case metric mean guards c == 0; Statistics.read keys by the full metric path and drops component sub-entries by aggregate-prefix instead of the hardcoded Dice special-case.

IO safety & runtime (konfai/utils/runtime.py, ITK.py)

• New safe_torch_load (prefers weights_only=True; local checkpoints may fall back, https:// never does), reused across trainer.py, bundle.py, and app.py.
• is_interactive_session guards stdout.isatty (tolerates a non-tty Log/MinimalLog proxy); synchronize_data works without CUDA when dist.is_initialized(); workflow seeding (np/random/torch.manual_seed) is applied from manual_seed; _invert_via_displacement_field rejects a missing reference image.

App server hardening & security (konfai-apps/konfai_apps/app_server.py, cli.py)

• Bearer check uses hmac.compare_digest; _require_configured_app restricts /repo_apps_* endpoints and job submission to the configured allowlist (SSRF/exfiltration guard); _configure_server_auth_env honours --token-env and clears KONFAI_API_TOKEN under --auth off.
• Uploads sharing a basename are disambiguated instead of overwritten (save_uploads); grouped uploads (save_upload_groups) and zip-slip-protected dataset extraction (extract_zip_safely); SSE switched from a hard TTL to a : keepalive heartbeat that terminates on real job state; log lines push thread-safely via emit_log (call_soon_threadsafe); q_put_drop_oldest catches QueueEmpty; GPU selection validates requested ids or auto-selects.

App client & CLI (konfai-apps/konfai_apps/app.py, cli.py, app_repository.py)

• Client-side path args resolved against the caller dir before chdir; fine_tune takes a single zipped dataset and forwards --lr; remote __ERROR__ markers raise RuntimeError; _supported_suffix picks the longest registered extension for Volume_i copies; supported-file listing is sorted so cases pair across groups; new konfai-apps download command + LocalAppRepositoryFromHF.download_files.

Packaging (pyproject.toml, per-app setup.py, .github/workflows/publish.yml)

• packages.find include changed to ["konfai", "konfai.*"] with konfai-apps* excluded, so the hyphenated sibling no longer ships in the konfai wheel while PEP 420 namespace subpackages (konfai.models.*) still do.
• Per-app / konfai-apps setup.py pin install_requires to the matching release (dynamic = [..., "dependencies"]); publish.yml gains a test job that gates build; pixi gains test-apps (folded into check); the local AUDIT.md is removed from the repo.

Testing

31 new regression tests under tests/unit/ (one per fix area): test_config_dict_roundtrip, test_config_dotted_dict_keys, test_patch_overlap_border, test_crop_transform_shape, test_augmentation_fixes, test_augmentation_label_interpolation, test_data_pipeline_audit, test_dataset_audit, test_ddp_shard_balance, test_get_infos_shape_order, test_hdf5_read_only, test_measure, test_dict_metric_logging, test_loss_weight_scheduling, test_perceptual_loss_targets, test_accepts_init_dispatch, test_batchnorm_init, test_named_forward_sibling, test_network_ddp_fixes, test_unet_attention, test_experimental_models, test_trainer_checkpoints, test_ema_convention, test_resume_https_checkpoint, test_safe_torch_load, test_evaluator_statistics, test_itk_transforms, test_interactive_session_guard, test_runtime_progress_ddp, test_transform_fixes, and test_wheel_packaging; plus updates to test_audit_fixes, test_early_stopping, and test_runtime_guards.

New konfai-apps/tests/unit/ suites: test_app_server_helpers, test_finetune_lr_override, test_repo_apps_allowlist, test_server_auth_env, test_upload_basename_collision, plus substantial additions to test_app_runtime.

CI: publish.yml runs pytest tests and pytest konfai-apps/tests as a test job that build depends on. The konfai-apps suite is a separate package and is not covered by pixi run test — run it via the new test-apps task (folded into check).

Review notes

• Numerics-affecting changes to double-check: the EMA blend convention flip in _avg_fn, BatchNorm gamma init 0.0 → 1.0, the Padding inverse-origin axis fix, Translate grid-coordinate normalization, and the get_infos size reversal for 2-D/4-D. These correct real bugs but will change numeric output for pipelines/checkpoints that depended on the old behaviour.
• Behaviour changes: selecting DDPM or 3-D VoxelMorph now raises NotImplementedError (previously crashed opaquely); https:// remote checkpoints load weights_only=True only and will fail if they need the unsafe unpickler (by design); the konfai wheel no longer contains konfai_apps.
• Server wire-format changes: fine_tune now expects a single zipped dataset upload rather than inputs, grouped uploads add *_groups size fields, and endpoints enforce the app allowlist. Clients on the old format need updating.
• Largest surface areas to focus review on: konfai-apps/konfai_apps/app_server.py, konfai/utils/dataset.py, konfai/trainer.py, konfai/network/network.py, and konfai/data/patching.py.