slime_plugins/rollout_buffer/generator/base_generator.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

import copy
import json
import random
import time
import uuid
from functools import partial
from multiprocessing import Process, Queue
from time import sleep

import requests
from openai import OpenAI
from tqdm import tqdm
from slime.rollout.rm_hub import get_deepscaler_rule_based_reward

TASK_TYPE = "math"

SAMPLING_PARAMS = {
    "top_p": 1,
}


def get_rule_based_math_reward(item):
    messages = item["messages"]
    label = item["label"]
    assert messages[-1]["role"] == "assistant", "last message must be assistant, but got {}".format(
        messages[-1]["role"]
    )

    response = messages[-1]["content"]
    if response is None or len(response) == 0:
        return 0

    reward = get_deepscaler_rule_based_reward(response, label)
    return reward


def query_single_turn(client, messages, sampling_params, tools=None):
    base_payload = {
        "messages": messages,
        **sampling_params,
        "model": "custom",
        "stream": False,
        "seed": random.randint(1, 10000000),
        "tools": tools,
    }

    text = None
    accumulated_tokens = 0
    finish_reason = "stop"

    for _attempt in range(6):
        try:
            # Create a fresh payload for each attempt
            current_payload = copy.deepcopy(base_payload)

            if text is not None:
                # Update messages with current progress
                current_messages = copy.deepcopy(messages)
                current_messages.append({"role": "assistant", "content": text})
                current_payload["messages"] = current_messages

                # Adjust max_tokens based on accumulated tokens
                if "max_tokens" in sampling_params:
                    current_payload["max_tokens"] = max(0, sampling_params["max_tokens"] - accumulated_tokens)

                # Add continue flag for partial rollouts
                current_payload["extra_body"] = {"continue_final_message": True}
            if current_payload["max_tokens"] == 0:
                break
            response = client.chat.completions.create(**current_payload)

            if len(response.choices) > 0:
                finish_reason = response.choices[0].finish_reason
                if finish_reason == "abort":
                    print(
                        f"query failed, reason: {response.choices[0].finish_reason}, currently generated: {response.usage.completion_tokens}"
                    )

                    accumulated_tokens += response.usage.completion_tokens

                    if text is None:
                        text = response.choices[0].message.content
                    else:
                        text += response.choices[0].message.content

                    sleep(10)
                    continue
                if text is None:
                    text = response.choices[0].message.content
                elif response.choices[0].message.content is not None:
                    text += response.choices[0].message.content
                break
            else:
                print(f"Error in query, status code: {response.status_code}")
                continue
        except Exception as e:
            print(f"query failed in single turn, error: {e}")
            continue

    # Update final messages
    if len(messages) > 0 and messages[-1]["role"] == "assistant":
        messages = messages[:-1]
    messages.append({"role": "assistant", "content": text})

    return messages, finish_reason


def worker_process(task_queue, done_queue, rollout_func, reward_func, client, sampling_params):

    for line in iter(task_queue.get, "STOP"):
        if isinstance(line, str):
            item = json.loads(line)
        else:
            item = line

        # try:
        messages, finish_reason = rollout_func(client, item["prompt"], sampling_params)

        item["uid"] = str(uuid.uuid4())
        item["messages"] = messages
        reward = reward_func(item)
        item["rollout_index"] = 1
        item["reward"] = reward
        item["extra_info"] = {}
        item.update(sampling_params)
        item["timestamp"] = str(time.time())
        item["round_number"] = len([_ for _ in item["messages"] if _["role"] == "assistant"])
        item["finish_reason"] = finish_reason

        output_item = {
            "uid": item.pop("uid"),
            "messages": messages,
            "reward": reward,
            "instance_id": item.pop("instance_id"),
            "extra_info": item,
        }

        done_queue.put(output_item)

    done_queue.put("COMPLETE")


