ezpz.dist⚓︎

def all_reduce(
    obj: Any,
    op: Optional[MPI.Op | torch.distributed.reduce_op] = None,  # type:ignore
    implementation: Optional[str] = None,
) -> Any:
    """All-reduce ``obj`` across all ranks using MPI.

    Args:
        obj: Picklable payload to reduce.
        op: Reduction operation; defaults to ``MPI.SUM``.
        implementation: Override to ``\"torch\"`` to use torch.distributed.

    Returns:
        The reduced value.

    Examples:
        >>> loss = 1.0 + get_rank()
        >>> total_loss = all_reduce(loss)  # sum across ranks
        >>> mean_loss = (
        ...     all_reduce(loss, implementation="torch") / get_world_size()
        ... )
    """
    if implementation is None or implementation.lower() == "mpi":
        op = MPI.SUM if op is None else op
        assert op is not None
        try:
            return MPI.COMM_WORLD.allreduce(obj, op=op)
        except Exception as exc:
            if get_rank() == 0:
                logger.warning(
                    "Unable to all-reduce with MPI, returning original object"
                )
                logger.exception(exc)
            # return obj
            raise exc
    elif implementation.lower() in {"torch", "pytorch", "pt"}:
        import torch.distributed as dist

        op = dist.ReduceOp.SUM if op is None else op  # type:ignore
        assert op is not None
        tensor = torch.tensor(obj)
        dist.all_reduce(tensor, op=op)  # type:ignore
        return tensor.item()
    else:
        raise ValueError(
            f"Unsupported all-reduce implementation: {implementation}"
        )

def barrier(
    device: Optional[torch.device | int | str] = None,
    group: (
        torch.distributed.ProcessGroup | None  # type:ignore
    ) = torch.distributed.GroupMember.WORLD,  # type:ignore
    async_op: bool = False,
    device_ids: str | Iterable | None = None,
    implementation: Optional[str] = None,
) -> torch.distributed.Work | None:  # type:ignore
    """
    Barrier for all processes in the group.

    This collective blocks processes until the whole group enters this function,
    if async_op is False, or if async work handle is called on wait().

    Args:
        device (torch.device | int | str, optional): The device to synchronize.
            If None, the default device will be used. Defaults to None.
        group (torch.distributed.ProcessGroup | None, optional): The process group to work on.
            If None, the default process group (WORLD) will be used.
            Defaults to torch.distributed.GroupMember.WORLD.
        async_op (bool, optional): If True, the barrier will be asynchronous.
        device_ids (str | Iterable | None, optional): The device IDs to synchronize.

    Returns:
        torch.distributed.Work | None: If async_op is True, returns a work handle.
            If async_op is False, returns None.

    Examples:
        >>> barrier()  # wait for all ranks
        >>> handle = barrier(async_op=True)  # launch async and wait later
        >>> if handle is not None:
        ...     handle.wait()
    """
    if implementation is not None and implementation.lower() not in {
        "mpi",
        "torch",
    }:
        raise ValueError(
            f"Unsupported barrier implementation: {implementation}"
        )
    if implementation is None or implementation.lower() == "mpi":
        try:
            MPI.COMM_WORLD.barrier()
        except Exception:
            if get_rank() == 0:
                logger.warning(
                    "Unable to use `MPI.COMM_WORLD.barrier` "
                    "for this process group. "
                    "Falling back to `torch.distributed` barrier."
                )
            _torch_barrier(group, async_op=async_op, device_ids=device_ids)
    if implementation is not None and implementation.lower() == "torch":
        try:
            _torch_barrier(group, async_op=async_op, device_ids=device_ids)
        except Exception:
            if get_rank() == 0:
                logger.warning(
                    "Unable to use `ezpz.dist._torch_barrier` "
                    "for this process group. "
                    "Falling back to `MPI.COMM_WORLD.barrier()` barrier."
                )
            MPI.COMM_WORLD.barrier()

def broadcast(
    obj: Any,
    root: int = 0,
) -> Any:
    """Broadcast ``obj`` from ``root`` to all ranks using MPI.

    Args:
        obj: Picklable payload to share.
        root: Rank that originates the value.

    Returns:
        The broadcast payload.

    Examples:
        >>> value = 42 if get_rank() == 0 else None
        >>> shared = broadcast(value, root=0)
        >>> assert shared == 42
    """
    try:
        return MPI.COMM_WORLD.bcast(obj, root=root)
    except Exception as exc:
        if get_rank() == 0:
            logger.warning(
                "Unable to broadcast with MPI, returning original object"
            )
            logger.exception(exc)
        # return obj
        raise exc

