TaskTracker.py

from enum import Enum
    
class TaskTracker:
    def __init__(self):
        self.running_tasks = {}  # task_id -> task_state
        self.completed_tasks = set()
    
    def start_task(self, task_id, running_time, timestamp, model_id=None):
        """启动任务，可选指定模型ID"""
        assert task_id not in self.completed_tasks, f"{task_id} is completed."
        self.running_tasks[task_id] = {
            "remain_time": running_time,
            "start_time": timestamp,
            "model_id": model_id
        }
    
    def complete_task(self, task_id):
        """完成任务"""
        assert self.running_tasks[task_id]["remain_time"] <= 0
        self.completed_tasks.add(task_id)
        #del self.running_tasks[task_id]

    def update_tasks(self):
        """更新任务状态"""
        completed_this_cycle = []
        
        for task_id, task_state in list(self.running_tasks.items()):
            task_state["remain_time"] -= 1
            if task_state["remain_time"] <= 0:
                completed_this_cycle.append(task_id)
        
        for task_id in completed_this_cycle:
            self.complete_task(task_id)

        return completed_this_cycle