Scheduler

Instead of placing separate queues within each processing module (e.g., LLM core(s) or memory manager), we centralize all the queues within the scheduler module. This aims to isolate the responsibility for request management from the individual modules, allowing each processing module to focus on its execution and simplifies the coordination of tasks across modules.

AIOS Scheduler manages the Request queue for each module, the Request queue includes both Queries (agent-to-kernel request) and Messages (module-to-module request)

The main purpose of the scheduler is to serve different agents queries in a fine-grained granularity (i.e., syscall level granuity) as below.

Scheduling example for syscalls for different agents (use LLM Syscall for illustration simplicity)

Below is the implementation for the basic class of the scheduler. Different schedulers will inherit this class and implement their own abstract methods.

class Scheduler:
    def __init__(
        self,
        llm,
        memory_manager,
        storage_manager,
        tool_manager,
        log_mode,
        get_llm_syscall: LLMRequestQueueGetMessage,
        get_memory_syscall: MemoryRequestQueueGetMessage,
        get_storage_syscall: StorageRequestQueueGetMessage,
        get_tool_syscall: ToolRequestQueueGetMessage,
    ):
        # self.agent_process_queue = Queue()
        self.get_llm_syscall = get_llm_syscall
        self.get_memory_syscall = get_memory_syscall
        self.get_storage_syscall = get_storage_syscall
        self.get_tool_syscall = get_tool_syscall
        self.active = False  # start/stop the scheduler
        self.log_mode = log_mode
        self.logger = self.setup_logger()
        self.syscall_processors = {
            "llm_syscall_processor": Thread(target=self.run_llm_syscall),
            "mem_syscall_processor": Thread(target=self.run_memory_syscall),
            "sto_syscall_processor": Thread(target=self.run_storage_syscall),
            "tool_syscall_processor": Thread(target=self.run_tool_syscall)
        }
        self.llm = llm
        self.memory_manager = memory_manager
        self.storage_manager = storage_manager
        self.tool_manager = tool_manager

    def start(self):
        """start the scheduler"""
        self.active = True
        for name, thread_value in self.syscall_processors.items():
            thread_value.start()

    def stop(self):
        """stop the scheduler"""
        self.active = False
        for name, thread_value in self.syscall_processors.items():
            thread_value.join()

    def setup_logger(self):
        logger = SchedulerLogger("Scheduler", self.log_mode)
        return logger

    def run_llm_syscall(self):
        pass

    def run_memory_syscall(self):
        pass
        
    def run_storage_syscall(self):
        pass
    
    def run_tool_syscall(self):
        pass