from vllm.executor.mlu_executor import MLUExecutor
from vllm_mlu.mlu_hijack_utils import MluHijackObject


def vllm__executor__mlu_executor__MLUExecutor__get_latency(self) -> float:
    '''
    requires that torch.mlu.synchronize() be executed before this function
    for getting an accurate reading
    '''
    latency = self.driver_worker.get_latency()
    return latency


def vllm__executor__mlu_executor__MLUExecutor__recapture_model(
        self,
        context_batch_size_to_capture,
        context_seq_len_to_capture
    ) -> None:
    return self.driver_worker.recapture_model(context_batch_size_to_capture,
                                              context_seq_len_to_capture)


def vllm__executor__mlu_executor__MLUExecutor__get_memory_usage(self):
    return self.driver_worker.get_memory_usage()


MluHijackObject.apply_hijack(MLUExecutor,
                             "get_latency",
                             vllm__executor__mlu_executor__MLUExecutor__get_latency)
MluHijackObject.apply_hijack(MLUExecutor,
                             "recapture_model",
                             vllm__executor__mlu_executor__MLUExecutor__recapture_model)
MluHijackObject.apply_hijack(MLUExecutor,
                             "get_memory_usage",
                             vllm__executor__mlu_executor__MLUExecutor__get_memory_usage)