lightseekorg--tokenspeed/python/tokenspeed/runtime/engine/request_handler.py

# Copyright (c) 2026 LightSeek Foundation
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

from __future__ import annotations

import logging
import os
from pathlib import Path
from typing import TYPE_CHECKING

import torch
import zmq
from viztracer import VizTracer

from tokenspeed.runtime.distributed.process_group_manager import (
    process_group_manager as pg_manager,
)
from tokenspeed.runtime.engine.generation_output_processor import RequestState
from tokenspeed.runtime.engine.io_struct import (
    AbortReq,
    FlushCacheReqInput,
    FlushCacheReqOutput,
    GetInternalStateReq,
    GetInternalStateReqOutput,
    GetLoadReqInput,
    GetLoadReqOutput,
    IsSchedulerPausedReqInput,
    IsSleepingReqInput,
    PauseSchedulerReqInput,
    ProfileReq,
    ProfileReqOutput,
    ProfileReqType,
    ReleaseMemoryOccupationReqInput,
    ResumeMemoryOccupationReqInput,
    ResumeSchedulerReqInput,
    SetInternalStateReq,
    SetInternalStateReqOutput,
    TokenizedGenerateReqInput,
)
from tokenspeed.runtime.engine.request_types import FINISH_ABORT
from tokenspeed.runtime.engine.scheduler_utils import make_spec
from tokenspeed.runtime.execution.forward_batch_info import ForwardMode
from tokenspeed.runtime.grammar.grammar_manager import GrammarManager
from tokenspeed.runtime.multimodal.shm_transport import sync_shm_features
from tokenspeed.runtime.pd.base.bootstrap import BootstrapInfo
from tokenspeed.runtime.utils import broadcast_pyobj
from tokenspeed.runtime.utils.dispatch import TypeBasedDispatcher
from tokenspeed.runtime.utils.env import envs
from tokenspeed.runtime.utils.hf_transformers_utils import get_tokenizer

if TYPE_CHECKING:
    from tokenspeed.runtime.utils.server_args import ServerArgs

logger = logging.getLogger(__name__)


