# scheduler_core.py
"""调度器核心数据结构（简化版）"""

from enum import Enum
from dataclasses import dataclass, field
from typing import Optional

class SequenceState(Enum):
    WAITING = "waiting"
    RUNNING = "running"
    SWAPPED = "swapped"
    FINISHED = "finished"

@dataclass
class SequenceGroup:
    """一组相关的序列（例如 beam search 中的多个 beam）"""
    request_id: str
    seqs: list  # 序列列表
    state: SequenceState
    priority: float = 0.0  # 优先级
    arrival_time: float = 0.0  # 到达时间
    lora_request: Optional[object] = None
    
    @property
    def num_unfinished_seqs(self) -> int:
        return sum(1 for s in self.seqs if not s.is_finished)

class Scheduler:
    def __init__(self, config):
        # 等待队列：尚未开始处理的请求
        self.waiting: list[SequenceGroup] = []
        
        # 运行队列：正在 GPU 上生成的请求
        self.running: list[SequenceGroup] = []
        
        # 交换队列：被抢占换出到 CPU 的请求
        self.swapped: list[SequenceGroup] = []
        
        # Block 分配器
        self.block_manager = BlockManager(config.num_gpu_blocks)
        
        # 调度配置
        self.max_num_seqs = config.max_num_seqs
        self.max_num_batched_tokens = config.max_num_batched_tokens

def schedule(self) -> SchedulerOutput:
    scheduled_seq_groups = []
    remaining_token_budget = self.max_num_batched_tokens
    remaining_seq_budget = self.max_num_seqs
    
    # 1. 先调度已运行的序列（每个仅需 1 个 token）
    new_running = []
    for seq_group in self.running:
        if remaining_token_budget <= 0 or remaining_seq_budget <= 0:
            break
        if not self.block_manager.can_append(seq_group):
            # 需要抢占
            break
        
        scheduled_seq_groups.append(seq_group)
        remaining_token_budget -= 1  # decode 只需 1 个 token
        remaining_seq_budget -= 1
        new_running.append(seq_group)
    
    # 2. 调度等待中的新序列（prefill 需要处理全部 prompt tokens）
    new_waiting = []
    for seq_group in self.waiting:
        num_prompt_tokens = seq_group.get_prompt_len()
        
        if num_prompt_tokens > remaining_token_budget:
            continue  # token 预算不够
        if remaining_seq_budget <= 0:
            break  # 序列数已满
        
        # 检查是否可以分配 block
        if self.block_manager.can_allocate(num_prompt_tokens):
            self.block_manager.allocate(seq_group)
            scheduled_seq_groups.append(seq_group)
            remaining_token_budget -= num_prompt_tokens
            remaining_seq_budget -= 1
            new_running.append(seq_group)
        else:
            new_waiting.append(seq_group)  # 资源不足，继续等待
    
    self.running = new_running
    self.waiting = new_waiting
    
    return SchedulerOutput(scheduled_seq_groups)

def preempt(seq_group, blocks_to_free):
    """选择抢占策略"""
    
    if can_swap_to_cpu():
        # 策略一：Swap（推荐）
        # 将 KV Cache 复制到 CPU 内存
        # 恢复时 swap 回 GPU
        cpu_blocks = swap_to_cpu(seq_group, gpu_blocks)
        return PreemptionAction.SWAP, cpu_blocks
    else:
        # 策略二：Recompute
        # 直接丢弃 KV Cache
        # 恢复时重新执行 prefill
        release_blocks(seq_group)
        return PreemptionAction.RECOMPUTE, None

def select_preemption_victims(num_blocks_needed: int) -> list[SequenceGroup]:
    """选择被抢占的序列（LRU 策略）"""
    
    # 按到达时间倒序（最新加入的先被抢占）
    candidates = sorted(
        self.running,
        key=lambda sg: sg.arrival_time,
        reverse=True,  # 最新到达的优先被抢占
    )
    
    victims = []
    freed_blocks = 0
    
    for seq_group in candidates:
        if freed_blocks >= num_blocks_needed:
            break
        victims.append(seq_group)
        freed_blocks += self.block_manager.get_num_blocks(seq_group)
    
    return victims

# priority_scheduler.py
"""带优先级的请求调度"""

from enum import IntEnum

class Priority(IntEnum):
    CRITICAL = 0    # 最高优先级
    HIGH = 1
    NORMAL = 2
    LOW = 3
    BATCH = 4       # 最低优先级（批量任务）

