chore: import upstream snapshot with attribution

python/tokenspeed/runtime/engine/request_handler.py

@@ -0,0 +1,500 @@
+# 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()