Fix batch inference issues. - #9
Conversation
Greptile SummaryThis PR adds robust large-scale batch inference support across both NV-Segment-CT and NV-Segment-CTMR bundles, replacing the simple single-directory glob with a recursive discovery utility ( Confidence Score: 4/5Safe to merge with one targeted fix: the hardcoded per-file timeout should be made configurable or raised before the script ships to production cohort processing. The core Python batch utility is well-designed (atomic cache, MONAI path prediction, clean multi-GPU coordination). The config changes are minimal and correct. The shell script issues (timeout, ERR trap, stderr merge) are real but won't cause data loss — failures are logged and the resume mechanism lets users retry. NV-Segment-CTMR/brain_t1_preprocess/run_brain_segmentation.sh — hardcoded timeout, ERR-trap leak, stderr/stdout merge. NV-Segment-CT/scripts/batch_inference_utils.py and its CTMR counterpart — input_suffix not forwarded. Important Files Changed
Sequence DiagramsequenceDiagram
participant U as User / torchrun
participant R0 as Rank 0 Process
participant RN as Rank N Process (N>0)
participant FS as Filesystem / Cache
U->>R0: torchrun ... monai.bundle run
U->>RN: torchrun ... monai.bundle run
Note over R0,RN: initialize: dist.init_process_group, set_device, set_determinism, checkpointloader
R0->>FS: Recursive glob **/*.nii.gz under input_dir
R0->>FS: Filter skip_dir_names / skip_dir_prefixes
R0->>FS: Check expected_output_path exists (resume skip)
R0->>FS: Write /tmp/nvseg_batch_input_hash.json (atomic rename)
RN->>FS: Poll cache file (up to batch_cache_wait_sec=120s)
FS-->>RN: Return cached path list
alt Nothing to run (all outputs exist)
R0->>R0: SystemExit(0)
RN->>RN: SystemExit(0)
else Files to process
R0->>R0: Build input_dicts → DistributedSampler
RN->>RN: Build input_dicts → DistributedSampler
R0->>FS: SaveImaged → output_dir/.../stem/stem_trans.nii.gz
RN->>FS: SaveImaged → output_dir/.../stem/stem_trans.nii.gz
end
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| def collect_input_paths( | ||
| input_root: Path, | ||
| pattern: str = "**/*.nii.gz", | ||
| *, | ||
| skip_dir_names: list[str] | None = None, | ||
| skip_dir_prefixes: list[str] | None = None, | ||
| ) -> list[Path]: | ||
| input_root = input_root.resolve() | ||
| out: list[Path] = [] | ||
| for p in sorted(input_root.glob(pattern)): | ||
| if not p.is_file(): | ||
| continue | ||
| if should_skip_path_by_parent_rules( | ||
| p, skip_dir_names=skip_dir_names, skip_dir_prefixes=skip_dir_prefixes | ||
| ): | ||
| continue | ||
| out.append(p) | ||
| return out |
There was a problem hiding this comment.
input_suffix config value is silently ignored
collect_input_paths hardcodes the glob pattern "**/*.nii.gz", which means the input_suffix key in configs/batch_inference.json (set to "*.nii.gz") has no effect on batch discovery. Two consequences:
- Users who change
input_suffixto, e.g.,"*.nii"will not see plain (non-gzipped) NIfTI files discovered—build_input_listnever passes the suffix through. - The same silent mis-match applies to the identical file in
NV-Segment-CTMR/scripts/batch_inference_utils.py.
Either accept and forward input_suffix (or a pattern parameter) from build_input_list down to collect_input_paths, or document prominently in batch_inference.json that input_suffix is unused for the batch utility and only **/*.nii.gz is ever matched.
| ((fail_count++)) | ||
| process_result=1 | ||
| fi | ||
| else | ||
| # Fallback: run without timeout wrapper | ||
| # Note: Without timeout command, we can't enforce the 5-minute limit | ||
| # but we can still catch failures |
There was a problem hiding this comment.
trap cleanup_on_error ERR is set inside process_single_file. The matching trap - ERR is placed at the very end of the success path (after all steps complete). However, when log_failure is called and return 1 exits the function early, trap - ERR is never reached. In Bash, trap is global to the shell (not scoped to functions), so the cleanup_on_error trap remains active in the calling shell after process_single_file returns with an error code.
In process_folder, subsequent arithmetic expansions like ((fail_count++)) evaluate to exit code 1 when the counter is 0, which would unexpectedly fire cleanup_on_error against the temp directory of the just-failed case (likely already removed, so harmless, but still unintended).
Reset the trap before every early return:
log_failure() {
local reason="$1"
echo -e "${RED}$reason: $input_file${NC}" >&2
if [[ -n "$failed_log" ]]; then
echo "$input_file" >> "$failed_log"
fi
cleanup_on_error
trap - ERR # reset here so the caller's shell is clean
return 1
}1 participant
Added supports for inferencing large scale datasets with batch processing.