tcod/tests/explorer/scheduler_test.py

import asyncio
import time
import unittest
from collections import defaultdict
from typing import Dict, List, Optional, Sequence

import ray
import torch
from parameterized import parameterized

from tests.tools import get_template_config
from trinity.common.config import ExperienceBufferConfig
from trinity.common.constants import StorageType, SyncStyle
from trinity.common.experience import EID, Experience
from trinity.common.models.model import InferenceModel, ModelWrapper
from trinity.common.workflows import WORKFLOWS, Task, Workflow
from trinity.explorer.scheduler import Scheduler


@WORKFLOWS.register_module("dummy_workflow")
class DummyWorkflow(Workflow):
    can_repeat: bool = True

    def __init__(self, *, task, model, auxiliary_models):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.step_num = task.workflow_args.get("step_num", 1)
        self.error_type = task.raw_task.get("error_type", "")
        self.seconds = None
        if "timeout" in self.error_type:
            parts = self.error_type.split("_")
            if len(parts) > 1:
                self.seconds = int(parts[-1])
            else:
                self.seconds = 10

    def set_repeat_times(self, repeat_times, run_id_base):
        self.repeat_times = repeat_times
        self.run_id_base = run_id_base

    def run(self) -> List[Experience]:
        if "timeout" in self.error_type:
            time.sleep(self.seconds)
        elif self.error_type == "exception":
            raise ValueError("Exception occurred")
        elif self.error_type == "exit":
            exit(1)
        elif self.error_type == "auxiliary_models":
            assert self.auxiliary_models is not None and len(self.auxiliary_models) == 2

        exps = []
        for i in range(self.repeat_times):
            run_level_metrics = {"run_metrics": float(i + self.run_id_base)}
            run_level_exps = []
            for step in range(self.step_num):
                run_level_exps.append(
                    Experience(
                        tokens=torch.zeros(5),
                        prompt_length=2,
                        prompt_text=self.error_type or "success",
                        eid=EID(run=i + self.run_id_base, step=step),
                        info={"repeat_times": self.repeat_times},
                    )
                )
            run_level_exps[-1].metrics = run_level_metrics
            exps.extend(run_level_exps)
        return exps


@WORKFLOWS.register_module("dummy_nonrepeat_workflow")
class DummyNonRepeatWorkflow(Workflow):
    can_reset: bool = True

    def __init__(self, *, task, model, auxiliary_models):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.reset_flag = False
        self.step_num = task.workflow_args.get("step_num", 1)
        self.metrics = task.workflow_args.get("metrics", [0])

    def reset(self, task: Task):
        self.task = task
        self.reset_flag = True
        self.step_num = task.workflow_args.get("step_num", 1)
        self.metrics = task.workflow_args.get("metrics", [0])

    def run(self) -> List[Experience]:
        exps = [
            Experience(
                eid=EID(run=self.run_id_base, step=step),
                tokens=torch.zeros(5),
                prompt_length=2,
                prompt_text="success",
                info={"reset_flag": self.reset_flag},
                metrics={
                    "run_metrics": self.metrics[step % len(self.metrics)],
                },
            )
            for step in range(self.step_num)
        ]
        return exps


@WORKFLOWS.register_module("dummy_async_workflow")
class DummyAsyncWorkflow(Workflow):
    can_repeat: bool = True
    is_async: bool = True

    def __init__(self, *, task, model, auxiliary_models):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.step_num = task.workflow_args.get("step_num", 1)

    def set_repeat_times(self, repeat_times, run_id_base):
        self.repeat_times = repeat_times
        self.run_id_base = run_id_base

    async def run_async(self):
        exps = []
        for i in range(self.repeat_times):
            run_level_metrics = {"run_metrics": float(i + self.run_id_base)}
            run_level_exps = []
            for step in range(self.step_num):
                run_level_exps.append(
                    Experience(
                        eid=EID(run=i + self.run_id_base, step=step),
                        tokens=torch.zeros(5),
                        prompt_length=2,
                        prompt_text="success",
                    )
                )
            run_level_exps[-1].metrics = run_level_metrics
            exps.extend(run_level_exps)
        return exps

    def run(self):
        raise RuntimeError("This method should not be called")


