from vllm.executor.multiproc_mlu_executor import MultiprocessingMLUExecutor
from vllm_mlu.mlu_hijack_utils import MluHijackObject

def vllm__executor__multiproc_mlu_executor__MultiprocessingMLUExecutor__recapture_model(
        self,
        context_batch_size_to_capture,
        context_seq_len_to_capture
    ) -> None:
    return self._run_workers("recapture_model",
                             context_batch_size_to_capture=context_batch_size_to_capture,
                             context_seq_len_to_capture=context_seq_len_to_capture)


MluHijackObject.apply_hijack(MultiprocessingMLUExecutor,
                             "recapture_model",
                             vllm__executor__multiproc_mlu_executor__MultiprocessingMLUExecutor__recapture_model)