simulator.py

from schedulers.scheduler_v3_1_main import Scheduler
from TaskTracker import TaskTracker
from ResourceManager import ResourceManager
from generator import load_workflows_from_directory

from source.logger import get_logger

logger = get_logger(__name__)

class Simulator:
    def __init__(self, scheduler):
        self.resource_manager = ResourceManager()
        self.task_tracker = TaskTracker()
        self.scheduler = scheduler
        self.current_timestamp = 0
        # self.workflows = {wf.wf_id: wf for wf in load_workflows_from_directory("multi_workflows_2048")}
        # self.workflows = {wf.wf_id: wf for wf in load_workflows_from_directory("multi_workflows_1000")}
        # self.workflows = {wf.wf_id: wf for wf in load_workflows_from_directory("multi_workflows_100")}
        self.workflows = {wf.wf_id: wf for wf in load_workflows_from_directory("multi_workflows")}
        self.wait_wf_ids = list(self.workflows.keys())
        self.task_info = {}  # task_id -> task_info (from workflows)
        
        # 初始化任务信息
        for wf_id, wf in self.workflows.items():
            self.task_info.update(wf.task_info)
        self.completed_tasks = []
        self.remain_tasks = set(self.task_info.keys())
        
        # 性能统计数据
        self.stats = {
            'workflow_metrics': {},  # wf_id -> {arrival_time, start_time, end_time}
            'task_metrics': {},      # task_id -> {ready_time, start_time, end_time}
            'resource_usage': {      # 每个时间步的资源使用情况
                'cpu_usage': [],
                'npu_usage': []
            },
            'model_operations': {    # 模型操作统计
                'deploy_count': 0,
                'unload_count': 0
            }
        }
        # 初始化workflow统计
        for wf_id, wf in self.workflows.items():
            self.stats['workflow_metrics'][wf_id] = {
                'arrival_time': wf.arrival_time,
                'start_time': None,
                'end_time': None
            }
    
    def _check_dependencies(self, task_id):
        """检查任务的所有前置依赖是否已完成"""
        for wf_id, wf in self.workflows.items():
            task_node = wf.nodes.get(task_id)
            if task_node:
                for parent_node in task_node.parents:
                    if parent_node.task_id not in self.task_tracker.completed_tasks:
                        return False, parent_node.task_id
                return True, None
        return True, None
    
    def check_all_workflows_completed(self):
        """检查是否所有工作流的所有任务都已完成"""
        if len(self.remain_tasks) == 0:
            logger.info(f"All workflows completed at timestamp {self.current_timestamp}.")
            return True
        return False

    def run(self):
        # 1.get new wf
        new_wf_ids = []
        for wf_id in self.wait_wf_ids:
            if self.workflows[wf_id].arrival_time <= self.current_timestamp:
                new_wf_ids.append(wf_id)
        for wf_id in new_wf_ids:
            self.wait_wf_ids.remove(wf_id)
        new_wfs = [self.workflows[wf_id] for wf_id in new_wf_ids]
    
        # 2. scheduler
        actions = self.scheduler.schedule(self.current_timestamp, self.resource_manager, new_wfs, self.completed_tasks)
        
        # 3.执行动作
        for action in actions:
            self.execute_action(action)
        
        # 4.更新数据
        self.completed_tasks = self.task_tracker.update_tasks()
        self.resource_manager.update_deployments(self.current_timestamp)
        
        # 5.处理完成的任务
        for task_id in self.completed_tasks:
            task = self.task_info.get(task_id)
            if not task:
                logger.error(f"[Simulator] Error: Completed task {task_id} not found in task_info.")
                continue
            self.remain_tasks.remove(task_id)
            
            # 记录任务完成时间
            if task_id in self.stats['task_metrics']:
                self.stats['task_metrics'][task_id]['end_time'] = self.current_timestamp
            
            # 检查是否是workflow的最后一个任务
            for wf_id, wf in self.workflows.items():
                if wf.final_node and wf.final_node.task_id == task_id:
                    self.stats['workflow_metrics'][wf_id]['end_time'] = self.current_timestamp
            
            # 释放模型资源
            if task.task_type in ["prefill", "decoding"]:
                task_state = self.task_tracker.running_tasks.get(task_id, {})
                model_id = task_state.get("model_id")
                if model_id:
                    self.resource_manager.complete_task(task_id, model_id)
            
            # 释放CPU资源（仅tool call任务）
            if task.task_type == "tool_call":
                self.resource_manager.release_cpu(task.cpus)
            del self.task_tracker.running_tasks[task_id]
        
        # 6.记录资源使用情况
        total_cpu, avail_cpu = self.resource_manager.get_cpu_info()
        used_cpu = total_cpu - avail_cpu
        self.stats['resource_usage']['cpu_usage'].append(used_cpu)
        
        npus = self.resource_manager.get_npu_info()
        used_npus = sum(1 for npu in npus.values() if not npu['available'])
        self.stats['resource_usage']['npu_usage'].append(used_npus)
        
        # 7.时间推进
        self.current_timestamp += 1
    
    def execute_action(self, action):
        """执行单个调度动作"""
        action_type = action["type"]
        
        if action_type == "deploy_model":
            self.execute_deploy_model(action)
        elif action_type == "unload_model":
            self.execute_unload_model(action)
        elif action_type == "assign_prefill":
            self.execute_assign_prefill(action)
        elif action_type == "assign_decoding":
            self.execute_assign_decoding(action)
        elif action_type == "assign_tool_call":
            self.execute_assign_tool_call(action)
    
    def execute_deploy_model(self, action):
        """执行部署模型动作"""
        model_name = action["model_name"]
        npu_ids = action["npu_ids"]
        model_id = self.resource_manager.deploy_model(
            self.current_timestamp, model_name, npu_ids
        )
        if model_id:
            self.stats['model_operations']['deploy_count'] += 1
            print(f"[Simulator] Deployed model {model_name} on NPUs {npu_ids} as {model_id}")
        else:
          print(f"[Simulator] Deployed model {model_name} on NPUs {npu_ids} error")
    
    def execute_unload_model(self, action):
        """执行卸载模型动作"""
        model_id = action["model_id"]
        success = self.resource_manager.unload_model(model_id)
        if success:
            self.stats['model_operations']['unload_count'] += 1
            print(f"[Simulator] Unloaded model {model_id}")
    
    def execute_assign_prefill(self, action):
        """执行分配prefill任务动作"""
        task_id = action["task_id"]
        model_id = action["model_id"]
        
        # 获取任务信息
        task = self.task_info.get(task_id)
        if not task:
            return
        # 检查任务依赖是否满足
        deps_satisfied, blocking_task = self._check_dependencies(task_id)
        if not deps_satisfied:
            logger.info(f"[Simulator] Error: Cannot assign prefill task {task_id}, dependency not satisfied. Waiting for task {blocking_task} to complete.")
            return
        
        # 检查模型是否可用
        if not self.resource_manager.can_assign_prefill(model_id):
            print(f"[Simulator] Cannot assign prefill task {task_id} to model {model_id}")
            return
        
        # 获取模型信息
        model = self.resource_manager.get_model_info().get(model_id)
        if not model:
            return
        
        # 获取模型配置
        model_config = self.resource_manager.get_model_config().get(model["model_name"])
        if not model_config:
            return
        
        # 获取第一个NPU的缩放因子
        npu_id = model["npu_ids"][0]
        npu = self.resource_manager.get_npu_info().get(npu_id)
        if not npu:
            return
        
        # 计算执行时间
        scaling_factor = npu["prefill_scaling_factor"]
        tokens_per_sec = model_config["prefill_tokens_per_second"] * scaling_factor
        tokens = task.actual_tokens if task.actual_tokens > 0 else task.predicted_tokens
        exec_time = max(1, int(tokens / tokens_per_sec))  # 至少1秒
        
        # 分配任务到模型
        if not self.resource_manager.assign_prefill(model_id, task_id):
            print(f"[Simulator] Failed to assign prefill task {task_id} to model {model_id}")
            return
        
        # 记录任务开始时间和workflow开始时间
        self.stats['task_metrics'][task_id] = {'start_time': self.current_timestamp, 'end_time': None}
        for wf_id, wf in self.workflows.items():
            if task_id in wf.nodes and self.stats['workflow_metrics'][wf_id]['start_time'] is None:
                self.stats['workflow_metrics'][wf_id]['start_time'] = self.current_timestamp
        
        # 启动任务
        self.task_tracker.start_task(task_id, exec_time, self.current_timestamp, model_id)
        print(f"[Simulator] Started prefill task {task_id} on model {model_id}, exec_time: {exec_time}s")
    
    def execute_assign_decoding(self, action):
        """执行分配decoding任务动作"""
        task_id = action["task_id"]
        model_id = action["model_id"]
        
        # 获取任务信息
        task = self.task_info.get(task_id)
        if not task:
            return
        
        # 检查模型是否可用
        if not self.resource_manager.can_assign_decoding(model_id):
            print(f"[Simulator] Cannot assign decoding task {task_id} to model {model_id}")
            return
        
        # 获取模型信息
        model = self.resource_manager.get_model_info().get(model_id)
        if not model:
            return
        # 检查任务依赖是否满足
        deps_satisfied, blocking_task = self._check_dependencies(task_id)
        if not deps_satisfied:
          logger.error(f"[Simulator] Error: Cannot assign decoding task {task_id}, dependency not satisfied. Waiting for task {blocking_task} to complete.")
          return
        # 获取模型配置
        model_config = self.resource_manager.get_model_config().get(model["model_name"])
        if not model_config:
            return
        
        # 获取第一个NPU的缩放因子
        npu_id = model["npu_ids"][0]
        npu = self.resource_manager.get_npu_info().get(npu_id)
        if not npu:
            return
        
        # 计算执行时间
        scaling_factor = npu["decoding_scaling_factor"]
        tokens_per_sec = model_config["decoding_tokens_per_second"] * scaling_factor
        tokens = task.actual_tokens if task.actual_tokens > 0 else task.predicted_tokens
        exec_time = max(1, int(tokens / tokens_per_sec))  # 至少1秒
        
        # 分配任务到模型
        if not self.resource_manager.assign_decoding(model_id, task_id):
            print(f"[Simulator] Failed to assign decoding task {task_id} to model {model_id}")
            return
        
        # 记录任务开始时间和workflow开始时间
        self.stats['task_metrics'][task_id] = {'start_time': self.current_timestamp, 'end_time': None}
        for wf_id, wf in self.workflows.items():
            if task_id in wf.nodes and self.stats['workflow_metrics'][wf_id]['start_time'] is None:
                self.stats['workflow_metrics'][wf_id]['start_time'] = self.current_timestamp
        
        # 启动任务
        self.task_tracker.start_task(task_id, exec_time, self.current_timestamp, model_id)
        print(f"[Simulator] Started decoding task {task_id} on model {model_id}, exec_time: {exec_time}s")
    
    def execute_assign_tool_call(self, action):
        """执行分配tool call任务动作"""
        task_id = action["task_id"]
        cpus = action["cpus"]
        
        # 获取任务信息
        task = self.task_info.get(task_id)
        if not task:
            return
        # 检查任务依赖是否满足
        deps_satisfied, blocking_task = self._check_dependencies(task_id)
        if not deps_satisfied:
          logger.error(f"[Simulator] Error: Cannot assign tool_call task {task_id}, dependency not satisfied. Waiting for task {blocking_task} to complete.")
          return
        # 分配CPU资源
        if not self.resource_manager.allocate_cpu(cpus):
            print(f"[Simulator] Failed to allocate {cpus} CPUs for tool call task {task_id}")
            return
        
        # 计算执行时间
        exec_time = max(1, int(task.actual_time if task.actual_time > 0 else task.predicted_time))
        
        # 记录任务开始时间和workflow开始时间
        self.stats['task_metrics'][task_id] = {'start_time': self.current_timestamp, 'end_time': None}
        for wf_id, wf in self.workflows.items():
            if task_id in wf.nodes and self.stats['workflow_metrics'][wf_id]['start_time'] is None:
                self.stats['workflow_metrics'][wf_id]['start_time'] = self.current_timestamp
        
        # 启动任务
        self.task_tracker.start_task(task_id, exec_time, self.current_timestamp)
        print(f"[Simulator] Started tool call task {task_id} with {cpus} CPUs, exec_time: {exec_time}s")
    
    def print_statistics(self):
        """输出性能统计报告"""
        import json
        from datetime import datetime
        
        print("\n" + "="*80)
        print("性能统计报告")
        print("="*80)
        
        # 1. Workflow级别指标
        print("\n【Workflow性能指标】")
        print("-" * 80)
        
        turnaround_times = []
        response_times = []
        waiting_times = []
        
        for wf_id, metrics in sorted(self.stats['workflow_metrics'].items()):
            arrival = metrics['arrival_time']
            start = metrics['start_time']
            end = metrics['end_time']
            
            if end is not None:
                turnaround = end - arrival
                turnaround_times.append(turnaround)
                
                if start is not None:
                    response = end - start
                    waiting = start - arrival
                    response_times.append(response)
                    waiting_times.append(waiting)
                    
                    print(f"  {wf_id}:")
                    print(f"    到达时间: {arrival}, 开始时间: {start}, 结束时间: {end}")
                    print(f"    周转时间: {turnaround}, 响应时间: {response}, 等待时间: {waiting}")
            else:
                print(f"  {wf_id}: 未完成")
        
        print(f"\n  汇总:")
        if turnaround_times:
            print(f"    平均周转时间: {sum(turnaround_times)/len(turnaround_times):.2f}")
            print(f"    最大周转时间: {max(turnaround_times)}")
            print(f"    最小周转时间: {min(turnaround_times)}")
        if response_times:
            print(f"    平均响应时间: {sum(response_times)/len(response_times):.2f}")
        if waiting_times:
            print(f"    平均等待时间: {sum(waiting_times)/len(waiting_times):.2f}")
        
        # 2. 资源利用率指标
        print("\n【资源利用率】")
        print("-" * 80)
        
        cpu_usage = self.stats['resource_usage']['cpu_usage']
        npu_usage = self.stats['resource_usage']['npu_usage']
        
        if cpu_usage:
            total_cpu = self.resource_manager.total_cpu_cores
            avg_cpu_usage = sum(cpu_usage) / len(cpu_usage)
            cpu_utilization = (avg_cpu_usage / total_cpu * 100) if total_cpu > 0 else 0
            print(f"  CPU平均使用: {avg_cpu_usage:.2f} / {total_cpu} 核")
            print(f"  CPU利用率: {cpu_utilization:.2f}%")
        
        if npu_usage:
            total_npus = len(self.resource_manager.get_npu_info())
            avg_npu_usage = sum(npu_usage) / len(npu_usage)
            npu_utilization = (avg_npu_usage / total_npus * 100) if total_npus > 0 else 0
            print(f"  NPU平均使用: {avg_npu_usage:.2f} / {total_npus} 个")
            print(f"  NPU利用率: {npu_utilization:.2f}%")
        
        # 3. 模型操作统计
        print("\n【模型操作统计】")
        print("-" * 80)
        print(f"  模型部署次数: {self.stats['model_operations']['deploy_count']}")
        print(f"  模型卸载次数: {self.stats['model_operations']['unload_count']}")
        
        # 4. 任务执行指标
        print("\n【任务执行指标】")
        print("-" * 80)
        
        total_tasks = len(self.task_info)
        completed_count = len(self.task_tracker.completed_tasks)
        print(f"  总任务数: {total_tasks}")
        print(f"  已完成任务数: {completed_count}")
        print(f"  完成率: {(completed_count/total_tasks*100):.2f}%")
        
        # 按任务类型统计
        task_types = {'prefill': [], 'decoding': [], 'tool_call': []}
        for task_id, metrics in self.stats['task_metrics'].items():
            if metrics['start_time'] is not None and metrics['end_time'] is not None:
                task = self.task_info.get(task_id)
                if task:
                    exec_time = metrics['end_time'] - metrics['start_time']
                    task_types[task.task_type].append(exec_time)
        
        print(f"\n  按任务类型统计:")
        for task_type, exec_times in task_types.items():
            if exec_times:
                print(f"    {task_type}: 完成{len(exec_times)}个, 平均执行时间: {sum(exec_times)/len(exec_times):.2f}")
        
        # 5. 系统吞吐量指标
        print("\n【系统吞吐量】")
        print("-" * 80)
        
        total_time = self.current_timestamp
        completed_wfs = sum(1 for m in self.stats['workflow_metrics'].values() if m['end_time'] is not None)
        
        print(f"  总执行时间: {total_time} 时间步")
        print(f"  完成的workflow数: {completed_wfs}")
        if total_time > 0:
            throughput = completed_wfs / total_time
            print(f"  平均吞吐量: {throughput:.4f} workflows/时间步")
        
        print("\n" + "="*80)
        
        # 保存详细数据到JSON文件（带时间戳）
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        stats_output = {
            'summary': {
                'total_time': total_time,
                'completed_workflows': completed_wfs,
                'total_tasks': total_tasks,
                'completed_tasks': completed_count,
                'avg_turnaround_time': sum(turnaround_times)/len(turnaround_times) if turnaround_times else 0,
                'avg_response_time': sum(response_times)/len(response_times) if response_times else 0,
                'cpu_utilization': cpu_utilization if cpu_usage else 0,
                'npu_utilization': npu_utilization if npu_usage else 0,
                'model_deploy_count': self.stats['model_operations']['deploy_count'],
                'model_unload_count': self.stats['model_operations']['unload_count']
            },
            'workflow_metrics': self.stats['workflow_metrics'],
            'task_metrics': self.stats['task_metrics']
        }
        
        output_file = f'log/performance_stats_{timestamp}.json'
        with open(output_file, 'w') as f:
            json.dump(stats_output, f, indent=2)
        print(f"详细统计数据已保存到 {output_file}")

if __name__ == "__main__":
    scheduler = Scheduler()
    simulator = Simulator(scheduler)
    for i in range(50000):
    # for i in range(1000):
        if simulator.check_all_workflows_completed():
            logger.info(f"All workflows are completed at timestamp {simulator.current_timestamp}.")
            break
        simulator.run()
    simulator.print_statistics()
    assert simulator.check_all_workflows_completed(), "Not all workflows are completed"