from functools import partial
import os
import queue
import threading
import traceback
from typing import Callable, Dict, List, Mapping, Optional, Type, Union
from zlib import ZLIB_VERSION
import torch
from vllm import envs
from vllm.config import DecodingConfig, ObservabilityConfig, VllmConfig
from vllm.core.scheduler import ScheduledSequenceGroup
from vllm.engine.llm_engine import _LOCAL_LOGGING_INTERVAL_SEC, LLMEngine, SchedulerContext, SchedulerOutputState
from vllm.engine.metrics_types import StatLoggerBase
from vllm.engine.output_processor.interfaces import SequenceGroupOutputProcessor
from vllm.logger import init_logger
from vllm.executor.executor_base import ExecutorBase
from vllm.inputs import INPUT_REGISTRY
from vllm.inputs.data import ProcessorInputs
from vllm.inputs.parse import split_enc_dec_inputs
from vllm.inputs.preprocess import InputPreprocessor
from vllm.inputs.registry import InputRegistry
from vllm.lora.request import LoRARequest
from vllm.model_executor.layers.sampler import SamplerOutput
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.registry import MultiModalRegistry
from vllm.outputs import PoolingRequestOutput, RequestOutput
from vllm.pooling_params import PoolingParams
from vllm.sampling_params import SamplingParams
from vllm.sequence import ExecuteModelRequest, ParallelSampleSequenceGroup, SequenceGroup, SequenceGroupBase, SequenceGroupMetadata
from vllm.tracing import init_tracer
from vllm.transformers_utils.tokenizer import AnyTokenizer
from vllm.version import __version__ as VLLM_VERSION
from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled
from vllm.utils import resolve_obj_by_qualname, weak_bind, Counter
from vllm.zero_overhead.sampler import SampleRecorder, get_last_sampler
from vllm.zero_overhead.sequence import ZeroOverheadSequence
from vllm.zero_overhead.stop_check import ZeroOverheadStopChecker
from vllm.zero_overhead.tokenizer import ZeroOverheadDetokenizer
from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled,
                                  usage_message)
from vllm.profiler.prof import profile
from vllm.zero_overhead.utils import SpecStepKind, get_accepted_token_ids, get_spec_step, is_zero_no_thread, set_spec_step, zero_overhead_stream

logger = init_logger(__name__)