def check(
    framework: str = "pytorch",
    backend: str = "deepspeed",
    port: int | str = "5432",
):
    """Check if the framework is installed and working"""
    from ezpz.configs import FRAMEWORKS

    if framework in FRAMEWORKS["pytorch"]:
        _ = setup_torch(
            backend=backend,
            port=str(port),
        )
    elif framework in FRAMEWORKS["tensorflow"]:
        _ = setup_tensorflow()
    else:
        raise ValueError(f"Unable to parse framework: {framework}")

def cleanup() -> None:
    """
    Cleanup the distributed environment.
    This function destroys the process group if it is initialized.

    Examples:
        >>> cleanup()
    """
    if wandb is not None and (run := getattr(wandb, "run")) is not None:
        logger.info(f"wandb.run=[{run.name}]({run.url})")
    if torch.distributed.is_initialized():  # type:ignore
        torch.distributed.destroy_process_group()  # type:ignore

def get_cpus_per_node() -> int:
    """Get the number of CPUs per node.

    Returns:
        Number of logical CPUs on the local node.

    Examples:
        >>> get_cpus_per_node() > 0
        True
    """
    from sh import getconf as sh_getconf  # type:ignore noqa

    return int(sh_getconf("_NPROCESSORS_ONLN").rstrip("\n"))

def get_device(
    type: Optional[str] = None, as_torch_device: Optional[bool] = None
) -> str | torch.device:
    """Alias for `get_torch_device`.

    Examples:
        >>> get_device()
        'cuda'
        >>> get_device(as_torch_device=True)
        device(type='cuda', index=0)
    """
    return get_torch_device(device_type=type, as_torch_device=as_torch_device)

def get_device_properties(
    device: Optional[str | torch.device | int] = None,
) -> dict | None:
    """Get the properties of the specified device.

    Args:
        device (str | torch.Device | int, optional): The device to get properties for.
            If None, the current device will be used. Defaults to None.

    Returns:
        dict: A dictionary containing the device properties.
    """
    # if device is None:
    #     device = get_torch_device()
    # if isinstance(device, str):
    #     device = torch.device(device)
    # if isinstance(device, int):
    #     device = torch.device(f"cuda:{device}")
    device_type: str = ezpz.get_torch_device_type()
    if device is None:
        device = torch.device(device_type)
    if isinstance(device, str):
        device = torch.device(device)
    if isinstance(device, int):
        device = torch.device(f"{device_type}:{device}")
    # else:
    #     raise ValueError(f"Unsupported device type: {device.type}")

    # return {
    #     "name": props.name,
    #     "total_memory": props.total_memory,
    #     "multi_processor_count": props.multi_processor_count,
    #     "major": props.major,
    #     "minor": props.minor,
    #     "device_id": device.index,
    # }
    # elif device.type == "xpu":
    #     props = torch.xpu.get_device_properties(device)
    #     for key, value in props.__dict__.items():
    #         logger.debug(f"xpu prop: {key} => {value}")
    #     return {
    #         "name": props.name,
    #         "total_memory": props.total_memory,
    #         "multi_processor_count": props.multi_processor_count,
    #         "major": props.major,
    #         "minor": props.minor,
    #         "device_id": device.index,
    #     }

    props = None
    if device_type == "cuda":
        props = torch.cuda.get_device_properties(device)
    if device_type == "xpu":
        props = torch.xpu.get_device_properties(device)

    if props is not None:
        return {
            k: getattr(props, k, None)
            for k in [i for i in props.__dir__() if not i.startswith("_")]
        }
    return {}

def get_dist_info(
    framework: Optional[str] = None,
    verbose: Optional[bool] = None,
    hostfile: Optional[PathLike] = None,
) -> dict[str, str | int | list]:
    """Get distributed info.

    Args:
        framework (str, optional): Framework to use. Defaults to None.
        verbose (bool, optional): Whether to print the info. Defaults to None.
        hostfile (PathLike, optional): Path to the hostfile. Defaults to None.

    Returns:
        dict: Dictionary containing the distributed info.
    """
    dist_info = _get_dist_info(
        hostfile=hostfile,
        framework=framework,
    )
    if verbose:
        import json

        logger.info(
            f"DistInfo={json.dumps(dist_info, indent=4, sort_keys=True)}"
        )
    return dist_info

