import time
import threading
class TokenBucket:
    def __init__(self, capacity: int, rate: float):
        self.capacity = capacity  # 令牌桶容量（最大并发）
        self.rate = rate          # 每秒生成令牌数
        self.tokens = capacity    # 当前令牌数
        self.last_time = time.time()
        self.lock = threading.Lock()
        self.total_requests = 0   # 总请求数
        self.passed_requests = 0  # 通过请求数
        self.rejected_requests = 0  # 拒绝请求数
    def acquire(self, tokens=1) -> bool:
        """获取令牌，成功返回True，失败返回False"""
        with self.lock:
            now = time.time()
            # 计算时间差，生成新令牌
            delta = now - self.last_time
            new_tokens = delta * self.rate
            self.tokens += new_tokens
            self.tokens = min(self.tokens, self.capacity)
            self.last_time = now
            self.total_requests += 1
            if self.tokens >= tokens:
                self.tokens -= tokens
                self.passed_requests += 1
                return True
            else:
                self.rejected_requests += 1
                return False
    def get_stats(self):
        """获取限流统计信息"""
        return {
            "current_tokens": round(self.tokens, 2),
            "total_requests": self.total_requests,
            "passed": self.passed_requests,
            "rejected": self.rejected_requests,
            "pass_rate": f"{self.passed_requests / max(self.total_requests, 1) * 100:.1f}%"
        }
# 大模型队列限流实例：每秒生成2个令牌，桶容量5（模拟GPU显存限制）
llm_rate_limiter = TokenBucket(capacity=5, rate=2)
def handle_llm_request(request_id: str, prompt: str = "", simulate_process: bool = True):
    """处理大模型请求"""
    timestamp = time.strftime("%H:%M:%S", time.localtime())
    # 请求先获取令牌
    if llm_rate_limiter.acquire():
        print(f"[{timestamp}] 请求{request_id}：✓ 获取令牌成功 | 当前令牌: {llm_rate_limiter.tokens:.1f} | 进入推理队列")
        # 模拟推理耗时（可选，用于演示令牌消耗）
        if simulate_process:
            time.sleep(0.05)  # 缩短处理时间，让限流效果更明显
        return True
    else:
        print(f"[{timestamp}] 请求{request_id}：✗ 限流触发 | 当前令牌: {llm_rate_limiter.tokens:.1f} | 排队等待或拒绝")
        return False
def burst_test_concurrent():
    """并发突发测试：多线程同时请求"""
    import concurrent.futures
    print("=" * 60)
    print("【场景1】并发突发流量测试：20个请求同时到达")
    print("=" * 60)
    print(f"配置：令牌桶容量={llm_rate_limiter.capacity}, 生成速率={llm_rate_limiter.rate}/秒")
    print("-" * 60)
    results = {"passed": 0, "rejected": 0}
    def worker(i):
        return handle_llm_request(f"concurrent_{i}", simulate_process=False)
    # 使用线程池并发执行20个请求
    with concurrent.futures.ThreadPoolExecutor(max_workers=20) as executor:
        futures = [executor.submit(worker, i) for i in range(20)]
        for future in concurrent.futures.as_completed(futures):
            if future.result():
                results["passed"] += 1
            else:
                results["rejected"] += 1
    print("-" * 60)
    stats = llm_rate_limiter.get_stats()
    print(f"统计：总请求={stats['total_requests']}, 通过={stats['passed']}, 拒绝={stats['rejected']}, 通过率={stats['pass_rate']}")
    print(f"说明：并发请求中仅前{llm_rate_limiter.capacity}个获取令牌，其余被拒绝")
# ===================== 测试场景 =====================
# 场景1：并发突发测试
burst_test_concurrent()
print("\n" + "=" * 60)
print("【场景2】令牌恢复测试：等待3秒后再次请求")
print("=" * 60)
# 先清空令牌
llm_rate_limiter.tokens = 0
print(f"当前令牌数已清空: {llm_rate_limiter.tokens:.1f}")
time.sleep(3)  # 等待令牌恢复
print(f"等待3秒后，当前令牌数: {llm_rate_limiter.tokens:.1f} (应恢复约{min(3 * llm_rate_limiter.rate, llm_rate_limiter.capacity)}个)")
for i in range(8):
    handle_llm_request(f"recover_{i}")
    time.sleep(0.05)