@WORKFLOWS.register_module("dummy_workflow_with_state")
class DummyWorkflowWithState(Workflow):
    can_repeat: bool = True
    is_async: bool = True

    def __init__(self, *, task, model: ModelWrapper, auxiliary_models):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)
        self.step_num = task.workflow_args.get("step_num", 1)

    def set_repeat_times(self, repeat_times, run_id_base):
        self.repeat_times = repeat_times
        self.run_id_base = run_id_base

    async def run_async(self) -> List[Experience]:
        exps = []
        for i in range(self.repeat_times):
            run_level_metrics = {"run_metrics": float(i + self.run_id_base)}
            run_level_exps = []
            for step in range(self.step_num):
                run_level_exps.append(
                    Experience(
                        eid=EID(run=i + self.run_id_base, step=step),
                        tokens=torch.zeros(5),
                        prompt_length=2,
                        prompt_text="success",
                    )
                )
            run_level_exps[-1].metrics = run_level_metrics
            self.logger.info(f"Setting workflow state to repeat_cnt={i}")
            await self.model.set_workflow_state({"repeat_cnt": i})
            await asyncio.sleep(1)
            exps.extend(run_level_exps)
        return exps


@WORKFLOWS.register_module("dummy_concurrent_workflow")
class DummyConcurrentWorkflow(Workflow):
    can_repeat: bool = False
    is_async: bool = True

    def __init__(self, *, task, model, auxiliary_models):
        super().__init__(task=task, model=model, auxiliary_models=auxiliary_models)

    async def run_async(self) -> List[Experience]:
        await asyncio.sleep(1)

        return [
            Experience(
                eid=EID(run=self.run_id_base, step=0),
                tokens=torch.zeros(5),
                prompt_length=2,
                prompt_text="success",
            )
        ]


@ray.remote
class DummyModel(InferenceModel):
    def __init__(self):
        from trinity.common.config import InferenceModelConfig

        super().__init__(InferenceModelConfig(model_path="dummy_model"))

    def sync_model(self, model_version, update_weight_args_list):
        return True

    async def prepare(self):
        return

    def get_model_version(self):
        return 0

    def init_process_group(
        self,
        master_address: str,
        master_port: int,
        rank_offset: int,
        world_size: int,
        group_name: str,
        backend: str = "nccl",
        timeout: int = 1200,
    ) -> None:
        pass

    def get_api_server_url(self) -> Optional[str]:
        return None

    async def chat(self, messages: List[Dict], lora_request=None, **kwargs) -> Sequence[Experience]:
        prompt_length = sum(len(msg["content"]) for msg in messages)
        return [
            Experience(
                tokens=torch.zeros(prompt_length + 10),
                prompt_length=prompt_length,
                logprobs=torch.zeros(10),
            )
        ]

    async def generate(self, prompt: str, lora_request=None, **kwargs) -> Sequence[Experience]:
        prompt_length = len(prompt)
        return [
            Experience(
                tokens=torch.zeros(prompt_length + 5),
                prompt_length=prompt_length,
                logprobs=torch.zeros(5),
            )
        ]


@ray.remote
class DummyAuxiliaryModel(InferenceModel):
    def sync_model(self, model_version, update_weight_args_list):
        return True

    def get_model_version(self):
        return 0

    def init_process_group(
        self,
        master_address: str,
        master_port: int,
        rank_offset: int,
        world_size: int,
        group_name: str,
        backend: str = "nccl",
        timeout: int = 1200,
    ) -> None:
        pass

    def get_api_server_url(self) -> str:
        return "http://localhost:12345"


def generate_tasks(
    total_num: int,
    timeout_num: int = 0,
    exception_num: int = 0,
    timeout_seconds: int = 10,
    repeat_times: int = 1,
    step_num: int = 1,
    repeatable: bool = True,
):
    """Generate some tasks for testing

    Args:
        total_num: number of normal tasks
        timeout_num: number of timeout tasks
        exception_num: number of exception tasks
        timeout_seconds: the timeout for timeout tasks
        repeat_times: number of times to repeat each task
        step_num: number of steps in each task
        repeatable: whether to use repeatableworkflow
    """
    workflow = DummyWorkflow if repeatable else DummyNonRepeatWorkflow
    tasks = [
        Task(
            workflow=workflow,  # type: ignore[type-abstract]
            workflow_args={"step_num": step_num},
            repeat_times=repeat_times,
            raw_task={},
        )
        for _ in range(total_num)
    ]

    tasks.extend(
        [
            Task(
                workflow=workflow,  # type: ignore[type-abstract]
                workflow_args={"step_num": step_num},
                repeat_times=repeat_times,
                raw_task={"error_type": f"timeout_{timeout_seconds}"},
            )
            for _ in range(timeout_num)
        ]
    )

    tasks.extend(
        [
            Task(
                workflow=workflow,  # type: ignore[type-abstract]
                workflow_args={"step_num": step_num},
                repeat_times=repeat_times,
                raw_task={"error_type": "exception"},
            )
            for _ in range(exception_num)
        ]
    )

    return tasks