def get_free_port() -> int:
    """
    Get a free port on the local machine.

    Returns:
        int: A free port number that can be used for communication.
    """
    sock = socket.socket()
    sock.bind(
        ("127.0.0.1", 0)
    )  # Bind to an available port on the loopback interface
    port = sock.getsockname()[1]
    sock.close()
    return port

def get_gpus_per_node() -> int:
    """Get the number of GPUs per node.

    Returns:
        Number of visible GPU devices on the local node.

    Examples:
        >>> get_gpus_per_node()  # returns 0 on CPU-only machines
        0
    """
    # return torch.cuda.device_count() if torch.cuda.is_available() else (
    #     (
    #         ipex.xpu.device_count() if ipex is not None else (
    #             get_cpus_per_node()
    #         )
    #     )
    # )
    # if _assert:
    #     raise RuntimeError(
    #         'No {X, G}pus found; but _assert specified. Returning !!'
    #     )
    # logger.warning('No {x,g}-pus found, returning' + f'{cpus_per_node}')
    ngpu_per_host = os.environ.get(
        "NGPU_PER_HOST",
        os.environ.get(
            "LOCAL_WORLD_SIZE",
            os.environ.get(
                "PMI_LOCAL_SIZE", os.environ.get("SLURM_NTASKS_PER_NODE")
            ),
        ),
    )
    if ngpu_per_host is not None:
        return int(ngpu_per_host)
    if torch.cuda.is_available():
        return torch.cuda.device_count()
    if torch.xpu.is_available():
        return torch.xpu.device_count()
    if ipex is not None and torch.xpu.is_available():
        return ipex.xpu.device_count()
    if torch.backends.mps.is_available():
        # XXX: Maybe we're running MPI with multiple MPS devices?
        return get_world_size_in_use()
    return 0

def get_hostfile_with_fallback(hostfile: Optional[PathLike] = None) -> Path:
    """Get the hostfile from the environment or create one if it doesn't exist.

    Args:
        hostfile: Optional explicit hostfile path.

    Returns:
        Path to a usable hostfile.

    Examples:
        >>> hostfile = (
        ...     get_hostfile_with_fallback()
        ... )  # uses scheduler env or writes localhost
        >>> Path(hostfile).exists()
        True
    """
    from ezpz.configs import get_scheduler

    scheduler = get_scheduler()
    if scheduler.lower() == "unknown":
        if get_rank() == 0:
            logger.debug("Unknown scheduler")
        hostfile = Path(os.getcwd()).joinpath("hostfile")
        hostfile.touch(exist_ok=True)
        write_localhost_to_hostfile(hostfile=hostfile)
    if scheduler.lower() == "slurm":
        hostfile = make_hostfile_from_slurm_env()
        assert Path(hostfile).is_file()
    if hostfile is None:
        hfp = os.environ.get(
            "PBS_NODEFILE",
            os.environ.get(
                "HOSTFILE",
                None,  # fallback_hostfile.as_posix()
            ),
        )
        if (
            hfp is None or not Path(hfp).is_file()
            # and scheduler == 'PBS'
        ):
            if scheduler == "PBS":
                # hfp = Path(get_pbs_nodefile_from_qstat())
                nodefile = ezpz.pbs.get_pbs_nodefile()
                assert nodefile is not None, (
                    "Unable to get PBS_NODEFILE from `qstat` or `ezpz.pbs`!"
                )
                hfp = Path(nodefile)
            else:
                # create makeshift hostfile containing 'localhost'
                hfp = Path(os.getcwd()).joinpath("hostfile")
                hfp.touch(exist_ok=True)
                write_localhost_to_hostfile(hfp)
    else:
        hfp = Path(hostfile)
    assert hfp is not None and Path(hfp).is_file()
    assert Path(hfp).is_file()
    hostfile = Path(hfp).as_posix()
    # if hfp is not None:
    # hostfile, hosts = get_hosts_from_hostfile(hostfile)
    # hosts = [h.split('.')[0] for h in hosts]
    # if scheduler == 'PBS':
    #     os.environ['PBS_NODEFILE'] = hostfile  # hfp.as_posix()
    hfname = f"{scheduler.upper()}_NODEFILE"
    if hfname not in os.environ:
        os.environ |= {hfname: hostfile}
    # os.environ[f'{scheduler.upper()}_NODEFILE'] = hostfile
    return Path(hfp)