class BaseGenerator:
    def __init__(
        self,
        remote_engine_url,
        remote_buffer_url,
        num_repeat_per_sample=1,
        queue_size=1000000,
        num_process=10,
        task_type="math",
        max_tokens=4096,
        num_repeats=10,
        skip_instance_ids: list[str] | None = None,
    ):
        self.queue_size = queue_size
        self.num_process = num_process
        self.remote_engine_url = remote_engine_url
        self.remote_buffer_url = remote_buffer_url
        self.num_repeat_per_sample = num_repeat_per_sample
        self.task_type = task_type
        self.max_tokens = max_tokens
        self.num_repeats = num_repeats
        # Ensure skip_instance_ids is a mutable list (copy to avoid modifying original)
        self.skip_instance_ids = list(skip_instance_ids) if skip_instance_ids is not None else None

        if self.skip_instance_ids is not None:
            print(f"BaseGenerator initialized with {len(self.skip_instance_ids)} instance_ids to skip")
            self.skip_instance_ids = self.skip_instance_ids * self.num_repeat_per_sample

        if "/v1" in remote_engine_url:
            self.client = OpenAI(api_key="test", base_url=remote_engine_url)
        else:
            remote_engine_url = remote_engine_url.strip("/") + "/v1"
            self.client = OpenAI(api_key="test", base_url=remote_engine_url)

    def send_data_to_buffer(self, data):
        remote_buffer_url = self.remote_buffer_url.rstrip("/") + "/buffer/write"

        for _ in range(2):
            try:
                response = requests.post(remote_buffer_url, json=data)
                if response.status_code == 200:
                    break
                else:
                    print(f"send data to buffer failed, status code: {response.status_code}")
                    continue
            except Exception as e:
                print(f"send data to buffer failed, error: {e}")
                continue

    def run(self, input_file, rollout_func, reward_func):
        task_queue, done_queue = Queue(maxsize=self.queue_size), Queue(maxsize=self.queue_size)

        def read_data_into_queue():
            cnt = 0
            items = []
            skipped_count = 0
            with open(input_file) as f:
                for i, line in enumerate(f):
                    item = json.loads(line)
                    if "instance_id" not in item:
                        item["instance_id"] = i
                    items.append(item)
            random.shuffle(items)

            for _ in range(self.num_repeats):

                for item in items:
                    for _ in range(self.num_repeat_per_sample):
                        item_repeat = copy.deepcopy(item)

                        if "uid" not in item_repeat:
                            item_repeat["uid"] = str(uuid.uuid4())

                        # Check if instance_id should be skipped
                        if self.skip_instance_ids is not None and item_repeat["instance_id"] in self.skip_instance_ids:
                            print(f"Skipping instance_id: {item_repeat['instance_id']}")
                            # Remove from skip list to handle potential duplicates in multiple epochs
                            self.skip_instance_ids.remove(item_repeat["instance_id"])
                            skipped_count += 1
                            continue

                        task_queue.put(item_repeat)
                    cnt += 1
                time.sleep(300)

            if skipped_count > 0:
                remaining_skip_count = len(self.skip_instance_ids) if self.skip_instance_ids is not None else 0
                print(
                    f"Rollout summary: skipped {skipped_count} instance_ids, {remaining_skip_count} still in skip list"
                )

            for _ in range(self.num_process):
                task_queue.put("STOP")

        processes = []
        SAMPLING_PARAMS["max_tokens"] = self.max_tokens

        for _ in range(self.num_process):
            process = Process(
                target=partial(worker_process, client=self.client, sampling_params=SAMPLING_PARAMS),
                args=(task_queue, done_queue, rollout_func, reward_func),
            )
            process.start()
            processes.append(process)

        process = Process(target=read_data_into_queue)
        process.start()

        progress_bar = tqdm()
        num_finished = 0
        while num_finished < self.num_process:
            item = done_queue.get()
            if item == "COMPLETE":
                num_finished += 1
            else:
                assert "reward" in item, f"reward not in item: {item}"
                assert "instance_id" in item, f"instance_id not in item: {item}"
                self.send_data_to_buffer(item)
                progress_bar.update(1)

        progress_bar.close()

        return "finished"

    def entry(self, input_file, rollout_func, reward_func, num_epoch=1):
        for _ in range(num_epoch):
            self.run(input_file, rollout_func, reward_func)


def run_rollout(data: dict):

    print(f"Starting math rollout with data: {data}")

    rollout_func = query_single_turn
    reward_func = get_rule_based_math_reward

    print("Waiting for 10 seconds for buffer server to start")
    time.sleep(10)
    global SAMPLING_PARAMS
    for k, v in data["sampling_params"].items():
        SAMPLING_PARAMS[k] = v
        print(f"Set {k} to {v}", type(v))

    generator = BaseGenerator(
        data["remote_engine_url"],
        data["remote_buffer_url"],
        num_repeat_per_sample=int(data["num_repeat_per_sample"]),
        queue_size=1000000,
        max_tokens=int(data["sampling_params"]["max_tokens"]),
        num_process=int(data.get("num_process", 100)),
        task_type=data["task_type"],
        skip_instance_ids=data.get("skip_instance_ids", None),
    )

    generator.entry(data["input_file"], rollout_func, reward_func, int(data.get("num_epoch", 1)))


def normalize_group_data(group, epsilon=1e-8, algo="grpo"):
    print(f"Using math-specific normalization for group {group[0]}")

    assert algo == "grpo", "Only 'grpo' is supported for now."

    instance_id = group[0]
    data = group[1]
    rewards = [item["reward"] for item in data]

    valid_rewards = [r for r in rewards if 1 >= r >= 0]

    if set(valid_rewards) == {0}:
        normalized_rewards = rewards
    else:
        mean_reward = sum(valid_rewards) / len(valid_rewards)
        std_reward = (sum((r - mean_reward) ** 2 for r in valid_rewards) / len(valid_rewards)) ** 0.5

        if std_reward < epsilon:
            print(f"[Math Info] Zero variance in group {instance_id}, setting all to 0.")
            normalized_rewards = [0.0 if 1 >= r >= 0 else r for r in rewards]
        else:
            normalized_rewards = [(r - mean_reward) / (std_reward + epsilon) if 1 >= r >= 0 else r for r in rewards]

    for i, item in enumerate(data):
        item["reward"] = normalized_rewards[i]
        item["raw_reward"] = rewards[i]

    return (instance_id, data)


def is_valid_group(group, min_valid_group_size, task_type="math"):
    # Handle both tuple and list inputs
    if isinstance(group, tuple):
        instance_id, items = group
    else:
        items = group

    # Count valid items (non-empty responses)
    valid_indices = []
    for i, item in enumerate(items):
        if item["messages"][-1]["content"].strip():
            valid_indices.append(i)

    group_size = len(items)
    valid_count = len(valid_indices)

    # A group is finished if it has reached the target size
    is_finished = group_size >= min_valid_group_size

    is_valid = is_finished and valid_count >= min_valid_group_size

    return is_valid