from vllm.engine.multiprocessing.client import MQLLMEngineClient
from vllm.logger import init_logger

from vllm_mlu.engine.multiprocessing import RPCSchedulerProfileRequest
from vllm_mlu.mlu_hijack_utils import MluHijackObject

logger = init_logger(__name__)

class MQLLMEngineClient_V2(MQLLMEngineClient):

    async def init_scheduler_view(self):
        """Send INIT_SCHEDULER_VIEW request to RPC Server."""

        await self._send_one_way_rpc_request(
            request=RPCSchedulerProfileRequest.INIT_SCHEDULER_VIEW,
            socket=self.input_socket)


    async def save_scheduler_view(self):
        """Send SAVE_SCHEDULER_VIEW request to RPC Server."""

        await self._send_one_way_rpc_request(
            request=RPCSchedulerProfileRequest.SAVE_SCHEDULER_VIEW,
            socket=self.input_socket)


MluHijackObject.apply_hijack(MQLLMEngineClient,
                             "init_scheduler_view",
                             MQLLMEngineClient_V2.init_scheduler_view)
MluHijackObject.apply_hijack(MQLLMEngineClient,
                             "save_scheduler_view",
                             MQLLMEngineClient_V2.save_scheduler_view)