# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils/10_pipeline_parallelism.ipynb.

# %% auto 0
__all__ = ['get_device', 'move_to_device', 'parallelize_module', 'unparallelize_module', 'parallelize_pipeline',
           'unparallelize_pipeline']

# %% ../../nbs/utils/10_pipeline_parallelism.ipynb 1
from collections import defaultdict

import torch
from torch import nn

# %% ../../nbs/utils/10_pipeline_parallelism.ipynb 4
def get_device(device: torch.device | str) -> torch.device:
    """Convert to torch.device object."""
    if isinstance(device, str):
        return torch.device(device)
    return device

# %% ../../nbs/utils/10_pipeline_parallelism.ipynb 5
def move_to_device(data, device):
    """
    Move data to the specified device.

    Args:
        data: The data to move.
        device: The device to move the data to.

    Returns:
        The data moved to the specified device.
    """
    device = get_device(device)
    if isinstance(data, (list, tuple, set)):
        return type(data)(move_to_device(d, device) for d in data)
    elif isinstance(data, dict):
        return {k: move_to_device(v, device) for k, v in data.items()}
    elif isinstance(data, torch.Tensor):
        return data.to(device)
    else:
        return data

# %% ../../nbs/utils/10_pipeline_parallelism.ipynb 8
def parallelize_module(
    module: nn.Module, processing_device: torch.device | str, output_device: torch.device | str | None = None
):
    # Get devices
    processing_device = get_device(processing_device)
    if output_device is None:
        output_device = processing_device
    output_device = get_device(output_device)

    # Move module to processing device
    module = module.to(processing_device)

    # Store module forward function without parallization
    module._forward_without_pipeline_parallelization = module.forward

    # Update with parallelized forward function
    def forward_with_pipeline_parallelization(*args, **kwargs):
        """Forward function that moves inputs to the processing device and outputs to the output device. To view the
        actual forward logic, use the `_forward_without_pipeline_parallelization` function."""
        args = move_to_device(args, processing_device)
        kwargs = move_to_device(kwargs, processing_device)
        output = module._forward_without_pipeline_parallelization(*args, **kwargs)
        return move_to_device(output, output_device)

    module.forward_with_pipeline_parallelization = forward_with_pipeline_parallelization
    module.forward = forward_with_pipeline_parallelization

    # Store module extra_repr function without parallization
    module._extra_repr_without_pipeline_parallelization = module.extra_repr

    # Update with extra_repr function
    def extra_repr_with_pipeline_parallelization(*args, **kwargs):
        """extra_repr function that also shows the processing device and the output device. To view the actual
        extra_repr function, use the `_extra_repr_without_pipeline_parallelization` function."""
        return (
            f"processing_device={processing_device}, output_device={output_device}"
            + module._extra_repr_without_pipeline_parallelization(*args, **kwargs)
        )

    module.extra_repr_with_pipeline_parallelization = extra_repr_with_pipeline_parallelization
    module.extra_repr = extra_repr_with_pipeline_parallelization

    # Add flag to module
    module._is_pipeline_parallelized = True

    return module

# %% ../../nbs/utils/10_pipeline_parallelism.ipynb 10
def unparallelize_module(module: nn.Module, common_device: torch.device | str | None = None):
    """Unparallelize a module that was parallelized with `parallelize_module`. This will move the module to the common
    device and restore the original forward and extra_repr functions."""
    # If common_device is provided, move the module to that device
    if common_device is not None:
        common_device = get_device(common_device)
        module = module.to(common_device)

    # Check if module is parallelized
    is_pipeline_parallelized = getattr(module, "_is_pipeline_parallelized", False)
    if not is_pipeline_parallelized:
        return module

    # Restore model's original forward function and delete extra functions
    module.forward = module._forward_without_pipeline_parallelization
    del module.forward_with_pipeline_parallelization, module._forward_without_pipeline_parallelization

    # Restore model's original extra_repr function and delete extra functions
    module.extra_repr = module._extra_repr_without_pipeline_parallelization
    del module.extra_repr_with_pipeline_parallelization, module._extra_repr_without_pipeline_parallelization

    # Update flag (although it can be deleted too, it's useful for debugging)
    module._is_pipeline_parallelized = False

    return module

# %% ../../nbs/utils/10_pipeline_parallelism.ipynb 12
def parallelize_pipeline(
    model: nn.Module, module_to_device: dict[str, torch.device | str | list[torch.device | str]]
) -> nn.Module:
    """
    Parallelize a model across multiple devices. It is recommended to pass a root device too i.e. {"": DEVICE, ...} to
    avoid errors in the calculations of the outermost modules.

    Args:
        model: The model to parallelize. An unparallelized model is expected, otherwise you will face recursion errors.
        module_to_device: A dictionary mapping module names to devices. Keys are modules names with nested modules
            separated by dots (e.g., "module.submodule"). Note that the parallelism is performed using Level Order
            Traversal (i.e. BFS) of the model. Therefore the device of the deepest module in the dictionary will be
            overwritten even if it's parent is also specified in the dictionary. Value is either a device or a 2-tuple
            of devices. The first device is the processing device, and the second device is the output device.

    Returns:
        The parallelized pipeline.
    """
    # Convert all dictionary values to lists
    for key, value in module_to_device.items():
        if not isinstance(value, list):
            module_to_device[key] = [value, value]

    # Parallelize the model if required
    if "" in module_to_device:
        # Identify the devices
        devices = module_to_device.pop("")
        if isinstance(devices, list):
            processing_device, output_device = devices
        else:
            processing_device, output_device = devices, devices

        # Parallelize
        parallelize_module(model, processing_device, output_device)

        # Update the output device of the rest of they modules in the dictionary so that they are able to work together
        for key in module_to_device:
            module_to_device[key][1] = processing_device

    # Identify children submodules and their devices
    children_submodules_to_device: dict[str, dict[str, torch.device | str]] = defaultdict(dict)
    for module_name, devices in sorted(module_to_device.items()):
        # Get the module's name and ensure it is present in the model
        module_name_split = module_name.split(".")
        child_name = module_name_split[0]
        if not hasattr(model, child_name):
            raise ValueError(f"Module {child_name} not found in the model.")

        # Add to children_submodules_to_device
        children_submodules_to_device[child_name][".".join(module_name_split[1:])] = devices

    # Iterate over the submodules and assign them to the appropriate devices
    for child_name, submodules_to_device in children_submodules_to_device.items():
        module = getattr(model, child_name)
        parallelize_pipeline(module, submodules_to_device)

    return model

# %% ../../nbs/utils/10_pipeline_parallelism.ipynb 14
def unparallelize_pipeline(model: nn.Module, common_device: torch.device | str | None = None) -> nn.Module:
    """
    Unparallelize a model that was parallelized with `parallelize_pipeline`. This will remove any pipeline
    parallelization and move the model to the common device.

    Args:
        model: The model to unparallelize.
        common_device: The device to move the model to. If None, the model will not be moved.

    Returns:
        The model after unparallelization and device transfer.
    """
    # If common_device is provided, get the device object
    if common_device is not None:
        common_device = get_device(common_device)

    # Unparallelize the model
    model = unparallelize_module(model, common_device)

    # Recurse over all submodules and unparallelize them
    for module in model.children():
        # Unparallelize the module
        unparallelize_pipeline(module, common_device)

    return model
