62 lines
2.2 KiB
Python
62 lines
2.2 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import asyncio
|
|
import json
|
|
from typing import Annotated
|
|
|
|
import typer
|
|
from openai import AsyncOpenAI
|
|
|
|
from slime.utils.data import read_file
|
|
|
|
|
|
# can unify w/ sglang_rollout.py later, e.g. add RM, if needed
|
|
def main(
|
|
prompt_data: Annotated[str, typer.Option()],
|
|
url: Annotated[str, typer.Option()] = "http://localhost:30000/v1",
|
|
input_key: Annotated[str, typer.Option()] = "input",
|
|
n_samples_per_prompt: Annotated[int, typer.Option()] = 1,
|
|
rollout_max_response_len: Annotated[int, typer.Option()] = 1024,
|
|
rollout_temperature: Annotated[float, typer.Option()] = 1.0,
|
|
rollout_top_p: Annotated[float, typer.Option()] = 1.0,
|
|
):
|
|
"""
|
|
Minimally send prompts to SGLang using OpenAI endpoints with arguments in the same format as main Slime.
|
|
|
|
Example usage:
|
|
python -m slime.utils.debug_utils.send_to_sglang --prompt-data /root/datasets/aime-2024/aime-2024.jsonl --input-key prompt --n-samples-per-prompt 16 --rollout-max-response-len 32768 --rollout-temperature 0.8 --rollout-top-p 0.7
|
|
"""
|
|
|
|
async def _main_async():
|
|
tasks = [
|
|
asyncio.create_task(_run_one(row, row_index=row_index, repeat_index=repeat_index))
|
|
for row_index, row in enumerate(read_file(prompt_data))
|
|
for repeat_index in range(n_samples_per_prompt)
|
|
]
|
|
outputs = await asyncio.gather(*tasks)
|
|
for output in outputs:
|
|
print(json.dumps(output))
|
|
|
|
async def _run_one(row, row_index: int, repeat_index: int):
|
|
resp = await client.chat.completions.create(
|
|
messages=row[input_key],
|
|
model="dummy_model",
|
|
max_tokens=rollout_max_response_len,
|
|
temperature=rollout_temperature,
|
|
top_p=rollout_top_p,
|
|
)
|
|
return dict(
|
|
row_index=row_index,
|
|
repeat_index=repeat_index,
|
|
**row,
|
|
response=resp.choices[0].message.content,
|
|
)
|
|
|
|
client = AsyncOpenAI(api_key="dummy_key", base_url=url)
|
|
asyncio.run(_main_async())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
typer.run(main)
|