print("\n" + "=" * 60)
print("【场景3】持续流量测试：请求速率 > 令牌生成速率")
print("=" * 60)
# 重置计数器
llm_rate_limiter.total_requests = 0
llm_rate_limiter.passed_requests = 0
llm_rate_limiter.rejected_requests = 0
llm_rate_limiter.tokens = llm_rate_limiter.capacity  # 重置令牌
print(f"配置：令牌生成速率={llm_rate_limiter.rate}/秒，请求间隔=0.3秒（即约3.3请求/秒）")
print("-" * 60)
# 模拟20个请求，间隔0.3秒（请求速率 > 令牌生成速率）
for i in range(20):
    handle_llm_request(f"steady_{i}")
    time.sleep(0.3)
print("-" * 60)
stats = llm_rate_limiter.get_stats()
print(f"最终统计：总请求={stats['total_requests']}, 通过={stats['passed']}, 拒绝={stats['rejected']}, 通过率={stats['pass_rate']}")
print(f"说明：当请求速率({1/0.3:.1f}/秒) > 令牌生成速率({llm_rate_limiter.rate}/秒)时，部分请求被拒绝")
print("=" * 60)

import queue
import time
import threading
from enum import Enum
class PriorityLevel(Enum):
    """优先级等级"""
    VIP = 0      # 最高优先级：VIP用户/付费用户
    NORMAL = 1   # 普通优先级：普通用户
    FREE = 2     # 最低优先级：免费试用用户
class PriorityLLMQueue:
    def __init__(self, maxsize: int = 100):
        """
        多级优先级队列
        maxsize: 每级队列最大容量
        """
        # 多级队列：0=VIP, 1=普通, 2=免费
        self.queues = {
            PriorityLevel.VIP: queue.Queue(maxsize=maxsize),
            PriorityLevel.NORMAL: queue.Queue(maxsize=maxsize),
            PriorityLevel.FREE: queue.Queue(maxsize=maxsize)
        }
        self.stats = {
            PriorityLevel.VIP: {"enqueued": 0, "dequeued": 0},
            PriorityLevel.NORMAL: {"enqueued": 0, "dequeued": 0},
            PriorityLevel.FREE: {"enqueued": 0, "dequeued": 0}
        }
        self.lock = threading.Lock()
    def enqueue(self, request: dict, priority: PriorityLevel) -> bool:
        """
        请求入队
        priority: PriorityLevel.VIP/NORMAL/FREE
        """
        try:
            self.queues[priority].put_nowait(request)
            with self.lock:
                self.stats[priority]["enqueued"] += 1
            return True
        except queue.Full:
            print(f"[警告] 优先级{priority.name}队列已满，请求{request['id']}被拒绝")
            return False
    def dequeue(self) -> tuple:
        """
        优先调度高优先级请求
        返回: (request, priority) 或 (None, None)
        """
        for priority in [PriorityLevel.VIP, PriorityLevel.NORMAL, PriorityLevel.FREE]:
            try:
                request = self.queues[priority].get_nowait()
                with self.lock:
                    self.stats[priority]["dequeued"] += 1
                return request, priority
            except queue.Empty:
                continue
        return None, None
    def get_queue_sizes(self) -> dict:
        """获取各级队列当前大小"""
        return {
            priority.name: self.queues[priority].qsize()
            for priority in PriorityLevel
        }
    def get_stats(self) -> dict:
        """获取统计信息"""
        return {
            priority.name: {
                "队列长度": self.queues[priority].qsize(),
                "已入队": self.stats[priority]["enqueued"],
                "已处理": self.stats[priority]["dequeued"]
            }
            for priority in PriorityLevel
        }