def get_hostname() -> str:
    """Get the hostname of the current machine.

    Returns:
        str: The hostname of the current machine.
    """
    import platform
    import socket

    def _normalize(name: str | None) -> str:
        if not name:
            return "localhost"
        return name.strip().lower()

    try:
        socket_hostname = socket.gethostname()
        if socket_hostname:
            try:
                resolved = socket.gethostbyaddr(socket_hostname)[0]
            except OSError:
                resolved = socket_hostname
            return _normalize(resolved)
    except Exception:
        pass

    env_hostname = os.environ.get("HOSTNAME") or os.environ.get("HOST")
    if env_hostname:
        return _normalize(env_hostname)

    platform_hostname = platform.node()
    if platform_hostname:
        return _normalize(platform_hostname)

    return "localhost"

def get_hosts_from_hostfile(
    hostfile: Optional[str | Path] = None,  # type:ignore[reportDeprecated]
) -> tuple[str, list[str]]:
    """
    Get hosts from the hostfile or environment variables.

    Args:
        hostfile (str | Path, optional): Path to the hostfile. Defaults to None.

    Returns:
        tuple[str, list[str]]: Tuple containing the hostfile path and a list of hosts.

    Examples:
        >>> get_hosts_from_hostfile("/path/to/hostfile")
        ('/path/to/hostfile', ['host1', 'host2', ...])
    """
    # hostfile = '' if hostfile is None else hostfile
    hostname = get_hostname()
    if hostfile is None:
        hostfile = os.environ.get(
            "HOSTFILE",
            os.environ.get(
                "PBS_NODEFILE",
                os.environ.get(
                    "COBALT_NODEFILE",
                    None,
                ),
            ),
        )
    hosts: list[str] = []
    if hostfile is not None and Path(hostfile).is_file():
        if get_rank() == 0:
            logger.debug("Reading hosts from %s", hostfile)
        hpath = Path(hostfile).resolve().absolute()
        with hpath.open("r") as f:
            hosts.extend([h.rstrip("\n") for h in f.readlines()])
        return hpath.as_posix(), hosts
    hosts.append(hostname)
    hostfile_path = Path(hostfile).as_posix() if hostfile else ""
    return hostfile_path, hosts

def get_local_rank() -> int:
    """Return `get_rank() % get_gpus_per_node()`

    Returns:
        int: The local rank of the current process within its node.
            This is calculated as the current rank modulo the number of GPUs per node.

    Examples:
        >>> local_rank = get_local_rank()
        >>> print(f"Local rank of the current process: {local_rank}")
    """
    local_rank = os.environ.get(
        "LOCAL_RANK",
        os.environ.get("PMI_LOCAL_RANK", os.environ.get("SLURM_LOCAL_ID")),
    )
    if local_rank is not None:
        return int(local_rank)
    return int(get_rank() % get_gpus_per_node()) if get_world_size() > 1 else 0

def get_machine(hostname: Optional[str] = None) -> str:
    """Get the machine name from the hostname.

    Args:
        hostname (str, optional): The hostname to check. Defaults to None.

    Returns:
        str: The machine name.

    Examples:
        >>> get_machine("frontier")
        "Frontier"
    """

    if hostname is None:
        try:
            hostname = socket.gethostbyaddr(socket.gethostname())[0]
        except Exception:
            try:
                hostname = socket.gethostname()
            except Exception:
                logger.warning("Unable to determine hostname!")
                hostname = "unknown"
    if hostname.startswith("frontier"):
        return "Frontier"
    if hostname.startswith("sophia"):
        return "Sophia"
    if hostname.startswith("theta"):
        return "ThetaGPU"
    if hostname.startswith("x1"):
        return "SunSpot"
    if hostname.startswith("x3"):
        if "sirius" in hostname:
            return "Sirius"
        return "Polaris"
    if hostname.startswith("x4"):
        return "Aurora"
    if hostname.startswith("login"):
        return "Perlmutter"
    if hostname.startswith("nid"):
        return "Perlmutter"
    return f"{hostname}"

def get_node_index() -> int:
    """Get the index of the current node in the hostfile"""
    # return get_rank() % get_num_nodes()
    return get_rank() // get_gpus_per_node()