class SchedulerTest(unittest.IsolatedAsyncioTestCase):
    def setUp(self):
        ray.init(ignore_reinit_error=True)
        self.config = get_template_config()
        self.config.explorer.max_retry_times = 1
        self.config.explorer.max_timeout = 5
        self.config.explorer.runner_per_model = 2
        self.config.buffer.train_batch_size = 2
        self.config.buffer.pad_token_id = 0
        self.config.buffer.explorer_output = (
            self.config.buffer.trainer_input.experience_buffer
        ) = ExperienceBufferConfig(
            name="test",
            storage_type=StorageType.QUEUE.value,
            schema_type="experience",
            path="",
        )
        self.config.buffer.trainer_input.experience_buffer.max_read_timeout = 1
        self.config.algorithm.repeat_times = 1
        self.config.check_and_update()

    async def test_get_results(self):
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()

        tasks = generate_tasks(8)
        scheduler.schedule(tasks, batch_id=0)

        statuses, exps = await scheduler.get_results(batch_id=0, min_num=8, timeout=20)
        self.assertEqual(len(statuses), 8)
        self.assertEqual(len(exps), 8)
        _, exps = await scheduler.get_results(batch_id=0, min_num=1, timeout=1)
        self.assertEqual(len(exps), 0)

        for result in statuses:
            self.assertTrue(result.ok)

        for batch_id in range(1, 4):
            tasks = generate_tasks(4)
            scheduler.schedule(tasks, batch_id=batch_id)

        for batch_id in range(1, 4):
            self.assertTrue(scheduler.has_step(batch_id))
            statuses, exps = await scheduler.get_results(batch_id=batch_id, min_num=4, timeout=10)
            self.assertEqual(len(statuses), 4)
            self.assertEqual(len(exps), 4)
            self.assertFalse(scheduler.has_step(batch_id))
        _, exps = await scheduler.get_results(batch_id=0, min_num=1, timeout=1)
        self.assertEqual(len(exps), 0)

        tasks = generate_tasks(3)
        scheduler.schedule(tasks, batch_id=4)
        self.assertTrue(scheduler.has_step(4))
        statuses, exps = await scheduler.get_results(batch_id=4)
        self.assertEqual(len(statuses), 3)
        self.assertEqual(len(exps), 3)
        self.assertFalse(scheduler.has_step(4))

        # test timeout
        tasks = generate_tasks(2, timeout_num=2, timeout_seconds=10)
        scheduler.schedule(tasks, batch_id=0)

        start_time = time.time()
        statuses, exps = await scheduler.get_results(batch_id=0, min_num=4, timeout=3)
        end_time = time.time()

        self.assertLessEqual(end_time - start_time, 5)
        self.assertEqual(len(statuses), 2)
        self.assertEqual(len(exps), 2)

        # test run tasks after timeout
        tasks = generate_tasks(4)
        scheduler.schedule(tasks, batch_id=0)

        # actor restart is slow, set a big timeout
        statuses, exps = await scheduler.get_results(batch_id=0, timeout=20)
        self.assertEqual(len(statuses), 4)

        success_count = sum(1 for r in statuses if r.ok)
        self.assertEqual(success_count, 4)
        self.assertEqual(len(exps), 4)
        _, exps = await scheduler.get_results(batch_id=0, min_num=1, timeout=1)
        self.assertEqual(len(exps), 0)

        # test exception tasks
        tasks = generate_tasks(1, exception_num=3)
        scheduler.schedule(tasks, batch_id=1)
        statuses, exps = await scheduler.get_results(batch_id=1, timeout=5)
        self.assertEqual(len(statuses), 4)

        success_count = sum(1 for r in statuses if r.ok)
        self.assertEqual(success_count, 1)
        self.assertEqual(len(exps), 1)
        _, exps = await scheduler.get_results(batch_id=1, min_num=1, timeout=1)
        self.assertEqual(len(exps), 0)

        # test clear_timeout_tasks
        tasks = generate_tasks(3, timeout_num=1, timeout_seconds=3)
        scheduler.schedule(tasks, batch_id=2)
        statuses, exps = await scheduler.get_results(
            batch_id=2, timeout=2, clear_timeout_tasks=False
        )
        self.assertEqual(len(statuses), 3)
        self.assertEqual(len(exps), 3)
        statuses, exps = await scheduler.get_results(
            batch_id=2, timeout=2, clear_timeout_tasks=False
        )
        self.assertEqual(len(statuses), 1)
        self.assertEqual(len(exps), 1)
        _, exps = await scheduler.get_results(batch_id=2, min_num=1, timeout=1)
        self.assertEqual(len(exps), 0)

        await scheduler.stop()

    async def test_wait_all(self):
        """Test wait all"""
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()

        tasks1 = generate_tasks(4)
        tasks2 = generate_tasks(3)
        scheduler.schedule(tasks1, batch_id=0)
        scheduler.schedule(tasks2, batch_id=1)

        start_time = time.time()
        await scheduler.wait_all(timeout=10.0)
        end_time = time.time()

        self.assertLess(end_time - start_time, 5.0)

        self.assertEqual(len(scheduler.pending_tasks), 0)
        self.assertEqual(len(scheduler.running_tasks), 0)

        status0, exps0 = await scheduler.get_results(batch_id=0, min_num=4, timeout=1)
        status1, exps1 = await scheduler.get_results(batch_id=1, min_num=3, timeout=1)
        self.assertEqual(len(status0), 4)
        self.assertEqual(len(status1), 3)

        # test timeout
        tasks = generate_tasks(2, timeout_num=2, timeout_seconds=10)
        scheduler.schedule(tasks, batch_id=0)

        start_time = time.time()
        with self.assertRaises(TimeoutError):
            await scheduler.wait_all(timeout=3.0)
        end_time = time.time()

        self.assertGreaterEqual(end_time - start_time, 2.8)
        self.assertLessEqual(end_time - start_time, 4.0)

        # test empty scenario

        start_time = time.time()
        await scheduler.wait_all(timeout=5.0)
        end_time = time.time()

        self.assertLess(end_time - start_time, 1.0)
        await scheduler.stop()

    async def test_wait_all_timeout_with_multi_batch(self):
        self.config.explorer.max_timeout = 5
        self.config.explorer.rollout_model.engine_num = 4
        self.config.explorer.runner_per_model = 1

        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()

        tasks = generate_tasks(1, timeout_num=3, timeout_seconds=3)
        scheduler.schedule(tasks, batch_id=0)
        tasks = generate_tasks(2, timeout_num=2, timeout_seconds=3)
        scheduler.schedule(tasks, batch_id=1)
        tasks = generate_tasks(3, timeout_num=1, timeout_seconds=3)
        scheduler.schedule(tasks, batch_id=2)
        start_time = time.time()
        await scheduler.wait_all()
        end_time = time.time()
        self.assertTrue(
            end_time - start_time > 9,
            f"wait time should be greater than 9, but got {end_time - start_time}",
        )

        await scheduler.stop()

    async def test_concurrent_operations(self):
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()

        async def schedule_tasks(batch_id, num_tasks):
            tasks = generate_tasks(num_tasks)
            scheduler.schedule(tasks, batch_id=batch_id)
            return await scheduler.get_results(batch_id=batch_id, min_num=num_tasks, timeout=10)

        results = await asyncio.gather(
            schedule_tasks(0, 3),
            schedule_tasks(1, 4),
            schedule_tasks(2, 2),
        )

        self.assertEqual(len(results[0][0]), 3)
        self.assertEqual(len(results[1][0]), 4)
        self.assertEqual(len(results[2][0]), 2)

        await scheduler.stop()

    async def test_scheduler_restart_after_stop(self):
        scheduler = Scheduler(self.config, [DummyModel.remote()])

        await scheduler.start()
        tasks = generate_tasks(2)
        scheduler.schedule(tasks, batch_id=0)
        results, exps = await scheduler.get_results(batch_id=0, min_num=2, timeout=10)
        self.assertEqual(len(results), 2)
        self.assertEqual(len(exps), 2)
        await scheduler.stop()

        await scheduler.start()
        tasks = generate_tasks(3, repeat_times=2)
        scheduler.schedule(tasks, batch_id=1)
        results, exps = await scheduler.get_results(batch_id=1, min_num=3, timeout=10)
        self.assertEqual(len(results), 3)
        self.assertEqual(len(exps), 3 * 2)
        await scheduler.stop()

    async def test_scheduler_all_methods(self):
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        tasks = generate_tasks(8)
        scheduler.schedule(tasks, batch_id=0)
        self.assertTrue(scheduler.has_step(0))
        statuses, exps = await scheduler.get_results(batch_id=0, min_num=8, timeout=20)
        self.assertEqual(len(statuses), 8)
        self.assertEqual(len(exps), 8)
        scheduler.schedule(tasks, batch_id=1)
        scheduler.schedule(tasks[:4], batch_id=2)
        self.assertFalse(scheduler.has_step(0))
        statuses, exps = await scheduler.get_results(batch_id=0, min_num=8)
        self.assertFalse(scheduler.has_step(0))
        self.assertEqual(len(statuses), 0)  # batch_id 0 has no more tasks
        self.assertEqual(len(exps), 0)
        self.assertFalse(scheduler.has_step(0))
        self.assertTrue(scheduler.has_step(1))
        self.assertTrue(scheduler.has_step(2))
        await scheduler.wait_all()
        st = time.time()
        statuses, exps = await scheduler.get_results(batch_id=1)
        et = time.time()
        self.assertTrue(et - st < 1.0)
        self.assertEqual(len(statuses), 8)
        self.assertEqual(len(exps), 8)
        self.assertFalse(scheduler.has_step(1))
        self.assertTrue(scheduler.has_step(2))
        st = time.time()
        statuses, exps = await scheduler.get_results(batch_id=2)
        et = time.time()
        self.assertTrue(et - st < 1.0)
        self.assertEqual(len(statuses), 4)
        self.assertEqual(len(exps), 4)
        self.assertFalse(scheduler.has_step(2))
        await scheduler.stop()

    async def test_split_tasks(self):
        self.config.explorer.max_repeat_times_per_runner = 2
        self.config.check_and_update()
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        exp_list = []

        tasks = generate_tasks(4, repeat_times=8)  # ceil(8 / 2) == 4
        scheduler.schedule(tasks, batch_id=1)
        statuses, exps = await scheduler.get_results(batch_id=1)
        self.assertEqual(len(statuses), 4)
        self.assertEqual(len(exps), 4 * 8)
        exp_list.extend(exps)
        _, exps = await scheduler.get_results(batch_id=1, min_num=1, timeout=1)
        self.assertEqual(len(exps), 0)

        tasks = generate_tasks(4, repeat_times=5)  # ceil(5 / 2) == 3
        scheduler.schedule(tasks, batch_id=2)
        statuses, exps = await scheduler.get_results(batch_id=2)
        self.assertEqual(len(statuses), 4)
        self.assertEqual(len(exps), 4 * 5)
        exp_list.extend(exps)
        _, exps = await scheduler.get_results(batch_id=2, min_num=1, timeout=1)
        self.assertEqual(len(exps), 0)

        tasks = generate_tasks(3, repeat_times=1)  # ceil(1 / 2) == 1
        scheduler.schedule(tasks, batch_id=3)
        statuses, exps = await scheduler.get_results(batch_id=3)
        self.assertEqual(len(statuses), 3)
        self.assertEqual(len(exps), 3 * 1)
        exp_list.extend(exps)
        _, exps = await scheduler.get_results(batch_id=3, min_num=1, timeout=1)
        self.assertEqual(len(exps), 0)

        # test task_id, run_id and unique_id
        group_ids = [exp.eid.tid for exp in exp_list]
        self.assertEqual(len(set(group_ids)), 11)  # 4 + 4 + 3
        run_ids = [exp.eid.rid for exp in exp_list]
        self.assertEqual(len(run_ids), len(set(run_ids)))
        unique_ids = [exp.eid.uid for exp in exp_list]
        self.assertEqual(len(unique_ids), len(set(unique_ids)))

        await scheduler.stop()

    async def test_multi_step_execution(self):
        self.config.explorer.max_repeat_times_per_runner = 1
        self.config.check_and_update()
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        tasks = generate_tasks(2, repeat_times=4)

        n_steps = 3
        for i in range(1, n_steps + 1):
            scheduler.schedule(tasks, batch_id=i)
            statuses, exps = await scheduler.get_results(batch_id=i)
            self.assertEqual(len(statuses), 2)
            self.assertEqual(len(exps), 2 * 4)

        await scheduler.stop()

    async def test_non_repeatable_workflow(self):
        self.config.explorer.max_repeat_times_per_runner = 2
        self.config.check_and_update()
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        task_num, repeat_times = 5, 4
        tasks = generate_tasks(task_num, repeat_times=repeat_times, repeatable=False)

        batch_num = 2
        exp_list = []
        for i in range(1, batch_num + 1):
            scheduler.schedule(tasks, batch_id=i)
            statuses, exps = await scheduler.get_results(batch_id=i)
            self.assertEqual(len(statuses), task_num)
            self.assertEqual(len(exps), task_num * repeat_times)
            exp_list.extend(exps)

        # test task_id, run_id and unique_id
        group_ids = [exp.eid.tid for exp in exp_list]
        self.assertEqual(len(set(group_ids)), batch_num * task_num)
        run_ids = [exp.eid.rid for exp in exp_list]
        self.assertEqual(len(set(run_ids)), batch_num * task_num * repeat_times)
        unique_ids = [exp.eid.uid for exp in exp_list]
        self.assertEqual(len(unique_ids), len(set(unique_ids)))

        # test reset used properly
        runner_num = (
            self.config.explorer.runner_per_model * self.config.explorer.max_repeat_times_per_runner
        )
        self.assertEqual(
            sum([exp.info["reset_flag"] for exp in exp_list]), len(exp_list) - runner_num
        )

    async def test_async_workflow(self):
        self.config.explorer.max_repeat_times_per_runner = 2
        self.config.check_and_update()
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        task_num, repeat_times, step_num = 5, 4, 3
        tasks = [
            Task(
                workflow=DummyAsyncWorkflow,  # type: ignore[type-abstract]
                workflow_args={"step_num": step_num},
                repeat_times=repeat_times,
                raw_task={},
            )
            for _ in range(task_num)
        ]

        batch_num = 2
        exp_list = []
        for i in range(1, batch_num + 1):
            scheduler.schedule(tasks, batch_id=i)
            statuses, exps = await scheduler.get_results(batch_id=i)
            self.assertEqual(len(statuses), task_num)
            self.assertEqual(len(exps), task_num * repeat_times * step_num)
            exp_list.extend(exps)

        # test task_id, run_id and unique_id
        group_ids = [exp.eid.tid for exp in exp_list]
        self.assertEqual(len(set(group_ids)), batch_num * task_num)
        run_ids = [exp.eid.rid for exp in exp_list]
        self.assertEqual(len(set(run_ids)), batch_num * task_num * repeat_times)
        unique_ids = [exp.eid.uid for exp in exp_list]
        self.assertEqual(len(unique_ids), len(set(unique_ids)))

    async def test_stepwise_experience_eid(self):
        task_num, repeat_times, step_num = 2, 4, 3
        self.config.buffer.batch_size = task_num
        self.config.buffer.train_batch_size = task_num * repeat_times * step_num
        self.config.explorer.max_repeat_times_per_runner = 2
        self.config.check_and_update()
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        batch_num = 2

        # repeatable stepwise workflow
        tasks = generate_tasks(
            task_num, step_num=step_num, repeat_times=repeat_times, repeatable=True
        )
        exp_list = []
        for i in range(1, batch_num + 1):
            scheduler.schedule(tasks, batch_id=i)
            statuses, exps = await scheduler.get_results(batch_id=i)
            self.assertEqual(len(statuses), task_num)
            self.assertEqual(len(exps), task_num * repeat_times * step_num)
            exp_list.extend(exps)

        # test task_id, run_id and unique_id
        group_ids = [exp.eid.tid for exp in exp_list]
        self.assertEqual(len(set(group_ids)), batch_num * task_num)
        run_ids = [exp.eid.rid for exp in exp_list]
        self.assertEqual(len(set(run_ids)), batch_num * task_num * repeat_times)
        unique_ids = [exp.eid.uid for exp in exp_list]
        self.assertEqual(len(unique_ids), len(set(unique_ids)))

        # Non-repeatable stepwise workflow
        tasks = generate_tasks(
            task_num, step_num=step_num, repeat_times=repeat_times, repeatable=False
        )
        exp_list = []
        for i in range(1, batch_num + 1):
            scheduler.schedule(tasks, batch_id=i)
            statuses, exps = await scheduler.get_results(batch_id=i)
            self.assertEqual(len(statuses), task_num)
            self.assertEqual(len(exps), task_num * repeat_times * step_num)
            exp_list.extend(exps)

        # test task_id, run_id and unique_id
        group_ids = [exp.eid.tid for exp in exp_list]
        self.assertEqual(len(set(group_ids)), batch_num * task_num)
        run_ids = [exp.eid.rid for exp in exp_list]
        self.assertEqual(len(set(run_ids)), batch_num * task_num * repeat_times)
        unique_ids = [exp.eid.uid for exp in exp_list]
        self.assertEqual(len(unique_ids), len(set(unique_ids)))

    @parameterized.expand(
        [
            (2,),
            (None,),
        ]
    )
    async def test_metric_calculation_with_repeatable_workflow(self, max_repeat_times_per_runner):
        self.config.explorer.max_repeat_times_per_runner = max_repeat_times_per_runner
        self.config.check_and_update()
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        tasks = []
        tasks.extend(generate_tasks(total_num=1, step_num=1, repeat_times=4, repeatable=True))
        tasks.extend(generate_tasks(total_num=1, step_num=4, repeat_times=8, repeatable=True))
        scheduler.schedule(tasks, batch_id=0)
        statuses, exps = await scheduler.get_results(batch_id=0)
        self.assertEqual(len(statuses), 2)
        self.assertEqual(len(exps), 1 * 4 * 1 + 1 * 8 * 4)
        self.assertAlmostEqual(statuses[0].metrics[0]["run_metrics"], 1.5)  # (0+1+2+3)/4
        self.assertAlmostEqual(statuses[1].metrics[0]["run_metrics"], 3.5)  # (0+1+2+3+4+5+6+7)/8

    @parameterized.expand(
        [
            (2,),
            (None,),
        ]
    )
    async def test_metric_calculation_with_non_repeatable_workflow(
        self, max_repeat_times_per_runner
    ):
        self.config.explorer.max_repeat_times_per_runner = max_repeat_times_per_runner
        self.config.check_and_update()
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        tasks = []
        tasks.extend(generate_tasks(total_num=1, step_num=3, repeat_times=4, repeatable=False))
        tasks[-1].workflow_args["metrics"] = [1.0, 2.0, 3.0]
        tasks.extend(generate_tasks(total_num=1, step_num=8, repeat_times=5, repeatable=False))
        tasks[-1].workflow_args["metrics"] = [2 * i for i in range(8)]
        scheduler.schedule(tasks, batch_id=0)
        statuses, exps = await scheduler.get_results(batch_id=0)
        self.assertEqual(len(statuses), 2)
        self.assertEqual(len(exps), 1 * 4 * 3 + 1 * 5 * 8)
        self.assertAlmostEqual(statuses[0].metrics[0]["run_metrics"], 2.0)  # (1+2+3)/3
        self.assertAlmostEqual(statuses[1].metrics[0]["run_metrics"], 7.0)  # (0+2+4+6+8+10+12+14)/8

    async def test_over_rollout_min_wait(self):
        self.config.explorer.over_rollout.ratio = 0.5
        self.config.explorer.over_rollout.wait_after_min = 3
        self.config.explorer.max_repeat_times_per_runner = None
        self.config.buffer.batch_size = 4
        self.config.synchronizer.sync_style = SyncStyle.DYNAMIC_BY_EXPLORER
        self.config.check_and_update()
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        tasks = []
        tasks.extend(generate_tasks(0, timeout_num=2, repeat_times=1, timeout_seconds=1))
        tasks.extend(generate_tasks(0, timeout_num=1, repeat_times=1, timeout_seconds=3))
        tasks.extend(generate_tasks(0, timeout_num=1, repeat_times=1, timeout_seconds=6))
        scheduler.schedule(tasks, batch_id=0)
        statuses, exps = await scheduler.get_results(batch_id=0, min_num=2)
        self.assertEqual(len(statuses), 3)
        self.assertEqual(len(exps), 3 * 1)

    async def test_dynamic_timeout(self):
        self.config.explorer.dynamic_timeout.enable = True
        self.config.explorer.dynamic_timeout.ratio = 3.0
        self.config.buffer.batch_size = 4
        self.config.explorer.max_timeout = 20
        self.config.explorer.max_retry_times = 0  # no retry here
        scheduler = Scheduler(self.config, [DummyModel.remote(), DummyModel.remote()])
        await scheduler.start()
        tasks = []
        tasks.extend(generate_tasks(0, timeout_num=4, repeat_times=1, timeout_seconds=1))
        for task in tasks:
            task.is_eval = True
        scheduler.schedule(
            tasks, batch_id="0/eval"
        )  # eval tasks will not count into dynamic timeout
        statuses, exps = await scheduler.get_results(batch_id="0/eval")
        self.assertEqual(len(statuses), 4)
        self.assertEqual(len(exps), 0)
        self.assertEqual(scheduler.total_running_time, 0)
        self.assertEqual(scheduler.total_completed_tasks, 0)
        tasks = []
        # generate 4 tasks that will run 1 second
        tasks.extend(generate_tasks(0, timeout_num=4, repeat_times=1, timeout_seconds=1))
        scheduler.schedule(tasks, batch_id=0)  # first step will not use dynamic timeout
        statuses, exps = await scheduler.get_results(batch_id=0)
        self.assertEqual(len(statuses), 4)
        # dynamic timeout will be set to 3.0 * 1.0 = 3.0 seconds for next step
        tasks = []
        tasks.extend(generate_tasks(0, timeout_num=4, repeat_times=1, timeout_seconds=4))
        st = time.time()
        scheduler.schedule(tasks, batch_id=1)
        statuses, exps = await scheduler.get_results(batch_id=1)
        et = time.time()
        self.assertTrue(
            et - st < 4
        )  # should wait about 1 * 3.0 seconds, here we set 4 seconds timeout
        self.assertEqual(len(exps), 0)
        self.assertEqual(len(statuses), 4)
        # tasks take 2 seconds, which is within the dynamic timeout 3.0 * 1.0 = 3.0 seconds
        tasks = []
        tasks.extend(generate_tasks(0, timeout_num=4, repeat_times=1, timeout_seconds=2))
        scheduler.schedule(tasks, batch_id=2)
        statuses, exps = await scheduler.get_results(batch_id=2)
        self.assertEqual(len(statuses), 4)
        self.assertEqual(len(exps), 4)

    def tearDown(self):
        try:
            ray.shutdown()
        except Exception:
            pass


