# scheduler.py:1658 - 优先级设置或验证
def _set_or_validate_priority(self, req: Req) -> bool:
    if self.enable_priority_scheduling and req.priority is None:
        # 设置默认优先级
        if self.schedule_low_priority_values_first:
            req.priority = sys.maxsize  # 默认最低优先级
        else:
            req.priority = -sys.maxsize - 1  # 默认最高优先级
    return True

# schedule_policy.py:276 - 优先级排序算法
@staticmethod
def _sort_by_priority_and_fcfs(waiting_queue: List[Req], priority_sign: int):
    waiting_queue.sort(
        key=lambda x: (
            x.priority * priority_sign,      # 主排序：优先级
            x.time_stats.wait_queue_entry_time,  # 次排序：到达时间（FCFS）
        )
    )

# scheduler.py:1668 - 抢占调度
def preempt_to_schedule(self, req: Req, server_args: ServerArgs) -> bool:
    """抢占正在运行的低优先级请求"""
    priority_diff = (req.priority - running_req.priority) * (-priority_sign)

    if priority_diff > self.priority_scheduling_preemption_threshold:
        # 执行抢占：将低优先级请求移回等待队列
        self.running_batch.remove_request(running_req)
        self.waiting_queue.append(running_req)
        return True
    return False

def _add_request_to_queue(self, req: Req):
    # 1. 验证或设置优先级
    if not self._set_or_validate_priority(req):
        return False

    # 2. 检查队列限制
    if self._abort_on_queued_limit(req):
        return False

    # 3. 预取 KV Cache（可选）
    self._prefetch_kvcache(req)

    # 4. 加入等待队列
    self.waiting_queue.append(req)

    # 5. 记录入队时间
    req.time_stats.wait_queue_entry_time = time.perf_counter()

    return True

def calc_priority(self, waiting_queue: List[Req]) -> bool:
    # 1. 确定活动策略（处理大队列降级）
    active_policy = self._determine_active_policy(waiting_queue)

    # 2. 应用调度策略
    if active_policy == CacheAgnosticPolicy.FCFS:
        if self.enable_priority_scheduling:
            # 按优先级 + FCFS 排序
            SchedulePolicy._sort_by_priority_and_fcfs(
                waiting_queue, self.priority_sign
            )
        # 否则保持 FCFS（已按入队时间排序）

    elif active_policy == CacheAwarePolicy.LPM:
        # 计算每个请求的前缀匹配长度
        for req in waiting_queue:
            req.prefix_len = self._calc_prefix_match_len(req)
        # 按前缀长度降序排序
        waiting_queue.sort(key=lambda x: -x.prefix_len)

    elif active_policy == CacheAwarePolicy.DFS_WEIGHT:
        # 按深度优先搜索权重排序
        SchedulePolicy._sort_by_dfs_weight(waiting_queue, self.tree_cache)

    elif active_policy == CacheAgnosticPolicy.LOF:
        # 按 max_new_tokens 降序排序
        waiting_queue.sort(key=lambda x: -x.max_new_tokens)

    elif active_policy == CacheAgnosticPolicy.RANDOM:
        # 随机打乱
        random.shuffle(waiting_queue)

    return True

def _determine_active_policy(self, waiting_queue: List[Req]) -> Policy:
    if self.policy == CacheAwarePolicy.LPM and len(waiting_queue) > 128:
        # LPM 在大队列时计算开销过大，降级为 FCFS
        return CacheAgnosticPolicy.FCFS
    return self.policy

@staticmethod
def _sort_by_dfs_weight(waiting_queue: List[Req], tree_cache: BasePrefixCache) -> None:
    """基于深度优先搜索权重排序等待队列"""
    last_node_to_reqs = defaultdict(list)
    for req in waiting_queue:
        last_node_to_reqs[req.last_node].append(req)

node_to_weight = defaultdict(int)
for node in last_node_to_reqs:
    node_to_weight[node] = len(last_node_to_reqs[node])
SchedulePolicy._calc_weight(tree_cache.root_node, node_to_weight)

@staticmethod
def _calc_weight(cur_node: TreeNode, node_to_weight: Dict[TreeNode, int]) -> None:
    """递归计算每个节点的权重（自身请求数 + 所有子节点的请求数）"""
    for child in cur_node.children.values():
        SchedulePolicy._calc_weight(child, node_to_weight)
        node_to_weight[cur_node] += node_to_weight[child]