class RequestHandler:
    """
    1. Recv Reqs from ZMQ
    2. manage sessions
    """

    def __init__(
        self,
        server_args: ServerArgs,
        hf_eos_token_id,
        max_req_len: int,
        vocab_size: int,
        recv_func,
        send_func,
        get_load_fn=None,
        architectures: list[str] | None = None,
        pause_controller=None,
        memory_controller=None,
    ) -> None:

        self.forward_ct = 0
        self.server_args = server_args
        # Owns pause/resume state; shared with the event loop. See pause.py.
        self.pause_controller = pause_controller
        # Owns release/resume_memory_occupation (data plane). See
        # memory_occupation.py. Shares the pause controller's drain machinery.
        self.memory_controller = memory_controller

        mapping = server_args.mapping
        self.attn_tp_size = mapping.attn.tp_size
        self.attn_tp_rank = mapping.attn.tp_rank
        self.attn_tp_cpu_group = pg_manager.get_process_group(
            "gloo", mapping.attn.tp_group
        )
        self.attn_tp_src_rank = mapping.attn.tp_group[0]

        self.hf_eos_token_id = hf_eos_token_id
        self.max_req_len = max_req_len
        self.vocab_size = vocab_size
        self.get_load_fn = get_load_fn

        self.tokenizer = get_tokenizer(
            server_args.tokenizer,
            tokenizer_mode=server_args.tokenizer_mode,
            trust_remote_code=server_args.trust_remote_code,
            revision=server_args.revision,
            architectures=architectures,
        )

        self.recv_func = recv_func
        self.send_func = send_func

        self.control_request_dispatcher = TypeBasedDispatcher(
            [(ProfileReq, self.profile)]
        )

        self.grammar_manager = GrammarManager(
            self.server_args, self.tokenizer, self.vocab_size
        )

        self.init_profiler()

    def recv_reqs(self) -> list:
        """Receive results at attn_tp_rank = 0 and broadcast it to all other TP ranks."""
        if self.attn_tp_rank == 0:
            recv_reqs = []

            while True:
                try:
                    recv_req = self.recv_func.recv_pyobj(zmq.NOBLOCK)
                except zmq.ZMQError:
                    break
                recv_reqs.append(recv_req)
        else:
            recv_reqs = None

        if self.attn_tp_size != 1:
            recv_reqs = broadcast_pyobj(
                recv_reqs,
                self.attn_tp_rank,
                self.attn_tp_cpu_group,
                src=self.attn_tp_src_rank,
            )

        if recv_reqs:
            sync_shm_features(recv_reqs, self.attn_tp_cpu_group, self.attn_tp_size)

        return recv_reqs

    def process_requests(self, recv_reqs: list):
        """Dispatch control requests and return new generate request specs and states."""
        new_req_specs, req_states, bootstrap_infos, abort_rids = [], [], [], []
        for recv_req in recv_reqs:
            if isinstance(recv_req, TokenizedGenerateReqInput):
                req_spec, req_state, bootstrap_info = self.handle_generate_request(
                    recv_req
                )

                new_req_specs.append(req_spec)
                req_states.append(req_state)
                bootstrap_infos.append(bootstrap_info)
            elif isinstance(recv_req, ProfileReq):
                output = self.control_request_dispatcher(recv_req)
                if output is not None:
                    self.send_func.send_pyobj(output)
            elif isinstance(recv_req, AbortReq):
                logger.debug("AbortReq for rid=%s", recv_req.rid)
                abort_rids.append(recv_req.rid)
            elif isinstance(recv_req, FlushCacheReqInput):
                # Prefix cache is owned by the scheduler path; acknowledge the
                # control request here so API callers still get a typed reply.
                self.send_func.send_pyobj(FlushCacheReqOutput(success=True))
            elif isinstance(recv_req, PauseSchedulerReqInput):
                # State change + reply (abort/wait replies are deferred by the
                # controller until the event loop observes a drained scheduler).
                self.pause_controller.handle_pause(recv_req)
            elif isinstance(recv_req, ResumeSchedulerReqInput):
                self.pause_controller.handle_resume(recv_req)
            elif isinstance(recv_req, IsSchedulerPausedReqInput):
                self.pause_controller.handle_is_paused(recv_req)
            elif isinstance(recv_req, ReleaseMemoryOccupationReqInput):
                # Deferred: pauses + drains, then frees GPU memory and replies.
                self.memory_controller.handle_release(recv_req)
            elif isinstance(recv_req, ResumeMemoryOccupationReqInput):
                self.memory_controller.handle_resume(recv_req)
            elif isinstance(recv_req, IsSleepingReqInput):
                self.memory_controller.handle_is_sleeping(recv_req)
            elif isinstance(recv_req, GetInternalStateReq):
                self.send_func.send_pyobj(GetInternalStateReqOutput(internal_state={}))
            elif isinstance(recv_req, SetInternalStateReq):
                self.send_func.send_pyobj(
                    SetInternalStateReqOutput(updated=False, server_args={})
                )
            elif isinstance(recv_req, GetLoadReqInput):
                if self.get_load_fn is not None:
                    self.send_func.send_pyobj(self.get_load_fn())
                else:
                    self.send_func.send_pyobj(GetLoadReqOutput())
            else:
                raise NotImplementedError(f"Unsupported request type: {type(recv_req)}")
        return new_req_specs, req_states, bootstrap_infos, abort_rids

    def handle_generate_request(
        self,
        recv_req: TokenizedGenerateReqInput,
    ):
        if recv_req.bootstrap_port is None:
            recv_req.bootstrap_port = self.server_args.disaggregation_bootstrap_port

        req_spec = make_spec(
            rid=recv_req.rid,
            tokens=recv_req.input_ids,
        )
        req_state = RequestState.from_recv_req(
            recv_req,
            tokenizer=self.tokenizer,
            eos_token_ids=self.hf_eos_token_id,
        )

        if (
            recv_req.session_params is not None
            and recv_req.session_params.id is not None
        ):
            req_state.finished_reason = FINISH_ABORT(
                f"Invalid request: session id {recv_req.session_params.id} does not exist"
            )
            return (
                req_spec,
                req_state,
                BootstrapInfo(
                    recv_req.bootstrap_host,
                    recv_req.bootstrap_port,
                    recv_req.bootstrap_room,
                ),
            )

        req_state.sampling_params.max_new_tokens = min(
            (
                req_state.sampling_params.max_new_tokens
                if req_state.sampling_params.max_new_tokens is not None
                else 1 << 30
            ),
            self.max_req_len - len(req_state.prompt_input_ids) - 1,
        )
        return (
            req_spec,
            req_state,
            BootstrapInfo(
                recv_req.bootstrap_host,
                recv_req.bootstrap_port,
                recv_req.bootstrap_room,
            ),
        )

    # ------------------------------------------------------------------
    # Profiling: torch / cuda / viztracer / mem-snapshot, driven by
    # /start_profile and /stop_profile control requests.
    # ------------------------------------------------------------------

    def init_profiler(self):
        self.torch_profiler = None
        self.profiler_output_dir: str | None = None
        self.profiler_activities: list[str] | None = None
        self.profile_id: str | None = None
        self.profiler_start_forward_ct: int | None = None
        self.profiler_target_forward_ct: int | None = None
        self.profiler_target_prefill_ct: int | None = None
        self.profiler_target_decode_ct: int | None = None
        self.profiler_prefill_ct: int | None = None
        self.profiler_decode_ct: int | None = None
        self.profile_by_stage: bool = False
        self.profile_in_progress: bool = False
        self.viztracer = None

    def init_profile(
        self,
        output_dir: str | None,
        start_step: int | None,
        num_steps: int | None,
        activities: list[str] | None,
        with_stack: bool | None,
        record_shapes: bool | None,
        profile_by_stage: bool,
        profile_id: str,
    ) -> ProfileReqOutput:
        if self.profile_in_progress:
            return ProfileReqOutput(
                success=False,
                message="Profiling is already in progress. Call /stop_profile first.",
            )

        self.profile_by_stage = profile_by_stage

        if output_dir is None:
            output_dir = envs.TOKENSPEED_PROFILER_DIR.get()
        if activities is None:
            activities = ["CPU", "GPU"]

        self.profiler_output_dir = output_dir
        self.torch_profiler_with_stack = with_stack
        self.torch_profiler_record_shapes = record_shapes
        self.profiler_activities = activities
        self.profile_id = profile_id

        if start_step:
            self.profiler_start_forward_ct = max(start_step, self.forward_ct + 1)

        if num_steps:
            if self.profile_by_stage:
                self.profiler_target_prefill_ct = num_steps
                self.profiler_target_decode_ct = num_steps
                self.profiler_prefill_ct = 0
                self.profiler_decode_ct = 0
            elif start_step:
                self.profiler_target_forward_ct = (
                    self.profiler_start_forward_ct + num_steps
                )
            else:
                self.profiler_target_forward_ct = self.forward_ct + num_steps
            # The caller will be notified when reaching profiler_target_forward_ct
        else:
            self.profiler_target_forward_ct = None

        return ProfileReqOutput(success=True, message="Succeeded")

    def start_profile(
        self, stage: ForwardMode | None = None
    ) -> ProfileReqOutput | None:
        stage_str = f" for {stage.name}" if stage else ""
        stage_suffix = f"-{stage.name}" if stage else ""

        activities = self.profiler_activities
        with_stack = self.torch_profiler_with_stack
        record_shapes = self.torch_profiler_record_shapes

        activity_map = {
            "CPU": torch.profiler.ProfilerActivity.CPU,
            "GPU": torch.profiler.ProfilerActivity.CUDA,
        }
        torchprof_activities = [
            activity_map[a] for a in activities if a in activity_map
        ]

        if torchprof_activities:
            self.torch_profiler = torch.profiler.profile(
                activities=torchprof_activities,
                with_stack=with_stack if with_stack is not None else True,
                record_shapes=record_shapes if record_shapes is not None else False,
            )
            self.torch_profiler.start()

        if "MEM" in activities:
            torch.cuda.memory._record_memory_history(max_entries=100000)

        if "CUDA_PROFILER" in activities:
            torch.cuda.cudart().cudaProfilerStart()

        if "VIZTRACER" in activities:
            Path(self.profiler_output_dir).mkdir(parents=True, exist_ok=True)
            self.viztracer = VizTracer(
                output_file=os.path.join(
                    self.profiler_output_dir,
                    f"{self.profile_id}-TP-{self.attn_tp_rank}{stage_suffix}.viztracer.json",
                ),
                min_duration=int(
                    os.environ.get("TOKENSPEED_VIZTRACER_MIN_DURATION_US", "100")
                ),
                log_async=True,
            )
            self.viztracer.start()

        if activities:
            if activities != ["CUDA_PROFILER"]:
                logger.info(
                    "Profiling starts%s. Traces will be saved to: %s (with profile id: %s)",
                    stage_str,
                    self.profiler_output_dir,
                    self.profile_id,
                )
            self.profile_in_progress = True

        return ProfileReqOutput(success=True, message="Succeeded")

    def stop_profile(self, stage: ForwardMode | None = None) -> ProfileReqOutput | None:
        if not self.profile_in_progress:
            return ProfileReqOutput(
                success=False,
                message="Profiling is not in progress. Call /start_profile first.",
            )

        Path(self.profiler_output_dir).mkdir(parents=True, exist_ok=True)

        stage_suffix = f"-{stage.name}" if stage else ""
        logger.info("Stop profiling%s...", stage_suffix)

        if self.torch_profiler is not None:
            self.torch_profiler.stop()
            self.torch_profiler.export_chrome_trace(
                os.path.join(
                    self.profiler_output_dir,
                    f"{self.profile_id}-TP-{self.attn_tp_rank}{stage_suffix}.trace.json.gz",
                )
            )
            torch.distributed.barrier(self.attn_tp_cpu_group)

        if self.profiler_activities is not None and "MEM" in self.profiler_activities:
            memory_profile_path = os.path.join(
                self.profiler_output_dir,
                f"{self.profile_id}-TP-{self.attn_tp_rank}-memory{stage_suffix}.pickle",
            )
            torch.cuda.memory._dump_snapshot(memory_profile_path)
            torch.cuda.memory._record_memory_history(enabled=None)

        if "CUDA_PROFILER" in self.profiler_activities:
            torch.cuda.cudart().cudaProfilerStop()

        if "VIZTRACER" in self.profiler_activities and self.viztracer is not None:
            self.viztracer.stop()
            self.viztracer.save()
            self.viztracer = None

        if self.profiler_activities and self.profiler_activities != ["CUDA_PROFILER"]:
            logger.info(
                "Profiling done. Traces are saved to: %s", self.profiler_output_dir
            )

        self.torch_profiler = None
        self.profile_in_progress = False
        self.profiler_start_forward_ct = None

        return ProfileReqOutput(success=True, message="Succeeded.")

    def _profile_batch_predicate(self, forward_mode=None):
        """Check and toggle profiling based on forward step count.

        Args:
            forward_mode: Optional ForwardMode for stage-based profiling.
                Not needed for step-count-based profiling.
        """
        if self.profile_by_stage and forward_mode is not None:
            if forward_mode.is_extend_or_mixed():
                if self.profiler_prefill_ct == 0:
                    self.start_profile(forward_mode)
                self.profiler_prefill_ct += 1
                if self.profiler_prefill_ct > self.profiler_target_prefill_ct:
                    if self.profile_in_progress:
                        self.stop_profile(stage=ForwardMode.EXTEND)
            elif forward_mode.is_decode():
                if self.profiler_decode_ct == 0:
                    if self.profile_in_progress:
                        self.stop_profile(ForwardMode.EXTEND)
                    self.start_profile(forward_mode)
                self.profiler_decode_ct += 1
                if self.profiler_decode_ct > self.profiler_target_decode_ct:
                    if self.profile_in_progress:
                        self.stop_profile(stage=ForwardMode.DECODE)
            elif forward_mode.is_idle():
                pass
        else:
            if (
                self.profiler_target_forward_ct
                and self.profiler_target_forward_ct <= self.forward_ct
            ):
                self.stop_profile()
            if (
                self.profiler_start_forward_ct
                and self.profiler_start_forward_ct == self.forward_ct
            ):
                self.start_profile()

    def profile(self, recv_req: ProfileReq):
        if recv_req.type == ProfileReqType.START_PROFILE:
            res = self.init_profile(
                recv_req.output_dir,
                recv_req.start_step,
                recv_req.num_steps,
                recv_req.activities,
                recv_req.with_stack,
                recv_req.record_shapes,
                recv_req.profile_by_stage,
                recv_req.profile_id,
            )
            if not res.success or recv_req.profile_by_stage or recv_req.start_step:
                return res
            return self.start_profile()
        else:
            return self.stop_profile()