class ZeroOverheadEngine(LLMEngine):
    def __init__(
        self,
        vllm_config: VllmConfig,
        executor_class: Type[ExecutorBase],
        log_stats: bool,
        usage_context: UsageContext = UsageContext.ENGINE_CONTEXT,
        stat_loggers: Optional[Dict[str, StatLoggerBase]] = None,
        mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY,
        use_cached_outputs: bool = False,
    ) -> None:
        if envs.VLLM_USE_V1:
            raise ValueError(
                "Using V0 LLMEngine, but envs.VLLM_USE_V1=True. "
                "This should not happen. As a workaround, try using "
                "LLMEngine.from_vllm_config(...) or explicitly set "
                "VLLM_USE_V1=0 or 1 and report this issue on Github.")

        self.vllm_config = vllm_config
        self.model_config = vllm_config.model_config
        self.cache_config = vllm_config.cache_config
        self.lora_config = vllm_config.lora_config
        self.parallel_config = vllm_config.parallel_config
        self.scheduler_config = vllm_config.scheduler_config
        self.device_config = vllm_config.device_config
        self.speculative_config = vllm_config.speculative_config  # noqa
        self.load_config = vllm_config.load_config
        self.decoding_config = vllm_config.decoding_config or DecodingConfig(  # noqa
        )
        self.prompt_adapter_config = vllm_config.prompt_adapter_config  # noqa
        self.observability_config = vllm_config.observability_config or ObservabilityConfig(  # noqa
        )

        logger.info(
            "Initializing a V0 LLM engine (v%s) with config: %s, "
            "use_cached_outputs=%s, ",
            VLLM_VERSION,
            vllm_config,
            use_cached_outputs,
        )

        self.log_stats = log_stats
        self.use_cached_outputs = use_cached_outputs
        self.thread_running = False

        if not self.model_config.skip_tokenizer_init:
            self.tokenizer = self._init_tokenizer()
            self.detokenizer = ZeroOverheadDetokenizer(self.tokenizer)
            tokenizer_group = self.get_tokenizer_group()
        else:
            self.tokenizer = None
            self.detokenizer = None
            tokenizer_group = None

        # Ensure that the function doesn't contain a reference to self,
        # to avoid engine GC issues
        def get_tokenizer_for_seq(sequence: ZeroOverheadSequence) -> AnyTokenizer:
            assert tokenizer_group, ("tokenizer_group cannot be None, "
                                     "make sure skip_tokenizer_init is False")
            return tokenizer_group.get_lora_tokenizer(sequence.lora_request)

        self.seq_counter = Counter()
        self.generation_config_fields = (
            self.model_config.try_get_generation_config())

        self.input_preprocessor = InputPreprocessor(self.model_config,
                                                    self.tokenizer,
                                                    mm_registry)

        self.model_executor = executor_class(vllm_config=vllm_config, )

        if self.model_config.runner_type != "pooling":
            self._initialize_kv_caches()

        # If usage stat is enabled, collect relevant info.
        if is_usage_stats_enabled():
            from vllm.model_executor.model_loader import (
                get_architecture_class_name)
            usage_message.report_usage(
                get_architecture_class_name(self.model_config),
                usage_context,
                extra_kvs={
                    # Common configuration
                    "dtype":
                    str(self.model_config.dtype),
                    "tensor_parallel_size":
                    self.parallel_config.tensor_parallel_size,
                    "block_size":
                    self.cache_config.block_size,
                    "gpu_memory_utilization":
                    self.cache_config.gpu_memory_utilization,

                    # Quantization
                    "quantization":
                    self.model_config.quantization,
                    "kv_cache_dtype":
                    str(self.cache_config.cache_dtype),

                    # Feature flags
                    "enable_lora":
                    bool(self.lora_config),
                    "enable_prompt_adapter":
                    bool(self.prompt_adapter_config),
                    "enable_prefix_caching":
                    self.cache_config.enable_prefix_caching,
                    "enforce_eager":
                    self.model_config.enforce_eager,
                    "disable_custom_all_reduce":
                    self.parallel_config.disable_custom_all_reduce,
                })

        self.cached_scheduler_outputs = [
            SchedulerOutputState()
            for _ in range(self.parallel_config.pipeline_parallel_size)
        ]

        self.scheduler_contexts = [
            SchedulerContext(multi_step_stream_outputs=self.scheduler_config.
                             multi_step_stream_outputs)
            for _ in range(self.parallel_config.pipeline_parallel_size)
        ]

        if self.model_config.use_async_output_proc:
            process_model_outputs = weak_bind(self._process_model_outputs)

            self.async_callbacks = [
                partial(process_model_outputs,
                        ctx=self.scheduler_contexts[v_id])
                for v_id in range(self.parallel_config.pipeline_parallel_size)
            ]
        else:
            self.async_callbacks = []

        # Currently used by AsyncLLMEngine to ensure quick append
        # of request outputs to asyncio queues
        self.process_request_outputs_callback: Optional[Callable] = None

        # Create the scheduler.
        # NOTE: the cache_config here have been updated with the numbers of
        # GPU and CPU blocks, which are profiled in the distributed executor.
        if isinstance(self.vllm_config.scheduler_config.scheduler_cls, str):
            Scheduler = resolve_obj_by_qualname(
                self.vllm_config.scheduler_config.scheduler_cls)
        else:
            Scheduler = self.vllm_config.scheduler_config.scheduler_cls
        self.scheduler = [
            Scheduler(
                self.scheduler_config, self.cache_config, self.lora_config,
                self.parallel_config.pipeline_parallel_size,
                self.async_callbacks[v_id]
                if self.model_config.use_async_output_proc else None)
            for v_id in range(self.parallel_config.pipeline_parallel_size)
        ]

        # Metric Logging.
        if self.log_stats:
            if stat_loggers is not None:
                self.stat_loggers = stat_loggers
            else:
                # Lazy import for prometheus multiprocessing.
                # We need to set PROMETHEUS_MULTIPROC_DIR environment variable
                # before prometheus_client is imported.
                # See https://prometheus.github.io/client_python/multiprocess/
                from vllm.engine.metrics import (LoggingStatLogger,
                                                 PrometheusStatLogger)

                self.stat_loggers = {
                    "logging":
                    LoggingStatLogger(
                        local_interval=_LOCAL_LOGGING_INTERVAL_SEC,
                        vllm_config=vllm_config),
                    "prometheus":
                    PrometheusStatLogger(
                        local_interval=_LOCAL_LOGGING_INTERVAL_SEC,
                        labels=dict(
                            model_name=self.model_config.served_model_name),
                        vllm_config=vllm_config),
                }
                self.stat_loggers["prometheus"].info("cache_config",
                                                     self.cache_config)

        self.tracer = None
        if self.observability_config.otlp_traces_endpoint:
            self.tracer = init_tracer(
                "vllm.llm_engine",
                self.observability_config.otlp_traces_endpoint)

        # Create sequence output processor, e.g. for beam search or
        # speculative decoding.
        self.output_processor = (
            SequenceGroupOutputProcessor.create_output_processor(
                self.scheduler_config,
                self.detokenizer,
                self.scheduler,
                self.seq_counter,
                get_tokenizer_for_seq,
                stop_checker=ZeroOverheadStopChecker(
                    self.scheduler_config.max_model_len,
                    get_tokenizer_for_seq,
                ),
            ))
        self.tree_decoding = os.environ.get('VLLM_TREE_DECODING') == '1'

        self.seq_id_to_seq_group: Dict[str, SequenceGroupBase] = {}

        # Flag to set when an input fails to process and the engine should run
        # the next step without re-scheduling.
        self._skip_scheduling_next_step = False
        self.async_d2h = None
        self.last_record = None
        self.async_event = torch.cuda.Event(enable_timing=False)
        self.q_recorder = queue.Queue()
        self.use_stream = zero_overhead_stream(self.model_executor.device_config.device)
        if not is_zero_no_thread():
            self.zero_thread = threading.Thread(target=self.thread_zero_overhead)
            self.thread_running = True
            self.sem_m2s = threading.Semaphore(0) # main to scheduler thread
            self.zero_thread.start()
        profile.StartTracer()
    
    def __del__(self):
        self.finish_thread()
        return super().__del__()

    def finish_thread(self):
        if self.thread_running:
            self.thread_running = False
            self.sem_m2s.release()

    
    def thread_zero_overhead(self):
        logger.info('zero overhead thread start!')
        last_sampler = get_last_sampler()
        last_sampler.seq_ids.clear()
        try:
            with torch.cuda.stream(self.use_stream):
                while True:
                    self.sem_m2s.acquire()
                    if not self.thread_running:
                        logger.debug("Stopping remote worker execution loop.")
                        self.model_executor.stop_remote_worker_execution_loop()

                        break

                    virtual_engine = 0
                    # Clear outputs for each new scheduler iteration
                    
                    # Schedule iteration
                    (seq_group_metadata_list, scheduler_outputs,
                        allow_async_output_proc
                        ) = self.scheduler[virtual_engine].schedule()
                    if self.last_record is not None:
                        last_sampler = self.last_record[1]
                        spec_step = get_spec_step()
                        if spec_step == SpecStepKind.KIND_DEFAULT:
                            if last_sampler.sampled_token_ids_tensor is not None:
                                self.async_d2h = last_sampler.sampled_token_ids_tensor.to('cpu', non_blocking=True)
                            else:
                                self.async_d2h = None
                        elif spec_step == SpecStepKind.SCORE_DECODE:
                            self.async_d2h = last_sampler.to('cpu', non_blocking=True)
                        self.async_event.record()
                        self.q_recorder.put(self.last_record)
                    else:
                        self.q_recorder.put(None)
                    if len(seq_group_metadata_list) == 0:
                        self.last_record = None
                        continue

                    finished_requests_ids = self.scheduler[
                        virtual_engine].get_and_reset_finished_requests_ids()

                    assert seq_group_metadata_list is not None
                    assert scheduler_outputs is not None
                    last_sampled_token_ids = \
                        self._get_last_sampled_token_ids(virtual_engine)

                    execute_model_req = ExecuteModelRequest(
                        seq_group_metadata_list=seq_group_metadata_list,
                        blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
                        blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
                        blocks_to_copy=scheduler_outputs.blocks_to_copy,
                        num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
                        running_queue_size=scheduler_outputs.running_queue_size,
                        finished_requests_ids=finished_requests_ids,
                        # We use ExecuteModelRequest to pass the last sampled_token_ids
                        # to each of the non-last PP stages for in-place prepare_input.
                        last_sampled_token_ids=last_sampled_token_ids)
                    outputs = self.model_executor.execute_model(
                        execute_model_req=execute_model_req)

                    for output in outputs:
                        self._advance_to_next_step(
                            output, seq_group_metadata_list,
                            scheduler_outputs.scheduled_seq_groups)
                    scheduler_outputs.scheduled_seq_groups = [item for item in scheduler_outputs.scheduled_seq_groups] #deep copy
                    last_sampler = None
                    spec_step = get_spec_step()
                    if spec_step == SpecStepKind.KIND_DEFAULT:
                        last_sampler = get_last_sampler()
                    elif spec_step == SpecStepKind.SCORE_DECODE:
                        last_sampler, _ = get_accepted_token_ids()
                    self.last_record = [outputs, last_sampler, seq_group_metadata_list, scheduler_outputs, spec_step]

        except Exception as e:
            print(f"thread_zero_overhead error : {e}")
            traceback.print_exc()

    def zero_overhead_step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]:
        if not self.thread_running:
            self.zero_thread.join()
            self.thread_running = True
            self.zero_thread = threading.Thread(target=self.thread_zero_overhead)
            self.zero_thread.start()
        self.sem_m2s.release()
        recode_output = self.q_recorder.get()
        if recode_output is None: # None is for the first step
            return None
        virtual_engine = 0
        ctx = self.scheduler_contexts[virtual_engine]
        ctx.request_outputs.clear()
        outputs, last_sampler, seq_group_metadata_list, scheduler_outputs, spec_step = recode_output
        ctx.seq_group_metadata_list = seq_group_metadata_list
        ctx.scheduler_outputs = scheduler_outputs
        if spec_step == SpecStepKind.KIND_DEFAULT:
            self.async_event.synchronize()
            if self.async_d2h is not None:
                self._fix_last_step(
                    outputs, last_sampler, seq_group_metadata_list,
                    scheduler_outputs.scheduled_seq_groups)
        elif spec_step == SpecStepKind.SCORE_DECODE:
            self.async_event.synchronize()
            self._fix_spec_decode_steps(
                outputs, seq_group_metadata_list,
                scheduler_outputs.scheduled_seq_groups)

        # is_first_step_output is True only when the num_steps of all
        # the sequences are 1. When the num_steps > 1,
        # multi_step_model_runner does the first-step output append.
        is_first_step_output: bool = False if not seq_group_metadata_list \
            else seq_group_metadata_list[0].state.num_steps == 1

        # Add results to the output_queue
        ctx.append_output(outputs=outputs,
                            seq_group_metadata_list=seq_group_metadata_list,
                            scheduler_outputs=scheduler_outputs,
                            is_async=True,
                            is_last_step=True,
                            is_first_step_output=is_first_step_output)

        # Check if need to run the usual non-async path
        #if not allow_async_output_proc:
        self._process_model_outputs(ctx=ctx)

        #profile.ProfRangeAutoPush('has_unfinish')
        if not self.has_unfinished_requests():
            # Drain async postprocessor (if exists)
            if len(ctx.output_queue) > 0:
                self._process_model_outputs(ctx=ctx)
            assert len(ctx.output_queue) == 0

            # Stop the execute model loop in parallel workers until there are
            # more requests to process. This avoids waiting indefinitely in
            # torch.distributed ops which may otherwise timeout, and unblocks
            # the RPC thread in the workers so that they can process any other
            # queued control plane messages, such as add/remove lora adapters.
            # logger.debug("Stopping remote worker execution loop.")
            # self.model_executor.stop_remote_worker_execution_loop()
            self.finish_thread()
        return ctx.request_outputs
    
    
    def _fix_last_step(
            self, output: List[SamplerOutput],
            last_sampler: SampleRecorder,
            seq_group_metadata_list: List[SequenceGroupMetadata],
            scheduled_seq_groups: List[ScheduledSequenceGroup]) -> None:
        
        #sample_out_list = output[0].sampler_out_tenosr.cpu().tolist()
        sample_out_list = self.async_d2h.tolist()
        sample_out_ids = last_sampler.seq_ids
        for seq_group_metadata, sequence_group_outputs, scheduled_seq_group in \
            zip(seq_group_metadata_list, output[0], scheduled_seq_groups):
            seq_group = scheduled_seq_group.seq_group

            if seq_group.is_finished():
                continue

            if seq_group_metadata.do_sample:
                sample = sequence_group_outputs.samples[0]

                assert len(seq_group.seqs) == 1
                seq : ZeroOverheadSequence = seq_group.seqs[0]
                for token_id, seq_id in zip(sample_out_list, sample_out_ids):
                    if seq.seq_id == seq_id:
                        if type(token_id) is list:
                            sample.output_token = token_id[0]
                        else:
                            sample.output_token = token_id
                        seq.fix_last_token_id(sample.output_token)
                        break
    
    def _fix_spec_decode_steps(
            self, output: List[SamplerOutput],
            seq_group_metadata_list: List[SequenceGroupMetadata],
            scheduled_seq_groups: List[ScheduledSequenceGroup]):
        
        sample_out_list = self.async_d2h.tolist()
        group_idx = 0
        for seq_group_metadata, accept_token_ids, scheduled_seq_group in \
            zip(seq_group_metadata_list, sample_out_list, scheduled_seq_groups):
            seq_group = scheduled_seq_group.seq_group

            if seq_group.is_finished():
                group_idx += 1
                continue

            if seq_group_metadata.do_sample:
                assert len(seq_group.seqs) == 1
                seq : ZeroOverheadSequence = seq_group.seqs[0]
                remove_count = 0
                for token_id in accept_token_ids:
                    if token_id == -1:
                        remove_count += 1
                    else:
                        seq.fix_last_token_id(token_id)
                seq.remove_last_place_holder(remove_count)
            group_idx += 1

    def no_thread_step(self):
        virtual_engine = 0
        # Clear outputs for each new scheduler iteration
        
        # Schedule iteration
        (seq_group_metadata_list, scheduler_outputs,
            allow_async_output_proc
            ) = self.scheduler[virtual_engine].schedule()
        if self.last_record is not None:
            last_sampler = self.last_record[1]
            self.async_d2h = last_sampler.sampled_token_ids_tensor.to('cpu', non_blocking=True)
            self.async_event.record()
            self.q_recorder.put(self.last_record)
        else:
            self.q_recorder.put(None)
        if len(seq_group_metadata_list) == 0:
            self.last_record = None
        else:
            finished_requests_ids = self.scheduler[
                virtual_engine].get_and_reset_finished_requests_ids()

            assert seq_group_metadata_list is not None
            assert scheduler_outputs is not None
            last_sampled_token_ids = \
                self._get_last_sampled_token_ids(virtual_engine)

            execute_model_req = ExecuteModelRequest(
                seq_group_metadata_list=seq_group_metadata_list,
                blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in,
                blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out,
                blocks_to_copy=scheduler_outputs.blocks_to_copy,
                num_lookahead_slots=scheduler_outputs.num_lookahead_slots,
                running_queue_size=scheduler_outputs.running_queue_size,
                finished_requests_ids=finished_requests_ids,
                # We use ExecuteModelRequest to pass the last sampled_token_ids
                # to each of the non-last PP stages for in-place prepare_input.
                last_sampled_token_ids=last_sampled_token_ids)
            outputs = self.model_executor.execute_model(
                execute_model_req=execute_model_req)

            if len(outputs) == 1:
                self._advance_to_next_step(
                    outputs[0], seq_group_metadata_list,
                    scheduler_outputs.scheduled_seq_groups)
            scheduler_outputs.scheduled_seq_groups = [item for item in scheduler_outputs.scheduled_seq_groups] #deep copy
            last_sampler = get_last_sampler()    
            self.last_record = [outputs, last_sampler, seq_group_metadata_list, scheduler_outputs]

        recode_output = self.q_recorder.get()
        if recode_output is None: # None is for the first step
            return None
        virtual_engine = 0
        ctx = self.scheduler_contexts[virtual_engine]
        ctx.request_outputs.clear()
        outputs, last_sampler, seq_group_metadata_list, scheduler_outputs = recode_output
        ctx.seq_group_metadata_list = seq_group_metadata_list
        ctx.scheduler_outputs = scheduler_outputs
        self.async_event.synchronize()
        self._fix_last_step(
            outputs, last_sampler, seq_group_metadata_list,
            scheduler_outputs.scheduled_seq_groups)

        # is_first_step_output is True only when the num_steps of all
        # the sequences are 1. When the num_steps > 1,
        # multi_step_model_runner does the first-step output append.
        is_first_step_output: bool = False if not seq_group_metadata_list \
            else seq_group_metadata_list[0].state.num_steps == 1

        # Add results to the output_queue
        ctx.append_output(outputs=outputs,
                            seq_group_metadata_list=seq_group_metadata_list,
                            scheduler_outputs=scheduler_outputs,
                            is_async=True,
                            is_last_step=True,
                            is_first_step_output=is_first_step_output)

        # Check if need to run the usual non-async path
        #if not allow_async_output_proc:
        self._process_model_outputs(ctx=ctx)

        #profile.ProfRangeAutoPush('has_unfinish')
        if not self.has_unfinished_requests():
            # Drain async postprocessor (if exists)
            if len(ctx.output_queue) > 0:
                self._process_model_outputs(ctx=ctx)
            assert len(ctx.output_queue) == 0

            # Stop the execute model loop in parallel workers until there are
            # more requests to process. This avoids waiting indefinitely in
            # torch.distributed ops which may otherwise timeout, and unblocks
            # the RPC thread in the workers so that they can process any other
            # queued control plane messages, such as add/remove lora adapters.
            logger.debug("Stopping remote worker execution loop.")
            self.model_executor.stop_remote_worker_execution_loop()
        return ctx.request_outputs
                    
    def step(self) -> List[Union[RequestOutput, PoolingRequestOutput]]:
        with torch.cuda.stream(self.use_stream):
            if is_zero_no_thread():
                out = self.no_thread_step()
                if out is None: #the first step need launch twice
                    out = self.no_thread_step()
            else:
                out = self.zero_overhead_step()
                if out is None: #the first step need launch twice
                    out = self.zero_overhead_step()
        return out
    
    def _add_processed_request(
        self,
        request_id: str,
        processed_inputs: ProcessorInputs,
        params: Union[SamplingParams, PoolingParams],
        arrival_time: float,
        lora_request: Optional[LoRARequest],
        trace_headers: Optional[Mapping[str, str]] = None,
        priority: int = 0,
    ) -> Optional[SequenceGroup]:
        """Add a processed request to the engine's request pool.
        return the created sequence group.
        """
        if isinstance(params, SamplingParams) and params.n > 1:
            ParallelSampleSequenceGroup.add_request(
                request_id,
                self,
                params,
                processed_inputs=processed_inputs,
                arrival_time=arrival_time,
                lora_request=lora_request,
                trace_headers=trace_headers,
                priority=priority,
            )
            return None

        self._validate_model_inputs(processed_inputs, lora_request)
        # Create the sequences.
        block_size = self.cache_config.block_size
        seq_id = next(self.seq_counter)
        eos_token_id = self.input_preprocessor.get_eos_token_id(lora_request)

        encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs)

        seq = ZeroOverheadSequence(seq_id, decoder_inputs, block_size, eos_token_id,
                       lora_request)

        encoder_seq = (None if encoder_inputs is None else ZeroOverheadSequence(
            seq_id, encoder_inputs, block_size, eos_token_id, lora_request))

        # Create a SequenceGroup based on SamplingParams or PoolingParams
        if isinstance(params, SamplingParams):
            seq_group = self._create_sequence_group_with_sampling(
                request_id,
                seq,
                params,
                arrival_time=arrival_time,
                lora_request=lora_request,
                trace_headers=trace_headers,
                encoder_seq=encoder_seq,
                priority=priority)
        elif isinstance(params, PoolingParams):
            seq_group = self._create_sequence_group_with_pooling(
                request_id,
                seq,
                params,
                arrival_time=arrival_time,
                lora_request=lora_request,
                encoder_seq=encoder_seq,
                priority=priority)
        else:
            raise ValueError(
                "Either SamplingParams or PoolingParams must be provided.")

        # Add the sequence group to the scheduler with least unfinished seqs.
        costs = [
            scheduler.get_num_unfinished_seq_groups()
            for scheduler in self.scheduler
        ]
        min_cost_scheduler = self.scheduler[costs.index(min(costs))]
        min_cost_scheduler.add_seq_group(seq_group)

        return seq_group