import torch
from typing import List

from vllm.worker.mlu_multi_step_model_runner import MLUMultiStepModelRunner
from vllm_mlu.mlu_hijack_utils import MluHijackObject
from vllm.logger import init_logger

logger = init_logger(__name__)


def vllm__worker__mlu_multi_step_model_runner__MLUMultiStepModelRunner__capture_model_with_context(
    self,
    kv_caches: List[List[torch.Tensor]],
    num_gpu_blocks: int
) -> None:
    return self._base_model_runner.capture_model_with_context(kv_caches, num_gpu_blocks)


def vllm__worker__mlu_multi_step_model_runner__MLUMultiStepModelRunner__reset_capture_context(self):
    return self._base_model_runner.reset_capture_context()


MluHijackObject.apply_hijack(MLUMultiStepModelRunner,
                             "capture_model_with_context",
                             vllm__worker__mlu_multi_step_model_runner__MLUMultiStepModelRunner__capture_model_with_context)
MluHijackObject.apply_hijack(MLUMultiStepModelRunner,
                             "reset_capture_context",
                             vllm__worker__mlu_multi_step_model_runner__MLUMultiStepModelRunner__reset_capture_context)
