chore: import upstream snapshot with attribution

python/tokenspeed/runtime/engine/io_struct.py

@@ -0,0 +1,954 @@
+# 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.
+
+"""
+The definition of objects transferred between different
+processes (TokenizerManager, DetokenizerManager, Controller).
+"""
+
+import copy
+import uuid
+from abc import ABC
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Literal
+
+from tokenspeed.runtime.engine.request_types import BaseFinishReason
+from tokenspeed.runtime.sampling.sampling_params import SamplingParams
+
+
+def _require(condition: bool, message: str) -> None:
+    if not condition:
+        raise ValueError(message)
+
+
+@dataclass
+class BaseReq(ABC):
+    rid: str | list[str] | None = field(default=None)
+    http_worker_ipc: str | None = field(default=None)
+
+    def regenerate_rid(self):
+        """Generate a new request ID and return it."""
+        if isinstance(self.rid, list):
+            self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))]
+        else:
+            self.rid = uuid.uuid4().hex
+        return self.rid
+
+
+@dataclass
+class SessionParams:
+    id: str | None = None
+    rid: str | None = None
+    offset: int | None = None
+    replace: bool | None = None
+
+
+@dataclass
+class GenerateReqInput:
+    # The input prompt. It can be a single prompt or a batch of prompts.
+    text: list[str] | str | None = None
+    # The token ids for text; one can specify either text or input_ids
+    input_ids: list[list[int]] | list[int] | None = None
+    input_multi_ids: list[list[int]] | list[list[int]] | None = None
+    # The embeddings for input_ids; one can specify either text or input_ids or input_embeds.
+    input_embeds: list[list[list[float]]] | list[list[float]] | None = None
+    # Pre-built MultimodalInputs (already produced by an upstream preprocessor,
+    # e.g. SMG's Rust crates/multimodal pipeline). The engine's InputProcessor
+    # uses this directly (it does no in-process image preprocessing). input_ids
+    # must already contain expanded image placeholder tokens at the right
+    # offsets — the gateway is responsible for that. Typed as Any to avoid a
+    # circular import on MultimodalInputs.
+    precomputed_multimodal_inputs: Any | None = None
+    # The sampling_params. See descriptions below.
+    sampling_params: list[dict] | dict | None = None
+    input_extra_infos: list[dict] | dict | None = None
+    # Optional client label for logging; defaults to `rid`. Safe to reuse.
+    user_rid: list[str] | str | None = None
+    # Routing id; always server-assigned during normalize, never caller-settable.
+    rid: list[str] | str | None = field(default=None, init=False)
+    # --- Logprob request (two dialects, one compute path) ---
+    # vLLM-compatible requests use ``sampling_params["logprobs"]``;
+    # SGLang-compatible requests use the legacy fields below. A request uses
+    # one dialect; the response is rendered to match (override with
+    # ``logprob_format``).
+    return_logprob: list[bool] | bool | None = None
+    # Start location in the prompt for prompt logprobs. -1 (default) = output
+    # tokens only.
+    logprob_start_len: list[int] | int | None = None
+    # Number of top logprobs per position.
+    top_logprobs_num: list[int] | int | None = None
+    # Specific token ids to score per position.
+    token_ids_logprob: list[list[int]] | list[int] | None = None
+    # Detokenize tokens in the returned logprobs.
+    return_text_in_logprobs: bool = False
+    # Output rendering dialect: "vllm" | "sglang" | "both". None = auto (match
+    # the request dialect: vllm if sampling_params.logprobs is set, else sglang).
+    logprob_format: list[str | None] | str | None = None
+    # Whether to stream output.
+    stream: bool = False
+    # Whether to log metrics for this request (e.g. health_generate calls do not log metrics)
+    log_metrics: bool = True
+
+    # Session info for continual prompting
+    session_params: list[dict] | dict | None = None
+
+    # Custom logit processor for advanced sampling control. Must be a serialized instance
+    # of `CustomLogitProcessor` in python/tokenspeed/runtime/sampling/custom_logit_processor.py
+    # Use the processor's `to_str()` method to generate the serialized string.
+    custom_logit_processor: list[str | None] | str | None = None
+
+    # Whether to return hidden states
+    return_hidden_states: bool = False
+
+    # For disaggregated inference
+    bootstrap_host: list[str] | str | None = None
+    bootstrap_port: list[int] | int | None = None
+    bootstrap_room: list[int] | int | None = None
+
+    def normalize_batch_and_arguments(self):
+        if (
+            self.text is None and self.input_ids is None and self.input_embeds is None
+        ) or (
+            self.text is not None
+            and self.input_ids is not None
+            and self.input_embeds is not None
+        ):
+            raise ValueError(
+                "Either text, input_ids or input_embeds should be provided."
+            )
+
+        # Derive the batch size
+        if self.text is not None:
+            if isinstance(self.text, str):
+                self.is_single = True
+                self.batch_size = 1
+            else:
+                self.is_single = False
+                self.batch_size = len(self.text)
+            self.input_embeds = None
+        elif self.input_ids is not None:
+            if isinstance(self.input_ids[0], int):
+                self.is_single = True
+                self.batch_size = 1
+            else:
+                self.is_single = False
+                self.batch_size = len(self.input_ids)
+            self.input_embeds = None
+        else:
+            _require(
+                isinstance(self.input_embeds, list), "input_embeds should be a list."
+            )
+            if isinstance(self.input_embeds[0][0], float):
+                # list[list[float]]
+                self.is_single = True
+                self.batch_size = 1
+            else:
+                # list[list[list[float]]]
+                _require(
+                    isinstance(self.input_embeds[0][0], list),
+                    "input_embeds should be a list of float lists.",
+                )
+                _require(
+                    isinstance(self.input_embeds[0][0][0], float),
+                    "input_embeds should contain floats.",
+                )
+                self.is_single = False
+                self.batch_size = len(self.input_embeds)
+
+        # Handle parallel sampling. Pop "n" out of sampling_params so the
+        # downstream SamplingParams(**dict) construction doesn't see it —
+        # "n" is a request-level fan-out knob, not a per-sample field.
+        if self.sampling_params is None:
+            self.parallel_sample_num = 1
+        elif isinstance(self.sampling_params, dict):
+            self.parallel_sample_num = self.sampling_params.get("n", 1)
+        else:  # isinstance(self.sampling_params, list):
+            self.parallel_sample_num = self.sampling_params[0].get("n", 1)
+            for sp in self.sampling_params[1:]:
+                _require(
+                    self.parallel_sample_num == sp.get("n", 1),
+                    "The parallel_sample_num should be the same for all samples in sample params.",
+                )
+
+        if self.parallel_sample_num > 1 and self.is_single:
+            self.is_single = False
+            if self.text is not None:
+                self.text = [self.text]
+            if self.input_ids is not None:
+                self.input_ids = [self.input_ids]
+            if self.input_multi_ids is not None:
+                self.input_multi_ids = [self.input_multi_ids]
+            if self.input_embeds is not None:
+                self.input_embeds = [self.input_embeds]
+
+        # Fill in default arguments
+        if self.is_single:
+            if self.sampling_params is None:
+                self.sampling_params = {}
+            if self.rid is None:
+                self.rid = uuid.uuid4().hex
+            if self.user_rid is None:
+                self.user_rid = self.rid
+            else:
+                if isinstance(self.user_rid, list):
+                    _require(
+                        len(self.user_rid) == 1,
+                        "user_rid list should have length 1 for single request.",
+                    )
+                    self.user_rid = self.user_rid[0]
+                _require(isinstance(self.user_rid, str), "user_rid should be a str.")
+            if self.return_logprob is None:
+                self.return_logprob = False
+            if self.logprob_start_len is None:
+                self.logprob_start_len = -1
+            if self.top_logprobs_num is None:
+                self.top_logprobs_num = 0
+            if not self.token_ids_logprob:  # covers both None and []
+                self.token_ids_logprob = None
+            if isinstance(self.input_extra_infos, dict):
+                self.input_extra_infos = [self.input_extra_infos]
+        else:
+            if self.parallel_sample_num == 1:
+                num = self.batch_size
+            else:
+                # Expand parallel_sample_num
+                num = self.batch_size * self.parallel_sample_num
+
+            if self.sampling_params is None:
+                self.sampling_params = [{} for _ in range(num)]
+            elif not isinstance(self.sampling_params, list):
+                self.sampling_params = [dict(self.sampling_params) for _ in range(num)]
+
+            if self.rid is None:
+                self.rid = [uuid.uuid4().hex for _ in range(num)]
+            else:
+                _require(isinstance(self.rid, list), "The rid should be a list.")
+            if self.user_rid is None:
+                self.user_rid = list(self.rid)
+            elif isinstance(self.user_rid, str):
+                self.user_rid = [self.user_rid] * num
+            else:
+                _require(
+                    isinstance(self.user_rid, list) and len(self.user_rid) == num,
+                    "user_rid should be a str or a list of matching length.",
+                )
+
+            if self.return_logprob is None:
+                self.return_logprob = [False] * num
+            elif not isinstance(self.return_logprob, list):
+                self.return_logprob = [self.return_logprob] * num
+            else:
+                _require(
+                    self.parallel_sample_num == 1,
+                    "return_logprob cannot be a list when n > 1.",
+                )
+
+            if self.logprob_start_len is None:
+                self.logprob_start_len = [-1] * num
+            elif not isinstance(self.logprob_start_len, list):
+                self.logprob_start_len = [self.logprob_start_len] * num
+            else:
+                _require(
+                    self.parallel_sample_num == 1,
+                    "logprob_start_len cannot be a list when n > 1.",
+                )
+
+            if self.top_logprobs_num is None:
+                self.top_logprobs_num = [0] * num
+            elif not isinstance(self.top_logprobs_num, list):
+                self.top_logprobs_num = [self.top_logprobs_num] * num
+            else:
+                _require(
+                    self.parallel_sample_num == 1,
+                    "top_logprobs_num cannot be a list when n > 1.",
+                )
+
+            if not self.token_ids_logprob:  # covers both None and []
+                self.token_ids_logprob = [None] * num
+            elif not isinstance(self.token_ids_logprob, list):
+                self.token_ids_logprob = [[self.token_ids_logprob] for _ in range(num)]
+            elif not isinstance(self.token_ids_logprob[0], list):
+                self.token_ids_logprob = [
+                    copy.deepcopy(self.token_ids_logprob) for _ in range(num)
+                ]
+            else:
+                _require(
+                    self.parallel_sample_num == 1,
+                    "token_ids_logprob cannot be nested lists when n > 1.",
+                )
+
+            if self.logprob_format is None or isinstance(self.logprob_format, str):
+                self.logprob_format = [self.logprob_format] * num
+
+            if self.custom_logit_processor is None:
+                self.custom_logit_processor = [None] * num
+            elif not isinstance(self.custom_logit_processor, list):
+                self.custom_logit_processor = [self.custom_logit_processor] * num
+            else:
+                _require(
+                    self.parallel_sample_num == 1,
+                    "custom_logit_processor cannot be a list when n > 1.",
+                )
+
+            if self.bootstrap_host is None:
+                self.bootstrap_host = [None] * num
+            elif not isinstance(self.bootstrap_host, list):
+                self.bootstrap_host = [self.bootstrap_host] * num
+            else:
+                _require(
+                    self.parallel_sample_num == 1,
+                    "bootstrap_host cannot be a list when n > 1.",
+                )
+
+            if self.bootstrap_port is None:
+                self.bootstrap_port = [None] * num
+            elif not isinstance(self.bootstrap_port, list):
+                self.bootstrap_port = [self.bootstrap_port] * num
+            else:
+                _require(
+                    self.parallel_sample_num == 1,
+                    "bootstrap_port cannot be a list when n > 1.",
+                )
+
+            if self.bootstrap_room is None:
+                self.bootstrap_room = [None] * num
+            elif not isinstance(self.bootstrap_room, list):
+                self.bootstrap_room = [self.bootstrap_room] * num
+            else:
+                _require(
+                    self.parallel_sample_num == 1,
+                    "bootstrap_room cannot be a list when n > 1.",
+                )
+
+        # Other checks
+        if self.session_params is not None:
+            _require(
+                isinstance(self.session_params, dict)
+                or isinstance(self.session_params[0], dict),
+                "session_params should be a dict or a list of dicts.",
+            )
+
+    def regenerate_rid(self):
+        self.rid = uuid.uuid4().hex
+        return self.rid
+
+    def __getitem__(self, i):
+        sub = GenerateReqInput(
+            text=self.text[i] if self.text is not None else None,
+            input_ids=self.input_ids[i] if self.input_ids is not None else None,
+            # precomputed_multimodal_inputs is a single prompt's MM; the SMG
+            # path only clears is_single via n>1 (batch_size == 1), so all n
+            # parallel samples correctly share it. Without this the image is
+            # silently dropped on the n>1 fan-out (placeholders -> text path).
+            precomputed_multimodal_inputs=self.precomputed_multimodal_inputs,
+            input_multi_ids=(
+                self.input_multi_ids[i] if self.input_multi_ids is not None else None
+            ),
+            input_embeds=(
+                self.input_embeds[i] if self.input_embeds is not None else None
+            ),
+            input_extra_infos=(
+                self.input_extra_infos[i]
+                if self.input_extra_infos is not None
+                else None
+            ),
+            sampling_params=self.sampling_params[i],
+            user_rid=self.user_rid[i],
+            return_logprob=self.return_logprob[i],
+            logprob_start_len=self.logprob_start_len[i],
+            top_logprobs_num=self.top_logprobs_num[i],
+            token_ids_logprob=self.token_ids_logprob[i],
+            return_text_in_logprobs=self.return_text_in_logprobs,
+            logprob_format=self.logprob_format[i],
+            stream=self.stream,
+            log_metrics=self.log_metrics,
+            custom_logit_processor=(
+                self.custom_logit_processor[i]
+                if self.custom_logit_processor is not None
+                else None
+            ),
+            return_hidden_states=self.return_hidden_states,
+            # if `__getitem__` is called, the bootstrap_host, bootstrap_port, bootstrap_room must be a list
+            bootstrap_host=(
+                self.bootstrap_host[i] if self.bootstrap_host is not None else None
+            ),
+            bootstrap_port=(
+                self.bootstrap_port[i] if self.bootstrap_port is not None else None
+            ),
+            bootstrap_room=(
+                self.bootstrap_room[i] if self.bootstrap_room is not None else None
+            ),
+        )
+        sub.rid = self.rid[i]
+        return sub
+
+
+@dataclass
+class TokenizedGenerateReqInput:
+    # The request id
+    rid: str
+    # The input text
+    input_text: str
+    # The input token ids
+    input_ids: list[int]
+    # The sampling parameters
+    sampling_params: SamplingParams
+    # Whether to return the sampled token's logprob for this request.
+    return_logprob: bool
+    # Internal carry-over fields kept for pipeline/PD compatibility. The vLLM
+    # output-logprob API only drives ``return_logprob``; InputProcessor sets
+    # these to neutral values (logprob_start_len=-1, top_logprobs_num=0,
+    # token_ids_logprob=None) since prompt logprobs, output top-k, and token-id
+    # logprobs are not supported.
+    logprob_start_len: int
+    top_logprobs_num: int
+    token_ids_logprob: list[int]
+    # Whether to stream output
+    stream: bool
+
+    # The input embeds
+    input_embeds: list[list[list[float]]] | list[list[float]] | None = None
+
+    # Session info for continual prompting
+    session_params: SessionParams | None = None
+
+    # Custom logit processor for advanced sampling control. Must be a serialized instance
+    # of `CustomLogitProcessor` in python/tokenspeed/runtime/sampling/custom_logit_processor.py
+    # Use the processor's `to_str()` method to generate the serialized string.
+    custom_logit_processor: str | None = None
+
+    # Whether to return hidden states
+    return_hidden_states: bool = False
+
+    # Time at object instantiated
+    created_time: float = 0.0
+
+    # For disaggregated inference
+    bootstrap_host: str | None = None
+    bootstrap_port: int | None = None
+    bootstrap_room: int | None = None
+
+    input_multi_ids: list[list[int]] = None
+    input_extra_infos: list[dict] | None = None
+    # Original prompt ids before multimodal pad/hash replacement. The scheduler
+    # uses input_ids, while detokenization must use these tokenizer-valid ids.
+    input_ids_unpadded: list[int] | None = None
+    multimodal_inputs: Any | None = None
+
+
+@dataclass
+class EmbeddingReqInput:
+    # The input prompt. It can be a single prompt or a batch of prompts.
+    text: list[str] | str | None = None
+    # The token ids for text; one can either specify text or input_ids.
+    input_ids: list[list[int]] | list[int] | None = None
+    # Optional client label for logging; defaults to `rid`. Safe to reuse.
+    user_rid: list[str] | str | None = None
+    # Routing id; always server-assigned during normalize, never caller-settable.
+    rid: list[str] | str | None = field(default=None, init=False)
+    # Optional placeholder so non-generation callers can still instantiate the
+    # shared request shape without real sampling params.
+    sampling_params: list[dict] | dict = None
+    # Optional placeholder for callers that do not provide input embeddings.
+    input_embeds: list[list[list[float]]] | list[list[float]] | None = None
+    # Whether to log metrics for this request (e.g. health_generate calls do not log metrics)
+    log_metrics: bool = True
+
+    def normalize_batch_and_arguments(self):
+        if (self.text is None and self.input_ids is None) or (
+            self.text is not None and self.input_ids is not None
+        ):
+            raise ValueError("Either text or input_ids should be provided.")
+
+        # Derive the batch size
+        if self.text is not None:
+            if isinstance(self.text, str):
+                self.is_single = True
+                self.batch_size = 1
+            else:
+                self.is_single = False
+                self.batch_size = len(self.text)
+        else:
+            if isinstance(self.input_ids[0], int):
+                self.is_single = True
+                self.batch_size = 1
+            else:
+                self.is_single = False
+                self.batch_size = len(self.input_ids)
+
+        # Fill in default arguments
+        if self.is_single:
+            if self.rid is None:
+                self.rid = uuid.uuid4().hex
+            if self.user_rid is None:
+                self.user_rid = self.rid
+            else:
+                if isinstance(self.user_rid, list):
+                    _require(
+                        len(self.user_rid) == 1,
+                        "user_rid list should have length 1 for single request.",
+                    )
+                    self.user_rid = self.user_rid[0]
+                _require(isinstance(self.user_rid, str), "user_rid should be a str.")
+            if self.sampling_params is None:
+                self.sampling_params = {}
+            self.sampling_params["max_new_tokens"] = 0
+        else:
+            if self.rid is None:
+                self.rid = [uuid.uuid4().hex for _ in range(self.batch_size)]
+            else:
+                _require(isinstance(self.rid, list), "The rid should be a list.")
+            if self.user_rid is None:
+                self.user_rid = list(self.rid)
+            elif isinstance(self.user_rid, str):
+                self.user_rid = [self.user_rid] * self.batch_size
+            else:
+                _require(
+                    isinstance(self.user_rid, list)
+                    and len(self.user_rid) == self.batch_size,
+                    "user_rid should be a str or a list of matching length.",
+                )
+
+            if self.sampling_params is None:
+                self.sampling_params = [{} for _ in range(self.batch_size)]
+            for i in range(self.batch_size):
+                self.sampling_params[i]["max_new_tokens"] = 0
+
+    def regenerate_rid(self):
+        self.rid = uuid.uuid4().hex
+        return self.rid
+
+    def __getitem__(self, i):
+        sub = EmbeddingReqInput(
+            text=self.text[i] if self.text is not None else None,
+            input_ids=self.input_ids[i] if self.input_ids is not None else None,
+            sampling_params=self.sampling_params[i],
+            user_rid=self.user_rid[i],
+        )
+        sub.rid = self.rid[i]
+        return sub
+
+
+@dataclass
+class TokenizedEmbeddingReqInput:
+    # The request id
+    rid: str
+    # The input text
+    input_text: str
+    # The input token ids
+    input_ids: list[int]
+    # Placeholder sampling params field so request metadata can share one shape
+    # with generation-oriented code paths.
+    sampling_params: SamplingParams
+    # Time at object instantiated
+    created_time: float
+
+
+@dataclass
+class BatchTokenIDOut:
+    # The request id
+    rids: list[str]
+    # The finish reason
+    finished_reasons: list[BaseFinishReason]
+    # For incremental decoding
+    decoded_texts: list[str]
+    decode_ids: list[list[int]]
+    read_offsets: list[int]
+    # Only used when `--skip-tokenizer-init` is on
+    output_ids: list[int] | None
+    output_multi_ids: list[int] | None
+    # Detokenization configs
+    skip_special_tokens: list[bool]
+    spaces_between_special_tokens: list[bool]
+    no_stop_trim: list[bool]
+
+    # Token counts
+    prompt_tokens: list[int]
+    completion_tokens: list[int]
+    cached_tokens: list[int]
+    spec_verify_ct: list[int]
+
+    # Logprobs
+    input_token_logprobs_val: list[float]
+    input_token_logprobs_idx: list[int]
+    output_token_logprobs_val: list[float]
+    output_token_logprobs_idx: list[int]
+    input_top_logprobs_val: list[list]
+    input_top_logprobs_idx: list[list]
+    output_top_logprobs_val: list[list]
+    output_top_logprobs_idx: list[list]
+    input_token_ids_logprobs_val: list[list]
+    input_token_ids_logprobs_idx: list[list]
+    output_token_ids_logprobs_val: list[list]
+    output_token_ids_logprobs_idx: list[list]
+
+    # Hidden states
+    output_hidden_states: list[list[float]]
+    batch_accept_draft_tokens: list[float]
+
+    # Store some custom information, such as decoding status in multimodal scenarios, etc.
+    output_extra_infos: list[dict[str, Any]]
+
+    generated_time: int
+
+
+@dataclass
+class BatchStrOut:
+    # The request id
+    rids: list[str]
+    # The finish reason
+    finished_reasons: list[dict]
+    # The output decoded strings
+    output_strs: list[str]
+    # The token ids
+    output_ids: list[int] | None
+
+    # Token counts
+    prompt_tokens: list[int]
+    completion_tokens: list[int]
+    cached_tokens: list[int]
+    spec_verify_ct: list[int]
+
+    # Logprobs
+    input_token_logprobs_val: list[float]
+    input_token_logprobs_idx: list[int]
+    output_token_logprobs_val: list[float]
+    output_token_logprobs_idx: list[int]
+    input_top_logprobs_val: list[list]
+    input_top_logprobs_idx: list[list]
+    output_top_logprobs_val: list[list]
+    output_top_logprobs_idx: list[list]
+    input_token_ids_logprobs_val: list[list]
+    input_token_ids_logprobs_idx: list[list]
+    output_token_ids_logprobs_val: list[list]
+    output_token_ids_logprobs_idx: list[list]
+
+    # Hidden states
+    output_hidden_states: list[list[float]]
+    batch_accept_draft_tokens: list[float]
+
+    # Store some custom information, such as decoding status in multimodal scenarios, etc.
+    output_extra_infos: list[dict[str, Any]]
+
+    generated_time: int
+
+
+@dataclass
+class BatchEmbeddingOut:
+    # The request id
+    rids: list[str]
+    # The finish reason
+    finished_reasons: list[BaseFinishReason]
+    # The output embedding
+    embeddings: list[list[float]] | list[dict]
+    # Token counts
+    prompt_tokens: list[int]
+
+
+@dataclass
+class FlushCacheReqInput:
+    pass
+
+
+@dataclass
+class FlushCacheReqOutput:
+    success: bool
+
+
+# How a pause should treat in-flight requests.
+# - "abort": kill in-flight requests immediately, then stop admitting new ones.
+# - "wait":  stop admitting new ones, keep stepping until running requests drain.
+# - "keep":  freeze everything in place; resume picks up where it left off.
+PauseMode = Literal["abort", "wait", "keep"]
+
+
+@dataclass
+class PauseSchedulerReqInput:
+    # See PauseMode for how each mode treats in-flight requests.
+    mode: PauseMode = "abort"
+
+
+@dataclass
+class PauseSchedulerReqOutput:
+    success: bool
+    message: str = ""
+
+
+@dataclass
+class ResumeSchedulerReqInput:
+    pass
+
+
+@dataclass
+class ResumeSchedulerReqOutput:
+    success: bool
+    message: str = ""
+
+
+@dataclass
+class IsSchedulerPausedReqInput:
+    pass
+
+
+@dataclass
+class IsSchedulerPausedReqOutput:
+    is_paused: bool
+
+
+@dataclass
+class UpdateWeightFromDiskReqInput:
+    # The model path with the new weights
+    model_path: str
+    # The format to load the weights
+    load_format: str | None = None
+
+
+@dataclass
+class UpdateWeightFromDiskReqOutput:
+    success: bool
+    message: str
+    # Number of paused requests during weight sync.
+    num_paused_requests: int | None = 0
+
+
+@dataclass
+class UpdateWeightsFromDistributedReqInput:
+    name: str
+    dtype: str
+    shape: list[int]
+
+
+@dataclass
+class UpdateWeightsFromDistributedReqOutput:
+    success: bool
+    message: str
+
+
+@dataclass
+class UpdateWeightsFromTensorReqInput:
+    serialized_named_tensors: bytes  # indeed Dict[str, torch.Tensor]
+    load_format: str | None
+    flush_cache: bool
+
+
+@dataclass
+class UpdateWeightsFromTensorReqOutput:
+    success: bool
+    message: str
+
+
+@dataclass
+class InitWeightsUpdateGroupReqInput:
+    # The master address
+    master_address: str
+    # The master port
+    master_port: int
+    # The rank offset
+    rank_offset: int
+    # The world size
+    world_size: int
+    # The group name
+    group_name: str = "weight_update_group"
+    # The backend
+    backend: str = "nccl"
+
+
+@dataclass
+class InitWeightsUpdateGroupReqOutput:
+    success: bool
+    message: str
+
+
+@dataclass
+class GetWeightsByNameReqInput:
+    name: str
+    truncate_size: int = 100
+
+
+@dataclass
+class GetWeightsByNameReqOutput:
+    parameter: list
+
+
+@dataclass
+class ReleaseMemoryOccupationReqInput:
+    # Memory regions to release. None ⇒ all ("weights" and "kv_cache").
+    tags: list[str] | None = None
+
+
+@dataclass
+class ReleaseMemoryOccupationReqOutput:
+    success: bool = True
+    message: str = ""
+
+
+@dataclass
+class ResumeMemoryOccupationReqInput:
+    # Memory regions to resume. None ⇒ all previously released tags.
+    tags: list[str] | None = None
+
+
+@dataclass
+class ResumeMemoryOccupationReqOutput:
+    success: bool = True
+    message: str = ""
+
+
+@dataclass
+class IsSleepingReqInput:
+    pass
+
+
+@dataclass
+class IsSleepingReqOutput:
+    is_sleeping: bool
+
+
+@dataclass
+class AbortReq:
+    # The request id
+    rid: str
+
+
+@dataclass
+class GetInternalStateReq:
+    pass
+
+
+@dataclass
+class GetInternalStateReqOutput:
+    internal_state: dict[Any, Any]
+
+
+@dataclass
+class SetInternalStateReq:
+    server_args: dict[str, Any]
+
+
+@dataclass
+class SetInternalStateReqOutput:
+    updated: bool
+    server_args: dict[str, Any]
+
+
+class ExpertDistributionReq(Enum):
+    START_RECORD = 1
+    STOP_RECORD = 2
+    DUMP_RECORD = 3
+
+
+@dataclass
+class ExpertDistributionReqOutput:
+    pass
+
+
+class ProfileReqType(Enum):
+    START_PROFILE = 1
+    STOP_PROFILE = 2
+
+
+@dataclass
+class ProfileReq:
+    type: ProfileReqType
+    output_dir: str | None = None
+    start_step: int | None = None
+    num_steps: int | None = None
+    activities: list[str] | None = None
+    profile_by_stage: bool = False
+    with_stack: bool | None = None
+    record_shapes: bool | None = None
+    profile_id: str | None = None
+
+
+@dataclass
+class ProfileReqOutput:
+    success: bool
+    message: str
+
+
+@dataclass
+class ConfigureLoggingReq:
+    log_requests: bool | None = None
+    log_requests_level: int | None = None
+    dump_requests_folder: str | None = None
+    dump_requests_threshold: int | None = None
+
+
+@dataclass
+class OpenSessionReqInput:
+    capacity_of_str_len: int
+    session_id: str | None = None
+
+
+@dataclass
+class CloseSessionReqInput:
+    session_id: str
+
+
+@dataclass
+class OpenSessionReqOutput:
+    session_id: str | None
+    success: bool
+
+
+@dataclass
+class HealthCheckOutput:
+    pass
+
+
+@dataclass
+class RpcReqInput:
+    method: str
+    parameters: dict | None = None
+
+
+@dataclass
+class RpcReqOutput:
+    success: bool
+    message: str
+
+
+@dataclass
+class GetLoadReqInput(BaseReq):
+    pass
+
+
+@dataclass
+class GetLoadReqOutput(BaseReq):
+    dp_rank: int = 0
+    num_reqs: int = 0
+    num_waiting_reqs: int = 0
+    num_pages: int = 0
+
+
+@dataclass
+class WatchLoadUpdateReq(BaseReq):
+    loads: list[GetLoadReqOutput] = field(default_factory=list)
+
+
+class BlockReqType(Enum):
+    BLOCK = 1
+    UNBLOCK = 2
+
+
+@dataclass
+class BlockReqInput(BaseReq):
+    type: BlockReqType = field(default_factory=BlockReqType.BLOCK)