def validate_data_is_mpi_safe()

in summarize_from_feedback/utils/dist_utils.py [0:0]

• 19 lines of code
• 8 McCabe index (conditional complexity)

def validate_data_is_mpi_safe(data, name="<unknown>"):
    known_safe_types = (int, float, str, bool, type(None), np.ndarray, np.generic)

    if isinstance(data, known_safe_types):
        pass
    elif isinstance(data, torch.Tensor):
        if data.is_cuda:
            raise ValueError(
                f"Data name={name} was a cuda tensor. MPI cannot handle CUDA tensors"
                f" as they result in unexpected CUDA OOMs."
            )
    elif isinstance(data, dict):
        for k, v in data.items():
            validate_data_is_mpi_safe(k)
            validate_data_is_mpi_safe(v, name=k)
    elif isinstance(data, Iterable):
        for item in data:
            validate_data_is_mpi_safe(item)
    else:
        raise ValueError(f"Data name={name} had unsupported type: {type(data)}")