@staticmethod
def _get_dfs_priority(
    cur_node: TreeNode,
    node_to_priority: Dict[TreeNode, int],
    last_node_to_reqs: Dict[TreeNode, List[Req]],
    q: List,
) -> None:
    """按权重深度优先遍历，优先处理权重大的分支"""
    childs = [child for child in cur_node.children.values()]
    childs.sort(key=lambda x: -node_to_priority[x])  # 子节点按权重降序

    for child in childs:
        SchedulePolicy._get_dfs_priority(child, node_to_priority, last_node_to_reqs, q)

    q.extend(last_node_to_reqs[cur_node])  # 叶子节点的请求加入队列

def get_next_batch_to_run(self) -> Optional[ScheduleBatch]:
    # 主调度循环

    # 1. 合并已完成的 prefill batch 到 running batch
    if self.prefill_batch is not None:
        self._merge_prefill_to_running()

    # 2. 如果没有运行中的 batch 或需要新 prefill
    if self.running_batch is None or self.running_batch.can_add_new_req():
        # 获取新的 prefill batch
        prefill_batch = self.get_new_batch_prefill()

        if prefill_batch is not None:
            self.prefill_batch = prefill_batch
            return self.prefill_batch

    # 3. 否则运行 decode batch
    if self.running_batch is not None:
        return self.running_batch

    return None

class PrefillAdder:
    def add_one_req(self, req: Req, has_chunked_req: bool,
                    truncation_align_size: Optional[int]) -> AddReqResult:
        """尝试添加一个请求到 prefill batch"""

        # 1. 检查内存预算
        result = self.budget_state()
        if result != AddReqResult.CONTINUE:
            return result

        # 2. 计算所需 token
        prefix_len = self.prefix_cache.match_prefix_len(req)
        input_tokens = req.req_input_indices.length()

        # 3. 检查是否需要分块 prefill
        if self.need_chunked_prefill(req, prefix_len, input_tokens):
            return self.handle_chunked_prefill(req, prefix_len, input_tokens)

        # 4. 检查 batch 大小限制
        if not self.can_add_more_requests():
            return AddReqResult.OTHER

        # 5. 添加到可运行列表
        self.can_run_list.append(req)

        # 6. 更新内存预算
        self._update_prefill_budget(prefix_len, input_tokens, req.max_new_tokens)

        return AddReqResult.CONTINUE

def budget_state(self) -> AddReqResult:
    if self.rem_total_tokens <= 0 or self.cur_rem_tokens <= 0:
        return AddReqResult.NO_TOKEN  # 总 token 预算不足

    if self.rem_input_tokens <= 0:
        return AddReqResult.OTHER  # 输入 token 预算不足

    return AddReqResult.CONTINUE  # 可以继续添加

self.running_batch: Optional[ScheduleBatch] = None  # 当前运行的 decode batch
self.last_batch: Optional[ScheduleBatch] = None     # 上一个批次

def preempt_to_schedule(self, req: Req, server_args: ServerArgs) -> bool:
    """
    抢占正在运行的低优先级请求以服务高优先级新请求

    返回: True 如果成功抢占，False 否则
    """
    if not self.enable_priority_scheduling:
        return False

    if self.running_batch is None or len(self.running_batch.reqs) == 0:
        return False

    # 计算优先级差异
    priority_sign = -1 if self.schedule_low_priority_values_first else 1
    for running_req in self.running_batch.reqs:
        priority_diff = (req.priority - running_req.priority) * (-priority_sign)

        # 检查是否超过抢占阈值
        if priority_diff > self.priority_scheduling_preemption_threshold:
            # 执行抢占
            self.running_batch.remove_request(running_req)
            self.waiting_queue.append(running_req)
            running_req.time_stats.wait_queue_entry_time = time.perf_counter()
            return True

    return False

from openai import OpenAI

client = OpenAI(
    base_url="http://localhost:8000/v1",
    api_key="EMPTY"
)

response = client.chat.completions.create(
    model="meta-llama/Llama-3.1-8B-Instruct",
    messages=[{"role": "user", "content": "Hello!"}],
    extra_body={
        "priority": 100  # 高优先级
    }
)

print(response.choices[0].message.content)

import sglang as sgl

@sgl.function
def hello(s):
    s += sgl.user("Hello!")
    s += sgl.assistant(sgl.gen("response"))

result = hello.run(
    model_path="meta-llama/Llama-3.1-8B-Instruct",
    priority=100  # 高优先级
)