38 lines
1.2 KiB
Python
38 lines
1.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-MLU project
|
|
|
|
from vllm.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
class MluProfilerWrapper:
|
|
def __init__(self) -> None:
|
|
self._profiler_running = False
|
|
# Note: lazy import to avoid dependency issues if MLU is not available.
|
|
import torch.mlu.profiler as mlu_profiler
|
|
|
|
self._mlu_profiler = mlu_profiler
|
|
|
|
def start(self) -> None:
|
|
try:
|
|
self._mlu_profiler.start()
|
|
self._profiler_running = True
|
|
logger.info_once("Started MLU profiler")
|
|
except Exception as e:
|
|
logger.warning_once("Failed to start MLu profiler: %s", e)
|
|
|
|
def stop(self) -> None:
|
|
if self._profiler_running:
|
|
try:
|
|
self._mlu_profiler.stop()
|
|
logger.info_once("Stopped MLU profiler")
|
|
except Exception as e:
|
|
logger.warning_once("Failed to stop MLU profiler: %s", e)
|
|
finally:
|
|
self._profiler_running = False
|
|
|
|
def shutdown(self) -> None:
|
|
"""Ensure profiler is stopped when shutting down."""
|
|
self.stop()
|