Ray 源码解析（一）：任务的状态转移和组织形式

enum class TaskState {
 // The task may be placed on a node.
 PLACEABLE,
 // The task has been placed on a node and is waiting for some object
 // dependencies to become local.
 WAITING,
 // The task has been placed on a node, all dependencies are satisfied, and is
 // waiting for resources to run.
 READY,
 // The task is running on a worker. The task may also be blocked in a ray.get
 // or ray.wait call, in which case it also has state BLOCKED.
 RUNNING,
 // The task has resources that cannot be satisfied by any node, as far as we
 // know.
 INFEASIBLE,
 // The task is an actor method and is waiting to learn where the actor was
 // created.
 WAITING_FOR_ACTOR_CREATION,
 // Swap queue for tasks that are in between states. This can happen when a
 // task is removed from one queue, and an async callback is responsible for
 // re-queuing the task. For example, a READY task that has just been assigned
 // to a worker will get moved to the SWAP queue while waiting for a response
 // from the worker. If the worker accepts the task, the task will be added to
 // the RUNNING queue, else it will be returned to READY.
 SWAP,
 // The number of task queues. All states that precede this enum must have an
 // associated TaskQueue in SchedulingQueue. All states that succeed
 // this enum do not have an associated TaskQueue, since the tasks
 // in those states may not have any associated task data.
 kNumTaskQueues,
 // The task is running but blocked in a ray.get or ray.wait call. Tasks that
 // were explicitly assigned by us may be both BLOCKED and RUNNING, while
 // tasks that were created out-of-band (e.g., the application created
 // multiple threads) are only BLOCKED.
 BLOCKED,
 // The task is a driver task.
 DRIVER,
};

class TaskQueue {
public:
 virtual ~TaskQueue() {}
 // 任务的增删改查操作
 virtual bool AppendTask(const TaskID &task_id, const Task &task);
 virtual bool RemoveTask(const TaskID &task_id,
                         std::vector<Task> *removed_tasks = nullptr);
 bool HasTask(const TaskID &task_id) const;
 const std::list<Task> &GetTasks() const;
 const Task &GetTask(const TaskID &task_id) const;
 // 获取队列中所需资源总和
 const ResourceSet &GetCurrentResourceLoad() const;
protected:
 // 链表+哈希组织，可以快速查找O(1)和线性遍历O(n)
 std::list<Task> task_list_;
 std::unordered_map<TaskID, std::list<Task>::iterator> task_map_;
 // 所有任务所需资源总和
 ResourceSet current_resource_load_;
};

class SchedulingQueue {
public:
 /// 创建一个空的调度队列，初始化各个状态对应的任务队列，就绪队列被单独拿出来用 ReadyQueue 做初始化。
 SchedulingQueue() : ready_queue_(std::make_shared<ReadyQueue>()) {
   for (const auto &task_state : {
            TaskState::PLACEABLE,
            TaskState::WAITING,
            TaskState::READY,
            TaskState::RUNNING,
            TaskState::INFEASIBLE,
            TaskState::WAITING_FOR_ACTOR_CREATION,
            TaskState::SWAP,
        }) {
     if (task_state == TaskState::READY) {
       task_queues_[static_cast<int>(task_state)] = ready_queue_;
    } else {
       task_queues_[static_cast<int>(task_state)] = std::make_shared<TaskQueue>();
    }
  }
}
 // 我觉得名字起得不好，他的实际操作是获取所有就绪任务资源需求之和
 ResourceSet GetResourceLoad() const;
 /// 单个任务的增删查
 bool HasTask(const TaskID &task_id) const;
 const Task &GetTaskOfState(const TaskID &task_id, TaskState task_state) const;
 bool RemoveTask(const TaskID &task_id, Task *removed_task, TaskState *removed_task_state = nullptr);
 // 按状态获取任务，对于就绪状态，还需要按资源进行聚集
 const std::unordered_map<ResourceSet, ordered_set<TaskID>> &GetReadyTasksWithResources() const;
 const std::list<Task> &GetTasks(TaskState task_state) const;
 // 一组任务的移来移去
 std::vector<Task> RemoveTasks(std::unordered_set<TaskID> &task_ids);
 void QueueTasks(const std::vector<Task> &tasks, TaskState task_state);
 void MoveTasks(std::unordered_set<TaskID> &tasks, TaskState src_state, TaskState dst_state);
 void FilterState(std::unordered_set<TaskID> &task_ids, TaskState filter_state) const;
 /// 这两个函数是按其他维度：Actor 和 Job 来获取一组任务
 std::unordered_set<TaskID> GetTaskIdsForJob(const JobID &job_id) const;
 std::unordered_set<TaskID> GetTaskIdsForActor(const ActorID &actor_id) const;
 /// 阻塞任务和用户进程增删改查
 const std::unordered_set<TaskID> &GetBlockedTaskIds() const;
 const std::unordered_set<TaskID> &GetDriverTaskIds() const;
 void AddBlockedTaskId(const TaskID &task_id);
 void RemoveBlockedTaskId(const TaskID &task_id);
 void AddDriverTaskId(const TaskID &task_id);
 void RemoveDriverTaskId(const TaskID &task_id);
 /// 用来调试和监控
 std::string DebugString() const;
 void RecordMetrics() const;
 /// 这个好像没啥用，都没实现
 ResourceSet GetReadyQueueResources() const;
private:
 /// 一个辅助函数，由于调度队列算是有两层索引 task state -> (task id -> task)，
 /// 因此经常需要定位到某个状态对应的任务队列，进而获取其中的某个任务。
 const std::shared_ptr<TaskQueue> &GetTaskQueue(TaskState task_state) const;
 /// 两个辅助函数，用来在指定状态的任务队列中删除或者过滤任务的
 void RemoveTasksFromQueue(ray::raylet::TaskState task_state,
                           std::unordered_set<ray::TaskID> &task_ids,
                           std::vector<ray::Task> *removed_tasks);
 void FilterStateFromQueue(std::unordered_set<ray::TaskID> &task_ids,
                           TaskState task_state) const;
 // kNumTaskQueues 作用便在此，所有int值在其之前的状态都有对应的任务队列
 std::array<std::shared_ptr<TaskQueue>, static_cast<int>(TaskState::kNumTaskQueues)>
     task_queues_;
 // 调度时候，就绪队列用的比较多，就单独维护一个指针在此
 const std::shared_ptr<ReadyQueue> ready_queue_;
 // 这两个状态（blocked 和 driver）没有对应的任务队列，只是用集合来保存id
 std::unordered_set<TaskID> blocked_task_ids_;
 std::unordered_set<TaskID> driver_task_ids_;
};