import json
import os

import torch

from roll.configs.worker_config import StrategyArguments
from roll.distributed.scheduler.initialize import init
from tests.distributed.strategy.generate.generate_pipeline import GenerateCmpPipeline
from tests.distributed.strategy.make_baseline_config import make_baseline_config
from roll.utils.logging import get_logger

logger = get_logger()


def vllm_generate_compare():
    os.environ["RAY_PROFILING"] = "1"

    init()

    ppo_config = make_baseline_config(config_path="./generate", config_name="generate_baseline_config")

    ppo_config.rollout_batch_size = 128
    ppo_config.actor_train.data_args.max_samples = 128
    ppo_config.actor_infer.model_args.model_name_or_path = "Qwen/Qwen2.5-7B-Instruct"
    ppo_config.actor_train.model_args.model_name_or_path = "Qwen/Qwen2.5-7B-Instruct"

    pipeline = GenerateCmpPipeline(pipeline_config=ppo_config)

    metric_list = pipeline.run()

    output_file = "generate_compare.json"
    with open(output_file, "w") as f:
        json.dump(metric_list, f, ensure_ascii=False)

    generate_times = [metric["time/generate"] for metric in metric_list]
    total_time = sum(generate_times)

    logger.info(f"{json.dumps({'total_time': total_time, 'time_list': generate_times})}")

    import ray

    ray.timeline(filename="timeline.json")
    ray._private.state.object_transfer_timeline(filename="object_timeline.json")


if __name__ == "__main__":
    vllm_generate_compare()
