from typing import Any, Dict, Optional

from vllm.executor.gpu_executor import GPUExecutor
from vllm_mlu.mlu_hijack_utils import MluHijackObject

def vllm__executor__gpu_executor__GPUExecutor___get_worker_kwargs(
    self,
    local_rank: int = 0,
    rank: int = 0,
    distributed_init_method: Optional[str] = None,
) -> Dict[str, Any]:
    """Return worker init args for a given rank."""
    if distributed_init_method is None:
        distributed_init_method = get_distributed_init_method(
            get_ip(), get_open_port())
    '''
    ==========================
    Modify by Context Parallel
    ==========================
    @brief: replace self.parallel_config.tensor_parallel_size with self.parallel_config.world_size.
    '''
    return dict(
        vllm_config=self.vllm_config,
        local_rank=local_rank,
        rank=rank,
        distributed_init_method=distributed_init_method,
        is_driver_worker=(not self.parallel_config)
        or (rank % self.parallel_config.world_size == 0),
    )
    '''
    =======================
    End of Context Parallel
    =======================
    '''


MluHijackObject.apply_hijack(
    GPUExecutor,
    GPUExecutor._get_worker_kwargs,
    vllm__executor__gpu_executor__GPUExecutor___get_worker_kwargs)