class PriorityScheduler(Scheduler):
    def __init__(self, config):
        super().__init__(config)
        self.priority_queue: dict[Priority, list[SequenceGroup]] = {
            p: [] for p in Priority
        }
    
    def add_request(self, seq_group: SequenceGroup, priority: Priority):
        """添加带优先级的请求"""
        self.priority_queue[priority].append(seq_group)
    
    def schedule(self) -> SchedulerOutput:
        """按优先级顺序调度"""
        # 先运行已运行的序列
        scheduled = self._schedule_running()
        
        # 按优先级从高到低调度新请求
        for priority in Priority:
            if not self.priority_queue[priority]:
                continue
            
            while self.priority_queue[priority]:
                seq_group = self.priority_queue[priority][0]
                
                if not self._can_schedule(seq_group):
                    break
                
                self._schedule_prefill(seq_group)
                self.priority_queue[priority].pop(0)
                scheduled.append(seq_group)
        
        return SchedulerOutput(scheduled)

# 通过 API 传递优先级
response = client.chat.completions.create(
    model="qwen-7b",
    messages=[{"role": "user", "content": "紧急问题"}],
    extra_body={"priority": "high"},
)

# 启用 Chunked Prefill
vllm serve model \
    --enable-chunked-prefill \
    --max-num-batched-tokens 512  # chunk 大小

def schedule_with_chunked_prefill(self) -> SchedulerOutput:
    """带 Chunked Prefill 的调度"""
    
    scheduled = []
    token_budget = self.max_num_batched_tokens
    
    # 预留 decode 需要的 token
    decode_budget = len(self.running)  # 每个运行序列 1 个 token
    prefill_budget = token_budget - decode_budget
    
    # 1. 处理 decode
    for seq_group in self.running:
        scheduled.append(ScheduledSeqGroup(seq_group, is_prefill=False))
    
    # 2. 处理 prefill（分块）
    for seq_group in self.waiting:
        remaining = seq_group.get_num_unprefilled_tokens()
        chunk_size = min(remaining, prefill_budget)
        
        if chunk_size <= 0:
            break
        
        scheduled.append(ScheduledSeqGroup(
            seq_group,
            is_prefill=True,
            chunk_size=chunk_size,
        ))
        prefill_budget -= chunk_size
        
        if seq_group.is_fully_prefilled():
            self.running.append(seq_group)
    
    return SchedulerOutput(scheduled)

# 高吞吐场景（离线批量处理）
llm = LLM(
    model="model",
    max_num_seqs=512,           # 大并发
    max_num_batched_tokens=16384,  # 大 token 预算
    max_model_len=2048,          # 限制序列长度
    gpu_memory_utilization=0.95,
)

# 低延迟场景（实时 API 服务）
llm = LLM(
    model="model",
    max_num_seqs=64,             # 适中并发
    max_num_batched_tokens=4096,
    enable_chunked_prefill=True,
    gpu_memory_utilization=0.9,
)

# 长上下文场景
llm = LLM(
    model="model",
    max_num_seqs=32,             # 较少并发
    max_num_batched_tokens=8192,
    max_model_len=32768,          # 长上下文
    gpu_memory_utilization=0.92,
)

# queue_management.py
"""请求队列管理"""

import asyncio
import time

class RequestQueue:
    """请求队列管理器"""
    
    def __init__(self, max_queue_size: int = 1000, timeout: float = 300):
        self.queue = asyncio.Queue(maxsize=max_queue_size)
        self.timeout = timeout
    
    async def enqueue(self, request) -> bool:
        """添加请求到队列"""
        try:
            self.queue.put_nowait(request)
            return True
        except asyncio.QueueFull:
            return False
    
    async def dequeue(self):
        """获取下一个请求（带超时）"""
        try:
            return await asyncio.wait_for(
                self.queue.get(), 
                timeout=self.timeout
            )
        except asyncio.TimeoutError:
            return None

# 调度统计信息
scheduler_stats = {
    "num_running_seqs": 42,       # 正在运行的序列数
    "num_waiting_seqs": 156,      # 等待中的序列数
    "num_swapped_seqs": 3,        # 被 swap 的序列数
    "gpu_cache_usage": 0.85,      # GPU 缓存使用率
    "cpu_cache_usage": 0.12,      # CPU 缓存使用率
    "num_preemptions": 2,         # 本批次抢占次数
}