distributed/parallel_state.py: process groups and GroupCoordinator

            from vllm.distributed.utils import (
                get_cpu_distributed_timeout_or_none,
                get_distributed_timeout_or_none,
            )

            timeout = get_cpu_distributed_timeout_or_none()
            device_timeout = get_distributed_timeout_or_none()

            for ranks in group_ranks:
                device_group = torch.distributed.new_group(
                    ranks,
                    backend=torch_distributed_backend,
                    timeout=device_timeout,
                )
                # a group with `gloo` backend, to allow direct coordination between
                # processes through the CPU.
                with suppress_stdout():
                    cpu_group = torch.distributed.new_group(
                        ranks, backend="gloo", timeout=timeout
                    )
                if self.rank in ranks:
                    self.ranks = ranks
                    self.world_size = len(ranks)
                    self.rank_in_group = ranks.index(self.rank)
                    self_device_group = device_group
                    self_cpu_group = cpu_group

        self.use_device_communicator = use_device_communicator
        self.device_communicator = None
        if use_device_communicator and self.world_size > 1:
            device_comm_cls = resolve_obj_by_qualname(
                current_platform.get_device_communicator_cls()
            )
            self.device_communicator = device_comm_cls(
                cpu_group=self.cpu_group,
                device=self.device,
                device_group=self.device_group,
                unique_name=self.unique_name,
                use_all2all=use_all2all,
            )

        from vllm.distributed.device_communicators.shm_broadcast import MessageQueue

        self.mq_broadcaster: MessageQueue | None = None
        if use_message_queue_broadcaster and self.world_size > 1:
            self.mq_broadcaster = MessageQueue.create_from_process_group(
                self.cpu_group, 1 << 22, 6
            )

        # TODO(#35915): Remove is_tpu() check once tpu_inference
        # overrides use_custom_op_collectives() to return True.
        self.use_custom_op_call = (
            current_platform.is_tpu() or current_platform.use_custom_op_collectives()
        )

    def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
        """User-facing all-reduce function before we actually call the
        all-reduce operation.

        We need this because Dynamo does not support passing an arbitrary
        object (`self` in this case) to a custom op. We need to pass the
         group name as a string, and then look up the group coordinator from
         the group name, dispatch the all-reduce operation to the group
         coordinator.

        In addition, PyTorch custom ops do not support mutation or returning
        a new tensor in the same op. So we always make the all-reduce operation
        out-of-place.
        """
        # Bypass the function if we are using only 1 GPU.
        if self.world_size == 1:
            return input_

        if self.use_custom_op_call:
            return torch.ops.vllm.all_reduce(input_, group_name=self.unique_name)
        else:
            return self._all_reduce_out_place(input_)

    def _all_reduce_out_place(self, input_: torch.Tensor) -> torch.Tensor:
        if self.device_communicator is None:
            raise ValueError("No device communicator found")
        return self.device_communicator.all_reduce(input_)

        group = self.device_group
        metadata_group = self.cpu_group

        metadata_list, tensor_list = _split_tensor_dict(tensor_dict)

        tensor_keys = [k for k, v in tensor_dict.items() if isinstance(v, torch.Tensor)]
        assert len(tensor_keys) == len(tensor_list)

        # Reap completed sends before posting new ones, bounding the
        # self-retention FIFO.
        self._reap_completed_isends()
        handles: list[Handle] = [self.isend_object(metadata_list, dst=dst)]

        sent_tensors: list[torch.Tensor] = []
        for key, tensor in zip(tensor_keys, tensor_list):
            if tensor.numel() == 0:
                continue

            if self._should_use_all_gather(
                key, tensor.numel(), all_gather_group, all_gather_tensors
            ):
                tensor = tensor.reshape(all_gather_size, -1)[all_gather_rank]

            comm_group = metadata_group if tensor.is_cpu else group
            handle = torch.distributed.isend(
                tensor, dst=self.ranks[dst], group=comm_group
            )
            if tensor.is_cuda:
                tensor.record_stream(torch.cuda.current_stream(tensor.device))
            handles.append(handle)
            sent_tensors.append(tensor)

        # Self-retain for fire-and-forget callers: keep the handles and
        # the source tensors alive until the sends complete (reaped
        # lazily by ``_reap_completed_isends``). Without this, dropping
        # the returned handles would GC the ``_RetainedHandle``'s
        # serialized metadata buffers mid-send, and non-CUDA devices
        # have no ``record_stream`` equivalent protecting the tensor
        # sources.
        pending = getattr(self, "_pending_isends", None)
        if pending is None:
            self._pending_isends = pending = deque()
        pending.append((handles, sent_tensors))

        return handles

    enable_elastic_ep = config is not None and config.parallel_config.enable_elastic_ep
    world_pg_store: Store | None = None
    if (
        config is not None
        and config.parallel_config.distributed_executor_backend != "external_launcher"
        and (
            config.parallel_config.nnodes > 1
            or config.parallel_config.data_parallel_size > 1
        )
        and not enable_elastic_ep
    ):
        parallel_config = config.parallel_config
        # adjust to take into account data parallelism
        # offset the rank by the data parallel rank
        rank = parallel_config.data_parallel_rank * world_size + rank
        # adjust the world size to take into account data parallelism
        world_size = parallel_config.world_size_across_dp

        # Use appropriate IP and port based on configuration

    # the layout order is: ExternalDP x DP x PP x PCP x TP
    # ExternalDP is the data parallel group that is not part of the model,
    # every dp rank can generate independently (in verl integration).
    # DP is the data parallel group that is part of the model,
    # all the ranks in the same DP group should generate simultaneously,
    # i.e. the `generate` call in the same DP group should be called together,
    # otherwise it will cause deadlock.
    # to get group_ranks for each dimension, transpose that dimension to the
    # last dimension, then reshape to 2D, then unbind the last dimension
    all_ranks = torch.arange(world_size).reshape(
        -1,
        data_parallel_size,
        pipeline_model_parallel_size,
        prefill_context_model_parallel_size,
        tensor_model_parallel_size,
    )  # noqa

    # Build the tensor model-parallel groups.
    global _TP
    assert _TP is None, "tensor model parallel group is already initialized"
    group_ranks = all_ranks.view(-1, tensor_model_parallel_size).unbind(0)

def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
    logger.debug(
        "[shutdown] Distributed: cleanup start shutdown_ray=%s",
        shutdown_ray,
    )
    # Reset environment variable cache
    envs.disable_envs_cache()

    # Reset rocm_aiter_ops class variables to match current os.environ.
    # These are class-level attributes that persist across tests and are
    # NOT restored by monkeypatch (which only restores os.environ).
    from vllm.platforms import current_platform

    if current_platform.is_rocm():
        from vllm._aiter_ops import rocm_aiter_ops

        rocm_aiter_ops.refresh_env_variables()

    # Ensure all objects are not frozen before cleanup
    gc.unfreeze()

    destroy_model_parallel()
    destroy_distributed_environment()
    if shutdown_ray:
        import ray  # Lazy import Ray

        ray.shutdown()
    gc.collect()