# ===================== 测试场景 =====================
def test_basic_priority():
    """基础优先级调度测试"""
    print("=" * 60)
    print("【场景1】基础优先级调度测试")
    print("=" * 60)
    print("说明：VIP用户优先于普通用户，普通用户优先于免费用户")
    print("-" * 60)
    llm_queue = PriorityLLMQueue()
    # 模拟混合请求入队（按乱序入队）
    requests = [
        {"id": "free_user_1", "prompt": "你好"},
        {"id": "vip_user_1", "prompt": "紧急问题"},
        {"id": "normal_user_1", "prompt": "普通咨询"},
        {"id": "free_user_2", "prompt": "谢谢"},
        {"id": "vip_user_2", "prompt": "重要业务"},
        {"id": "normal_user_2", "prompt": "一般问题"},
    ]
    priorities = [
        PriorityLevel.FREE,
        PriorityLevel.VIP,
        PriorityLevel.NORMAL,
        PriorityLevel.FREE,
        PriorityLevel.VIP,
        PriorityLevel.NORMAL
    ]
    print("【入队阶段】")
    for req, prio in zip(requests, priorities):
        llm_queue.enqueue(req, prio)
    print(f"\n当前队列状态: {llm_queue.get_queue_sizes()}")
    print("\n【调度执行阶段】")
    processed = []
    while True:
        req, prio = llm_queue.dequeue()
        if not req:
            break
        print(f"[处理] {req['id']} (优先级: {prio.name}) - 提示: {req['prompt']}")
        processed.append((req['id'], prio.name))
        time.sleep(0.2)
    print(f"\n处理顺序: {' -> '.join([p[0] for p in processed])}")
    print(f"优先级顺序: {' -> '.join([p[1] for p in processed])}")
    print("-" * 60)
def test_burst_priority():
    """突发流量下的优先级保障测试"""
    print("\n" + "=" * 60)
    print("【场景2】突发流量下的优先级保障")
    print("=" * 60)
    print("说明：大量免费用户请求涌入时，VIP用户请求仍能被优先处理")
    print("-" * 60)
    llm_queue = PriorityLLMQueue()
    # 模拟突发：10个免费用户 + 2个VIP用户
    print("【突发入队】10个免费用户请求 + 2个VIP用户请求")
    for i in range(10):
        llm_queue.enqueue({"id": f"free_user_{i}", "prompt": f"免费问题{i}"}, PriorityLevel.FREE)
    # VIP用户稍后到达
    time.sleep(0.1)
    llm_queue.enqueue({"id": "vip_user_urgent", "prompt": "紧急业务"}, PriorityLevel.VIP)
    llm_queue.enqueue({"id": "vip_user_important", "prompt": "重要咨询"}, PriorityLevel.VIP)
    print(f"\n当前队列状态: {llm_queue.get_queue_sizes()}")
    print("\n【调度执行】")
    processed = []
    for i in range(12):
        req, prio = llm_queue.dequeue()
        if req:
            print(f"[{i+1:2d}] {req['id']:20s} | 优先级: {prio.name:6s}")
            processed.append(prio.name)
    # 统计
    vip_count = processed.count("VIP")
    free_count = processed.count("FREE")
    print(f"\n统计：VIP请求处理{vip_count}个（占比{vip_count/len(processed)*100:.1f}%）")
    print(f"      免费请求处理{free_count}个（占比{free_count/len(processed)*100:.1f}%）")
    print("说明：虽然免费请求先入队，但VIP请求被优先处理")
    print("-" * 60)
def test_queue_full():
    """队列满时的处理测试"""
    print("\n" + "=" * 60)
    print("【场景3】队列容量限制测试")
    print("=" * 60)
    print("说明：每级队列容量有限，防止无限堆积")
    print("-" * 60)
    # 创建小容量队列
    llm_queue = PriorityLLMQueue(maxsize=3)
    print("【填满VIP队列】")
    for i in range(5):
        success = llm_queue.enqueue({"id": f"vip_user_{i}"}, PriorityLevel.VIP)
        print(f"  VIP请求{i}: {'入队成功' if success else '入队失败(队列满)'}")
    print(f"\n最终队列状态:\n{llm_queue.get_stats()}")
    print("-" * 60)
