Socket timeout during metric.aggregate()

[csheaff]

Hello,

Currently I am using Monai with AWS SageMaker DataParallel. I have a validation loop that uses something like:

import smdistributed.dataparallel.torch.distributed as dist
dist.init_process_group()

def get_model(args):
    if dist.get_rank() == 0:
        logger.info('Defining network')    
    model = monai.networks.nets.SwinUNETR(
        img_size=(96, 96, 32),
        in_channels=1,
        out_channels=1,
        feature_size=12, # must be divisible by 12
        use_checkpoint=True,
    ).to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-5)
    model = DistributedDataParallel(model)
    torch.cuda.set_device(args.local_rank)
    model.cuda(args.local_rank)
    return model, optimizer, device

train_loader, val_loader = get_data_loaders(args=args)
model, optimizer, device = get_model(args)

dice_metric = monai.metrics.DiceMetric(include_background=True, reduction="mean", get_not_nans=False)
with torch.no_grad():
    for (val_inputs, val_labels) in val_loader:
         val_outputs = monai.inferers.sliding_window_inference(val_inputs, roi_size, batch_size, model)
         dice_metric(y_pred=val_outputs, y=val_labels)
    mean_dice_val = dice_metric.aggregate().item()
    dice_metric.reset()

I've noticed that under a few different conditions i get the following error:

File "train.py", line 167, in validation
 mean_dice_val = dice_metric.aggregate().item()
 File "/opt/conda/lib/python3.9/site-packages/monai/metrics/meandice.py", line 97, in aggregate
 data = self.get_buffer()
 File "/opt/conda/lib/python3.9/site-packages/monai/metrics/metric.py", line 275, in get_buffer
 self._sync()
 File "/opt/conda/lib/python3.9/site-packages/monai/metrics/metric.py", line 255, in _sync
 raise TypeError(f"{e}. unable to sync buffer contents: {self._buffers}.") from e
 TypeError: Socket Timeout. unable to sync buffer contents: [[tensor([0.0054], device='cuda:0'), tensor([0.0001], device='cuda:0'), tensor([0.0666], device='cuda:0'), tensor([0.0334], device='cuda:0'), tensor([0.0430], device='cuda:0'), tensor([0.0112], device='cuda:0'), tensor([0.0115], device='cuda:0'), tensor([0.0306], device='cuda:0')]].

For instance, this occurs if my batch_size = 1 rather than 2 or 4. I've also thrown in some debugging statements that write info to files which have caused this.

Question: is there some sort of timeout value that I can adjust to prevent this error?

I'm not sure this is necessarily a Monai issue, but i figured this would be the best place to start. I see something similar discussed in #6120, but i'm not sure it's the same issue and the solutions proposed do not appear to help me.

Any assistance would be greatly appreciated.

Monai v1.1.0
Pytorch version: 1.13.1+cu117
Pytorch Ignite version: 0.4.11
Python 3.9

[KumoLiu]

Hi @csheaff, I just noticed that the error will only happened when batch_size=1. And I'm not sure if you are only using one data, how to using DDP, maybe just remove DDP and if you are using more data then you could add DDP.
Hope it can help you, thanks!

[csheaff]

Thanks for the response @KumoLiu . Yes, I can just use batch size > 1, however I believe this has happened under other scenarios as well (e.g. writing debug statements to logs on S3). I will try to reproduce.

If there is someway to resolve the issue directly that would be great.

[csheaff]

A condition where this timeout occurs: monitoring model weights and biases using

import smdebug.pytorch as smd
hook=smd.get_hook(create_if_not_exists=True)
hook.register_module(model)

These are getting saved to S3 as TensorBoard metrics. I wonder if the network activity is somehow blocking the retrieval of the monai metric.

update: reducing the rate at which i'm logging data during the validation step appears to resolve this.