Bug in `src/utils/inference_utils.py`

[seanlaw]

Currently, there appears to be a bug in the following function:

GRID/src/utils/inference_utils.py

Lines 226 to 257 in 4d83873

	def on_predict_end(
	self,
	trainer: Trainer,
	pl_module: LightningModule,
	) -> None:
	super().on_predict_end(trainer, pl_module)
	if self.should_merge_files_on_main:
	# if we use multiple workers, we need to wait for all of them to finish writing
	# before merging the files
	if trainer.global_rank != None:
	torch.distributed.barrier()
	if self.global_rank == 0:
	log.info("Merging pickle files on main process.")
	self._merge_files()
	# other processes can continue after merging
	if trainer.global_rank != None:
	torch.distributed.barrier()
	# conducting post-processing functions on the files
	for process_func in self.post_processing_functions:
	all_files = [f for f in os.listdir(self.output_dir)]
	for file in all_files:
	file_path = os.path.join(self.output_dir, file)
	if process_func.get("main_only", False):
	if self.global_rank == 0:
	process_func["function"](file_path)
	else:
	process_func["function"](file_path)
	if trainer.global_rank != None:
	torch.distributed.barrier()

as, during the initial embedding step, it warns (and eventually fails) with Default process group has not been initialized, please make sure to call init_process_group., Retrying src.utils.inference_utils:on_predict_end in 48 seconds... and then followed by Failed for the last time: Default process group has not been initialized, please make sure to call init_process_group. I believe the fix is to replace the function with something like (see two sections with # <-- Added this annotation):

    def on_predict_end(
        self,
        trainer: Trainer,
        pl_module: LightningModule,
    ) -> None:
        super().on_predict_end(trainer, pl_module)

        if not torch.distributed.is_initialized():  # <-- Added this
            torch.distributed.init_process_group(
                backend="nccl",
                rank=trainer.global_rank,
                world_size=trainer.world_size,
                device_id=torch.device("cuda", trainer.global_rank)
            )

        if self.should_merge_files_on_main:
            # if we use multiple workers, we need to wait for all of them to finish writing
            # before merging the files
            if trainer.global_rank != None:
                torch.distributed.barrier()
            if self.global_rank == 0:
                log.info("Merging pickle files on main process.")
                self._merge_files()

            # other processes can continue after merging
            if trainer.global_rank != None:
                torch.distributed.barrier()

        # conducting post-processing functions on the files
        for process_func in self.post_processing_functions:
            all_files = [f for f in os.listdir(self.output_dir)]
            for file in all_files:
                file_path = os.path.join(self.output_dir, file)
                if process_func.get("main_only", False):
                    if self.global_rank == 0:
                        process_func["function"](file_path)
                else:
                    process_func["function"](file_path)
                if trainer.global_rank != None:
                    torch.distributed.barrier()

        torch.distributed.destroy_process_group()  # <-- Added this

I hope that helps others! Note that the backend="nccl" parameter is probably bad if no GPUs are available. Maybe @jumxglhf can offer some guidance or come up with a permanent fix.