def get_nodes_from_hostfile(
    hostfile: PathLike,
) -> list[str]:
    """Get the nodes from the hostfile.

    Args:
        hostfile (PathLike): The path to the hostfile.

    Returns:
        list[str]: A list of nodes from the hostfile.
    """
    # cobalt_nodefile = get_cobalt_nodefile()
    fpath = Path(hostfile)
    assert fpath.is_file()
    with fpath.open("r") as f:
        nodes = [i.rstrip("\n") for i in f.readlines()]
    return nodes

def get_num_nodes(hostfile: Optional[PathLike] = None) -> int:
    """Get the number of nodes from the hostfile.

    Args:
        hostfile: Optional hostfile path to count nodes from.

    Examples:
        >>> get_num_nodes()  # counts lines in hostfile or SLURM env
        1
    """
    num_nodes = os.environ.get("SLURM_NNODES", None)
    if num_nodes is not None:
        return int(num_nodes)
    hfp = get_hostfile_with_fallback(hostfile)
    hosts = [h.split(".")[0] for h in get_nodes_from_hostfile(hfp)]
    return len(hosts)

def get_rank() -> int:
    """Get current MPI rank.

    Returns:
        int: The rank of the current process in the MPI world.

    Examples:
        >>> rank = get_rank()
        >>> print(f"Current MPI rank: {rank}")
    """
    return int(MPI.COMM_WORLD.Get_rank())

def get_torch_backend() -> str:
    """
    Get the current PyTorch backend.

    Returns:
        str: The current PyTorch backend.
    """
    # return (
    #     "nccl"
    #     if torch.cuda.is_available()
    #     else ("xccl" if (torch.xpu.is_available() and torch.distributeelse "gloo")
    # )
    backend_from_env = os.environ.get("TORCH_BACKEND", None)
    if backend_from_env is not None:
        return backend_from_env
    if torch.cuda.is_available() and torch.distributed.is_backend_available(
        "nccl"
    ):
        return "nccl"
    if torch.xpu.is_available():
        if torch.distributed.is_backend_available("xccl"):
            return "xccl"
        return "ccl"
    return "gloo"

def get_torch_backend_on_xpu() -> str:
    """
    ! DEPRECATED FUNCTION !

    Deal with breaking change introduced in torch 2.6:

    See: https://github.com/pytorch/pytorch/pull/141856

    Examples:

        >>> get_torch_backend_on_xpu()
        "xccl"
    """
    # torch_version = get_torch_version_as_float()
    # assert torch.xpu.is_available()
    # return "xccl" if torch_version > 2.5 else "ccl"
    # XXX: This is broken for whatever reason, so hardcoding for now
    #
    #     ```python
    #     >>> import ezpz
    #     >>> ezpz.get_torch_version_as_float()
    #     2.11
    #     >>> ezpz.get_torch_backend_on_xpu()
    #     'ccl'
    #     ```
    # ???
    return "xccl"

def get_torch_device(
    *,
    device_type: Optional[str] = None,
    as_torch_device: Optional[bool] = None,
) -> str | torch.device:
    """Get the current PyTorch device.

    Args:
        device_type (str, optional): The type of device to return.
            If None, it will be determined automatically.
            Defaults to None.
        as_torch_device (bool, optional): If True, return a torch.device object.
            If False, return a string representing the device type.
            Defaults to False.

    Returns:
        str | torch.device: The current PyTorch device.
            If as_torch_device is True, returns a torch.device object.
            If as_torch_device is False, returns a string representing the device type.

    Examples:
        >>> get_torch_device()  # Returns the current device type as a string
    """
    # env_info = _get_env_torch_device()
    # if env_info is not None:
    #     device_str = env_info[0]
    #     return torch.device(device_str) if as_torch_device else device_str
    env_info = _get_env_torch_device()
    if env_info is not None:
        _apply_env_torch_device(env_info)
        return (
            torch.device(env_info[0]) if as_torch_device else env_info[0]
        )
    if device_type is None:
        device_type = get_torch_device_type(device_type)
        return torch.device(device_type) if as_torch_device else device_type
    return torch.device(device_type) if as_torch_device else device_type

