Race condition in nnUNetTrainer.do_split(): concurrent trainings can read a partially written splits_final.json

There's a race condition in nnUNetTrainer.do_split() in that multiple processes can attempt to write to and read from the same splits file. If one happens to read a partially written one, it can crash with errors such as json.decoder.JSONDecodeError: Expecting ',' delimiter: line 318 column 25 (char 8215) raised from load_json(splits_file) in do_split().

Version: Current master commit

Reproduction: Happens when -num-gpus is used for the first training run of the dataset (when no splits file is present yet) as in: nnUNetv2_train 220 3d_fullres 0 -p nnUNetResEncUNetLPlans -tr nnUNetTrainer -num_gpus 2. Presumably can also happen when running multiple folds at the same time with one GPU each, as is the recommended pattern. Unfortunately, reproduction is unreliable due to the stochastic nature of the problem.

Frequently leads to multiple log lines written by self.print_to_log_file("Creating new 5-fold cross-validation split...") indicating that a race happened and occasionally leads to crashes due to the JSON decoder failing after reading partial read: json.decoder.JSONDecodeError: Expecting ',' delimiter: line 318 column 25 (char 8215) raised from load_json(splits_file) in do_split().

My proposed fix is creating a temporary file first, per process, and then renaming it atomically as in:

            if not isfile(splits_file):
                self.print_to_log_file("Creating new 5-fold cross-validation split...")
                all_keys_sorted = list(np.sort(list(dataset.identifiers)))
                splits = generate_crossval_split(all_keys_sorted, seed=12345, n_splits=5)

                # Concurrent trainings share this file: DDP ranks, or one process per fold.
                # Publish by rename so no reader sees a partial write.
                tmp_file = f"{splits_file}.tmp.{uuid4().hex}"
                save_json(splits, tmp_file)
                os.replace(tmp_file, splits_file)

This would remove the need to run a single-process/single-gpu training run first to safely create the splits file before running any parallel setup. (Unless in cases where this is still necessary due to numpy arrays being unpacked on the first run).