# ===================== 运行所有测试 =====================
if __name__ == "__main__":
    test_basic_priority()
    test_burst_priority()
    test_queue_full()
    print("\n" + "=" * 60)
    print("测试完成！")
    print("=" * 60)

import queue
import threading
import time
from dataclasses import dataclass
# 定义请求结构体（大模型标准请求格式）
@dataclass
class LLMRequest:
    req_id: str
    prompt: str
    max_tokens: int
    create_time: float = None
    def __post_init__(self):
        self.create_time = time.time()
class LLMLongShortRequestSplitter:
    def __init__(self, short_token_threshold=512):
        """
        长短请求拆分器
        :param short_token_threshold: 短请求最大token数，超过即为长请求
        """
        self.threshold = short_token_threshold
        # 两个独立队列：短请求队列、长请求队列
        self.short_queue = queue.Queue(maxsize=50)
        self.long_queue = queue.Queue(maxsize=30)
        # 启动独立消费者
        self._start_consumers()
    def _estimate_time(self, max_tokens):
        """模拟大模型推理耗时预估：token越多耗时越长"""
        base_time = 0.2
        token_time = max_tokens * 0.015
        return round(base_time + token_time, 2)
    def classify_request(self, req: LLMRequest):
        """请求分类：自动打入长/短队列"""
        estimated_cost = self._estimate_time(req.max_tokens)
        is_short = req.max_tokens <= self.threshold
        if is_short:
            try:
                self.short_queue.put_nowait(req)
                print(f"[短请求] 入队 | ID={req.req_id} | tokens={req.max_tokens} | 预估耗时={estimated_cost}s")
            except queue.Full:
                print(f"[短请求队列溢出] 拒绝请求 {req.req_id}")
        else:
            try:
                self.long_queue.put_nowait(req)
                print(f"[长请求] 入队 | ID={req.req_id} | tokens={req.max_tokens} | 预估耗时={estimated_cost}s")
            except queue.Full:
                print(f"[长请求队列溢出] 拒绝请求 {req.req_id}")
    def _consume_short_request(self):
        """短请求消费者：低延迟、快速处理"""
        while True:
            try:
                req = self.short_queue.get(timeout=1)
                cost = self._estimate_time(req.max_tokens)
                print(f"[短请求执行] ID={req.req_id} | 推理耗时={cost}s")
                time.sleep(cost)
                self.short_queue.task_done()
            except queue.Empty:
                continue
    def _consume_long_request(self):
        """长请求消费者：慢速、不影响短请求"""
        while True:
            try:
                req = self.long_queue.get(timeout=1)
                cost = self._estimate_time(req.max_tokens)
                print(f"[长请求执行] ID={req.req_id} | 推理耗时={cost}s")
                time.sleep(cost)
                self.long_queue.task_done()
            except queue.Empty:
                continue
    def _start_consumers(self):
        """启动消费者线程"""
        # 短请求用2个线程（低延迟、高吞吐）
        threading.Thread(target=self._consume_short_request, daemon=True).start()
        threading.Thread(target=self._consume_short_request, daemon=True).start()
        # 长请求用1个线程（避免占用过多资源）
        threading.Thread(target=self._consume_long_request, daemon=True).start()
# ====================== 测试 ======================
if __name__ == "__main__":
    # 初始化拆分器：<=512 token 为短请求
    splitter = LLMLongShortRequestSplitter(short_token_threshold=512)
    # 模拟一批混合长短请求
    test_requests = [
        LLMRequest("req_001", "你好", max_tokens=64),
        LLMRequest("req_002", "写一篇人工智能综述", max_tokens=2048),
        LLMRequest("req_003", "什么是大模型", max_tokens=128),
        LLMRequest("req_004", "生成1万字产品文档", max_tokens=10000),
        LLMRequest("req_005", "介绍Python", max_tokens=256),
        LLMRequest("req_006", "写一份年度总结报告", max_tokens=3000),
    ]
    print("=== 开始提交请求 ===")
    for req in test_requests:
        splitter.classify_request(req)
        time.sleep(0.2)
    # 保持程序运行
    input("\n按回车退出\n")