def get_torch_device_type(device_type: Optional[str] = None) -> str:
    """Get the current PyTorch device type.

    Args:
        device_type (str, optional): The type of device to return.
            If None, it will be determined automatically.
            Defaults to None.

    Returns:
        str: The current PyTorch device type.
            Possible values are "cpu", "mps", "xpu", or "cuda".

    Examples:
        >>> get_torch_device_type()  # returns 'cuda' if available
        >>> os.environ["TORCH_DEVICE"] = "cpu"
        >>> get_torch_device_type()
        'cpu'
    """
    if device_type is not None:
        assert device_type in _SUPPORTED_DEVICE_TYPES
        if get_rank() == 0:
            logger.warning(
                " ".join(
                    [
                        f"device_type: {device_type} passed to",
                        "ezpz.dist.get_torch_device_type",
                    ]
                )
            )
        return device_type
    env_info = _get_env_torch_device()
    if env_info is not None:
        _apply_env_torch_device(env_info)
        return env_info[1]
    return (
        "xpu"
        if torch.xpu.is_available()
        else (
            "cuda"
            if torch.cuda.is_available()
            else ("mps" if torch.backends.mps.is_available() else "cpu")
        )
    )

def get_torch_version_as_float() -> float:
    """Get the PyTorch version as a float.

    Returns:
        float: The PyTorch version as a float, e.g., 1.10.0 -> 1.10
    """
    return float(".".join(torch.__version__.split(".")[:2]))

def get_world_size(
    total: Optional[bool] = None,
    in_use: Optional[bool] = None,
    implementation: str = "mpi",
) -> int:
    """
    Get the total number of *PUs available or currently in use.
    Args:
        total (bool, optional): If True, return the total number of *PUs available.
            Defaults to None.
        in_use (bool, optional): If True, return the number of *PUs currently in use.
            Defaults to None.

    Returns:
        int: The total number of *PUs available or currently in use.
            If both `total` and `in_use` are None, it returns the number of *PUs
            currently in use by the MPI communicator.

    Examples:
        >>> world_size = get_world_size(total=True)
        >>> print(f"Total number of *PUs available: {world_size}")
        >>> world_size_in_use = get_world_size(in_use=True)
        >>> print(f"Number of *PUs currently in use: {world_size_in_use}")
    """
    if total:
        return get_world_size_total()
    if in_use:
        return get_world_size_in_use()
    # TODO: Deal with subtlety between:
    # 1. 'world_size' == total AVAILABLE gpus (for record keeping)
    # 2. 'world_size' == number of gpus CURRENTLY IN USE (from {`mpi`, ...})
    # ¯\_(ツ)_/¯
    impl = implementation.lower()
    mpi_size: Optional[int] = None
    torch_size: Optional[int] = None
    if torch.distributed.is_available() and torch.distributed.is_initialized():
        try:
            torch_size = int(torch.distributed.get_world_size())  # type: ignore
        except Exception:
            torch_size = None
    try:
        mpi_size = int(MPI.COMM_WORLD.Get_size())
    except Exception:
        mpi_size = None

    if mpi_size is not None and torch_size is not None and mpi_size != torch_size:
        if get_rank() == 0:
            logger.warning(
                "MPI world size (%s) differs from torch.distributed (%s)",
                mpi_size,
                torch_size,
            )

    if impl in {"torch", "pytorch", "pt"}:
        if torch_size is not None:
            return torch_size
        if mpi_size is not None:
            return mpi_size
    if mpi_size is not None:
        return mpi_size
    if torch_size is not None:
        return torch_size
    num_nodes = get_num_nodes()
    gpus_per_node = get_gpus_per_node()
    world_size = num_nodes * gpus_per_node
    if get_rank() == 0:
        logger.warning(
            "MPI not initialized !!"
            "Calculating (and using!! ??) "
            "[world_size]=[(num_nodes) x (num_*pus_per_node)]=[num_*pus_total]"
            f"[{world_size}]=[({num_nodes}) x ({gpus_per_node})]"
        )
    # if world_size == 1:
    #     gpus_per_node = get_gpus_per_node()
    #     num_nodes = get_num_nodes()
    #     world_size = num_nodes * gpus_per_node
    return world_size

def get_world_size_in_use() -> int:
    """Get number of currently in use MPI ranks

    Returns:
        int: The number of currently in use MPI ranks.
            This is the size of the MPI communicator.
            It is the number of processes that are currently running
            and participating in the distributed computation.

    Examples:
        >>> world_size_in_use = get_world_size_in_use()
        >>> print(f"Number of currently in use MPI ranks: {world_size_in_use}")
    """
    return int(MPI.COMM_WORLD.Get_size())