[V1] Fix torch profiling for offline inference (#11125)

Signed-off-by: Roger Wang <ywang@roblox.com>
This commit is contained in:
Roger Wang 2024-12-12 07:51:53 -08:00 committed by GitHub
parent 85362f028c
commit 4816d20aa4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 21 additions and 14 deletions

View File

@ -1,4 +1,5 @@
import os
import time
from vllm import LLM, SamplingParams
@ -15,13 +16,15 @@ prompts = [
# Create a sampling params object.
sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
if __name__ == "__main__":
# Create an LLM.
llm = LLM(model="facebook/opt-125m", tensor_parallel_size=1)
llm.start_profile()
# Generate texts from the prompts. The output is a list of RequestOutput objects
# that contain the prompt, generated text, and other information.
# Generate texts from the prompts. The output is a list of RequestOutput
# objects that contain the prompt, generated text, and other information.
outputs = llm.generate(prompts, sampling_params)
llm.stop_profile()
@ -31,3 +34,7 @@ for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")
# Add a buffer to wait for profiler in the background process
# (in case MP is on) to finish writing profiling output.
time.sleep(10)

View File

@ -105,7 +105,7 @@ class InprocClient(EngineCoreClient):
def __del__(self):
self.shutdown()
async def profile(self, is_start=True) -> None:
def profile(self, is_start=True) -> None:
self.engine_core.profile(is_start)
@ -212,7 +212,7 @@ class SyncMPClient(MPClient):
def abort_requests(self, request_ids: List[str]) -> None:
self._send_input(EngineCoreRequestType.ABORT, request_ids)
async def profile(self, is_start=True) -> None:
def profile(self, is_start=True) -> None:
self._send_input(EngineCoreRequestType.PROFILE,
EngineCoreProfile(is_start))