File size: 2,205 Bytes
6011e08 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 | # 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)
|