import queue
import threading
import time
from dataclasses import dataclass
# 定义大模型请求结构
@dataclass
class LLMRequest:
    req_id: str
    priority: int   # 0最高 1中等 2最低
    prompt: str
    create_time: float = None
    def __post_init__(self):
        self.create_time = time.time()
class LLMBlockQueueOverflowProtect:
    def __init__(self, max_queue_size=30, wait_timeout=5):
        """
        阻塞队列溢出防护初始化
        :param max_queue_size: 队列最大容量
        :param wait_timeout: 请求队列最大等待超时(秒)
        """
        self.max_queue_size = max_queue_size
        self.wait_timeout = wait_timeout
        # 线程安全阻塞队列
        self.block_queue = queue.Queue(maxsize=self.max_queue_size)
        # 启动推理消费线程
        self._start_consumer()
    def get_queue_usage(self):
        """获取队列使用率"""
        current_size = self.block_queue.qsize()
        usage = current_size / self.max_queue_size
        return current_size, round(usage * 100, 2)
    def alert_warning(self):
        """队列告警：使用率80%预警，100%溢出告警"""
        _, usage = self.get_queue_usage()
        if usage >= 100:
            print(f"【严重告警】队列已满！使用率：{usage}%，触发溢出防护")
        elif usage >= 80:
            print(f"【预警】队列即将积压，使用率：{usage}%")
    def try_enqueue(self, req: LLMRequest) -> bool:
        """
        尝试入队 + 溢出防护
        返回True：入队成功  False：入队被拒绝/超时
        """
        self.alert_warning()
        current_size, _ = self.get_queue_usage()
        # 策略1：队列已满，低优先级直接丢弃
        if current_size >= self.max_queue_size:
            if req.priority >= 1:
                print(f"【溢出丢弃】队列已满，低优请求丢弃 ID:{req.req_id} 优先级:{req.priority}")
                return False
            else:
                # 高优请求尝试限时等待入队
                try:
                    self.block_queue.put(req, timeout=self.wait_timeout)
                    print(f"【高优插队】高优先级请求限时等待入队成功 ID:{req.req_id}")
                    return True
                except queue.Full:
                    print(f"【溢出拒绝】高优请求等待超时，入队失败 ID:{req.req_id}")
                    return False
        # 队列未满，正常入队
        try:
            self.block_queue.put(req)
            print(f"【正常入队】ID:{req.req_id} 优先级:{req.priority} 队列当前:{current_size+1}/{self.max_queue_size}")
            return True
        except Exception as e:
            print(f"【入队异常】ID:{req.req_id} 错误:{str(e)}")
            return False
    def _consumer_loop(self):
        """推理消费线程：模拟LLM推理耗时"""
        while True:
            try:
                # 从阻塞队列取任务，空则阻塞
                req = self.block_queue.get(timeout=2)
                # 模拟推理耗时
                infer_cost = 0.5 + (req.priority * 0.3)
                time.sleep(infer_cost)
                print(f"【推理完成】ID:{req.req_id} 耗时:{infer_cost}s")
                self.block_queue.task_done()
            except queue.Empty:
                continue
    def _start_consumer(self):
        # 启动2个推理消费线程
        for _ in range(2):
            t = threading.Thread(target=self._consumer_loop, daemon=True)
            t.start()
# ====================== 测试演示 ======================
if __name__ == "__main__":
    # 初始化：队列最大容量30，请求最大等待超时5秒
    protect_queue = LLMBlockQueueOverflowProtect(max_queue_size=30, wait_timeout=5)
    # 模拟连续涌入40个请求，超过队列容量，触发溢出防护
    for i in range(40):
        # 随机优先级 0/1/2
        prio = i % 3
        req = LLMRequest(
            req_id=f"req_{i:03d}",
            priority=prio,
            prompt=f"测试请求内容{i}"
        )
        protect_queue.try_enqueue(req)
        time.sleep(0.1)
    # 保持程序运行观察消费过程
    input("\n回车退出程序\n")