Multi-GPU prediction gather callback.

[konrad-karanowski]

Hello, I have a super-resolution model that operates on patches. In the validation step, I would like to reassembly selected predictions and save, so I can run some external program to get metrics. It works fine on single GPU. but on multiple GPU (DDP) I have deadlock after this step. Do you have any ideas how to even approach this program?

import json
import os
import shutil
from typing import Any, Dict, Optional

import numpy as np
import torch
from lightning.pytorch import Callback, LightningModule, Trainer

from src.evaluation_utils.reassembly import (
    load_block_indices,
    reassemble_from_arrays,
    _write_image,
)
from src.utils import RankedLogger

logger = RankedLogger(__name__)


def _parse_block_idx(input_path: str) -> int:
    """Extract block index from a block filename: ``image_42.png`` → 42."""
    stem = os.path.splitext(os.path.basename(input_path))[0]
    return int(stem.rsplit("_", 1)[-1])


def _copy_image(src: str, dst: str) -> None:
    try:
        shutil.copyfile(src, dst)
    except Exception as e:
        logger.warning(f"Could not copy {src} -> {dst}: {e}")


class FullImageValidationCallback(Callback):

    def __init__(
        self,
        every_n_epochs: int = 5,
        collection_dir: str = "",
        block_indices_root: str = "",
        external_metric_cmd: Optional[str] = None,  # reserved, currently unused
        save_input_image: bool = False,
        save_target_image: bool = False,
        dir_name: str = "images",
    ) -> None:
        super().__init__()
        self.every_n_epochs = every_n_epochs
        self.collection_dir = collection_dir
        self.block_indices_root = block_indices_root
        self.save_input_image = save_input_image
        self.save_target_image = save_target_image
        self.dir_name = dir_name

        self._predictions: Dict[str, Dict[int, np.ndarray]] = {}


    def on_validation_epoch_start(
        self, trainer: Trainer, pl_module: LightningModule
    ) -> None:
        self._predictions = {}

    def on_validation_batch_end(
        self,
        trainer: Trainer,
        pl_module: LightningModule,
        outputs: Any,
        batch: Any,
        batch_idx: int,
        dataloader_idx: int = 0,
    ) -> None:
        if not self._should_run(trainer) or not outputs:
            return

        full_val_image_info = self._full_val_image_info(trainer)
        if not full_val_image_info:
            return

        preds = outputs.get("preds") 
        image_ids = outputs.get("image_id") 
        input_paths = outputs.get("input_path") 
        if preds is None or image_ids is None or input_paths is None:
            return

        for i, (img_id, input_path) in enumerate(zip(image_ids, input_paths)):
            if img_id not in full_val_image_info:
                continue
            block_idx = _parse_block_idx(input_path)
            if img_id not in self._predictions:
                self._predictions[img_id] = {}
            
            self._predictions[img_id][block_idx] = (
                preds[i].detach().float().cpu().numpy()
            )

    def on_validation_epoch_end(
        self, trainer: Trainer, pl_module: LightningModule
    ) -> None:
        if not self._should_run(trainer):
            return

        full_val_image_info = self._full_val_image_info(trainer)
        if not full_val_image_info:
            return

        predictions = self._gather_predictions(trainer)

        if not trainer.is_global_zero:
            return

        output_dir = os.path.join(
            trainer.log_dir or ".", self.dir_name, f"epoch_{trainer.current_epoch:04d}"
        )
        os.makedirs(output_dir, exist_ok=True)

        for img_id, info in full_val_image_info.items():
            img_preds = predictions.get(img_id)
            if not img_preds:
                logger.warning(f"Image-{img_id}: no predictions collected — skipping reassembly.")
                continue
            if len(img_preds) != info["num_blocks"]:
                logger.warning(
                    f"Image-{img_id}: expected {info['num_blocks']} blocks "
                    f"but collected {len(img_preds)} — reassembly may be incomplete."
                )
            self._reassemble_entry(img_id, info, img_preds, output_dir)

    def _should_run(self, trainer: Trainer) -> bool:
        return not trainer.sanity_checking and (trainer.current_epoch + 1) % self.every_n_epochs == 0

    def _full_val_image_info(self, trainer: Trainer) -> Optional[dict]:
        dm = trainer.datamodule
        if dm is None or not hasattr(dm, "full_val_image_info"):
            return None
        return dm.full_val_image_info or {}

    def _gather_predictions(
        self, trainer: Trainer
    ) -> Dict[str, Dict[int, np.ndarray]]:
        if trainer.world_size <= 1:
            return self._predictions

        container = [None] * trainer.world_size
        torch.distributed.all_gather_object(container, self._predictions)

        merged: Dict[str, Dict[int, np.ndarray]] = {}
        for rank_preds in container:
            for img_id, blocks in rank_preds.items():
                if img_id not in merged:
                    merged[img_id] = {}
                merged[img_id].update(blocks)
        return merged

    def _reassemble_entry(
        self,
        img_id: str,
        info: dict,
        img_preds: Dict[int, np.ndarray],
        output_dir: str,
    ) -> None:
        # 1. Load precomputed block indices
        block_indices_path = os.path.join(self.block_indices_root, f"{img_id}.pt")
        if not os.path.exists(block_indices_path):
            logger.error(f"Image-{img_id}: block indices not found at {block_indices_path}")
            return

        all_block_indices = load_block_indices(block_indices_path)

        # 2. Output shape from manifest
        manifest_path = os.path.join(
            self.collection_dir, f"{img_id}", "manifest.json"
        )
        with open(manifest_path, "r") as f:
            manifest = json.load(f)
        output_shape = tuple(manifest["processed_image_shape"])

        # 3. Reassemble from in-memory arrays
        try:
            sorted_idxs = sorted(img_preds.keys())
            blocks = [img_preds[i] for i in sorted_idxs]
            indices = [all_block_indices[i] for i in sorted_idxs]

            assembled = reassemble_from_arrays(blocks, indices, output_shape)

            pred_out = os.path.join(output_dir, f"{img_id}_pred.png")
                        _write_image(pred_out, assembled)
            logger.info(f"Image-{img_id}: prediction image -> {pred_out}")

        except Exception:
            import traceback
            logger.error(f"Image-{img_id}: reassembly failed\n{traceback.format_exc()}")

        # 4. Optionally copy the reference images
        if self.save_input_image:
            _copy_image(
                info["input_image_path"],
                os.path.join(output_dir, f"{img_id}_input.png"),
            )
        if self.save_target_image:
            _copy_image(
                info["target_image_path"],
                os.path.join(output_dir, f"{img_id}_target.png"),
            )

[iamgooold]

Hey @konrad-karanowski, this deadlock happens because all_gather_object is a collective — all ranks must call it, even if they have no data.

Use Lightning's strategy.all_gather instead. It's safer and handles the empty-rank case:

def on_validation_epoch_end(self, trainer: Trainer, pl_module: LightningModule) -> None:
if not self._should_run(trainer):
return

# Gather across all ranks
local_preds = self._predictions
gathered = trainer.strategy.all_gather(local_preds) # List[dict] from all ranks

if not trainer.is_global_zero:
    return
    
# Merge on rank 0 only
merged = {}
for rank_dict in gathered:
    for img_id, blocks in rank_dict.items():
        merged.setdefault(img_id, {}).update(blocks)

#... rest of your reassembly code using 'merged'

Let me know if this fixes the deadlock. If it works, could you mark this as the Answer? Helps others with DDP.

[konrad-karanowski]

Solved my case. Thank you so much!