class TestRunnerStateCollection(unittest.IsolatedAsyncioTestCase):
    async def test_runner_state_collection(self):
        ray.init(ignore_reinit_error=True)
        config = get_template_config()
        config.explorer.runner_per_model = 2
        config.explorer.runner_state_report_interval = 0.5
        config.explorer.max_repeat_times_per_runner = 2
        config.check_and_update()
        scheduler = Scheduler(config, [DummyModel.remote(), DummyModel.remote()])
        # 4 runner in side the scheduler
        await scheduler.start()

        tasks = [
            Task(
                workflow=DummyWorkflowWithState,  # type: ignore[type-abstract]
                workflow_args={"step_num": 2},
                repeat_times=4,
                raw_task={},
            )
            for _ in range(4)
        ]
        scheduler.schedule(tasks, batch_id=0)

        async def monitor_routine():
            runner_0_state_history = defaultdict(set)
            await asyncio.sleep(0.5)  # wait for first report
            for _ in range(16):
                await asyncio.sleep(0.3)
                states = scheduler.get_all_state()
                self.assertEqual(len(states), 4)
                for state in states.values():
                    self.assertIn("workflow_id", state)
                    self.assertIn("model_version", state)
                    self.assertIn("begin_time", state)
                    self.assertIn("terminate_time", state)
                    self.assertIn("repeat_cnt", state)
                ids = scheduler.get_key_state("workflow_id")
                self.assertEqual(len(ids), 4)
                self.assertEqual(len(set(ids.values())), 4)
                runner_0_state = scheduler.get_runner_state(0)
                for key, value in runner_0_state.items():
                    runner_0_state_history[key].add(value)
            self.assertEqual(len(runner_0_state_history["repeat_cnt"]), 2)  # max_repeat_times is 2
            self.assertEqual(len(runner_0_state_history["model_version"]), 1)
            self.assertEqual(
                len(runner_0_state_history["workflow_id"]), 2
            )  # split into 2 sub tasks
            self.assertEqual(len(runner_0_state_history["begin_time"]), 2)

        await asyncio.gather(
            monitor_routine(),
            scheduler.get_results(batch_id=0),
        )