Ray 源码分析系列(5)—Raylet

Jim.Li

已于 2025-01-03 23:05:01 修改

阅读量1.2k

点赞数 19

文章标签：架构 c++ ai 分布式人工智能

于 2024-12-31 07:56:57 首次发布

本文链接：https://blog.youkuaiyun.com/weixin_43956669/article/details/144838968

版权

Raylet

Raylet 是 Ray 框架中的核心组件，负责任务调度和对象管理。作为一个轻量级的进程，Raylet 运行在每个工作节点上，协调计算资源的分配和任务的执行。它通过高效的调度算法，确保任务能够在可用的工作节点上快速执行，同时管理对象的生命周期，包括对象的存储、传输和释放。

Raylet 通过与对象存储（如 Plasma）进行交互，处理对象的创建和访问请求，确保数据在分布式环境中的一致性和可用性。它还支持动态资源管理，能够根据工作负载的变化自动调整资源分配。此外，Raylet 提供了远程过程调用（RPC）接口，使得不同节点之间能够高效地通信。

时序图

大家看一眼时序图基本能get到raylet以及gcs在整个任务执行过程中是如何交互的了。

代码概览

src/ray/raylet

文件名	作用描述
`node_manager.h/cc`	核心组件，负责管理节点上的资源分配、任务调度、对象管理和工作进程。包含以下主要功能： - 工作进程的生命周期管理 - 资源调度和分配 - 任务调度和执行 - 对象存储管理 - 与 GCS 的通信协调
`raylet.h/cc`	Raylet 服务的主要实现，负责初始化和协调各个组件，包括： - 启动和停止服务 - 处理客户端连接 - 注册到 GCS - 管理节点状态
`wait_manager.h/cc`	管理对象等待操作，处理客户端的等待请求： - 跟踪对象可用性 - 处理超时 - 通知等待完成
`worker_pool.h/cc`	管理工作进程池： - 创建和销毁工作进程 - 分配工作进程给任务 - 维护进程状态
`local_task_manager.h/cc`	管理任务的执行和调度： - 任务队列管理 - 任务调度策略 - 任务状态跟踪
`local_object_manager.h/cc`	管理对象的存储和传输： - 对象生命周期管理 - 对象传输协调 - 内存管理
`placement_group_resource_manager.h/cc`	管理放置组资源： - 资源预留 - 资源分配 - 放置组策略
`agent_manager.h/cc`	管理 Ray 的代理进程： - Dashboard 代理 - Runtime 环境代理 - 代理进程生命周期管理
`runtime_env_agent_client.h/cc`	提供与运行时环境代理的通信接口： - 环境设置 - 环境清理 - 状态查询
`worker_killing_policy.h/cc`	定义工作进程终止策略： - OOM 处理 - 资源回收 - 进程清理

src/ray/raylet/scheduling

文件名	主要职责	核心功能	关键类/接口
cluster_resource_scheduler.h/cc	集群资源调度管理	- 资源分配和跟踪 - 节点资源管理 - 负载均衡	ClusterResourceScheduler ResourceSet NodeResources
cluster_task_manager.h/cc	集群任务调度管理	- 任务队列管理 - 任务依赖处理 - 跨节点任务协调	ClusterTaskManager TaskQueue TaskDependencies
scheduling_ids.h	调度ID定义和管理	- ID生成 - ID序列化 - ID类型定义	SchedulingID SchedulingKey TaskID
scheduling_policy.h/cc	调度策略定义和实现	- 策略接口定义 - 默认策略实现 - 策略扩展支持	SchedulingPolicy DefaultPolicy CustomPolicy
scheduling_queue.h/cc	调度队列管理	- 优先级队列 - 任务排序 - 队列状态管理	SchedulingQueue PriorityQueue TaskQueue
scheduling_stats.h/cc	调度统计和监控	- 性能指标收集 - 统计数据管理 - 监控接口	SchedulingStats MetricsCollector StatsExporter

核心代码详解

raylet

构造函数

Raylet(instrumented_io_context &main_service,    // 主事件循环服务
       const NodeID &self_node_id,               // 节点唯一标识
       const std::string &socket_name,           // Unix domain socket 名称
       const std::string &node_ip_address,       // 节点 IP 地址
       const std::string &node_name,             // 节点名称
       const NodeManagerConfig &node_manager_config,    // 节点管理器配置
       const ObjectManagerConfig &object_manager_config, // 对象管理器配置
       std::shared_ptr<gcs::GcsClient> gcs_client,     // GCS 客户端
       int metrics_export_port,                  // 指标导出端口
       bool is_head_node,                        // 是否为头节点
       std::function<void()> shutdown_handler);  // 关闭处理函数

核心功能方法

class Raylet {
   
 public:
  // 启动 Raylet 服务
  void Start();

  // 停止 Raylet 服务
  void Stop();

  // 从 GCS 注销自己
  void UnregisterSelf(const rpc::NodeDeathInfo &node_death_info,
                     std::function<void()> unregister_done_callback);

  // 处理新的客户端连接
  void ProcessNewClient(ClientConnection &client);

  // 获取节点 ID
  NodeID GetNodeId() const {
    return self_node_id_; }

  // 获取节点管理器引用
  NodeManager &node_manager() {
    return node_manager_; }

私有成员变量

 private:
  // 节点标识
  NodeID self_node_id_;
  
  // 节点信息
  GcsNodeInfo self_node_info_;

  // GCS 客户端连接
  std::shared_ptr<gcs::GcsClient> gcs_client_;

  // 节点管理器实例
  NodeManager node_manager_;

  // Unix domain socket 名称
  std::string socket_name_;

  // 客户端连接接收器
  boost::asio::basic_socket_acceptor<local_stream_protocol> acceptor_;
  
  // 用于新连接的 socket
  local_stream_socket socket_;

主要工作流程

1. **构造时初始化各组件**
   - 在创建 Raylet 实例时，初始化所需的各个组件和配置。

2. **Start() 启动服务**
   - 调用 `Start()` 方法，启动 Raylet 服务。

3. **RegisterGcs() 向 GCS 注册**
   - 在 `Start()` 方法中，调用 `RegisterGcs()` 向全局控制服务（GCS）注册节点信息。

4. **DoAccept() 开始接受客户端连接**
   - 调用 `DoAccept()` 方法，开始异步接受来自客户端的连接请求。

5. **ProcessNewClient() 处理新客户端**
   - 在 `HandleAccept()` 方法中，调用 `ProcessNewClient()` 处理新连接的客户端。

6. **通过 node_manager_ 处理任务和资源**
   - 使用 `node_manager_` 处理任务调度、资源管理和对象管理。

7. **Stop() 停止服务**
   - 当需要停止服务时，调用 `Stop()` 方法，停止节点管理器并关闭 socket。

8. **UnregisterSelf() 从 GCS 注销**
   - 在服务停止时，调用 `UnregisterSelf()` 从 GCS 注销节点信息。

local_task_manager

LocalTaskManager 是 Ray 中负责本地任务管理的核心组件，它：

管理任务的完整生命周期
处理任务依赖关系
协调资源分配
管理工作进程
处理任务调度和执行
实现任务溢出机制
基本定义与职责
构造函数

LocalTaskManager(
    const NodeID &self_node_id,                  // 本地节点ID
    ClusterResourceScheduler &cluster_resource_scheduler,  // 集群资源调度器
    TaskDependencyManagerInterface &task_dependency_manager,  // 任务依赖管理器
    std::function<bool(const WorkerID &, const NodeID &)> is_owner_alive,  // 检查所有者存活
    internal::NodeInfoGetter get_node_info,      // 获取节点信息
    WorkerPoolInterface &worker_pool,            // 工作进程池
    absl::flat_hash_map<WorkerID, std::shared_ptr<WorkerInterface>> &leased_workers,  // 租用的工作进程
    std::function<bool(const std::vector<ObjectID> &object_ids,
                     std::vector<std::unique_ptr<RayObject>> *results)>
        get_task_arguments,                      // 获取任务参数
    size_t max_pinned_task_arguments_bytes,      // 最大固定任务参数字节数
    std::function<int64_t(void)> get_time_ms,   // 获取当前时间
    int64_t sched_cls_cap_interval_ms);         // 调度类容量间隔

核心方法

class LocalTaskManager {
 public:
  // 队列和调度任务
  void QueueAndScheduleTask(std::shared_ptr<internal::Work> work);

  // 调度和分发任务
  void ScheduleAndDispatchTasks();

  // 当任务依赖解除阻塞时调用
  void TasksUnblocked(const std::vector<TaskID> &ready_ids);

  // 任务完成时调用
  void TaskFinished(std::shared_ptr<WorkerInterface> worker, RayTask *task);

  // 取消任务
  bool CancelTasks(std::function<bool(const std::shared_ptr<internal::Work> &)> predicate,
                   rpc::RequestWorkerLeaseReply::SchedulingFailureType failure_type,
                   const std::string &scheduling_failure_message);

资源管理方法

  // 释放工作进程资源
  void ReleaseWorkerResources(std::shared_ptr<WorkerInterface> worker);

  // 从被阻塞的工作进程释放 CPU 资源
  bool ReleaseCpuResourcesFromBlockedWorker(std::shared_ptr<WorkerInterface> worker);

  // 返回 CPU 资源给解除阻塞的工作进程
  bool ReturnCpuResourcesToUnblockedWorker(std::shared_ptr<WorkerInterface> worker);

  // 计算普通任务资源
  ResourceSet CalcNormalTaskResources() const;

任务调度和跟踪

  // 设置工作进程积压
  void SetWorkerBacklog(SchedulingClass scheduling_class,
                       const WorkerID &worker_id,
                       int64_t backlog_size);

  // 清除工作进程积压
  void ClearWorkerBacklog(const WorkerID &worker_id);

  // 获取要分发的任务
  const absl::flat_hash_map<SchedulingClass, std::deque<std::shared_ptr<internal::Work>>>
      &GetTaskToDispatch() const;

  // 获取积压跟踪器
  const absl::flat_hash_map<SchedulingClass, absl::flat_hash_map<WorkerID, int64_t>>
      &GetBackLogTracker() const;

内部调度类信息结构

  struct SchedulingClassInfo {
    // 跟踪此调度类中运行的任务
    absl::flat_hash_set<TaskID> running_tasks;
    // 此调度类可以运行的任务总数
    const uint64_t capacity;
    // 此调度类下一次可以分发新任务的时间
    int64_t next_update_time;
  };

私有变量

 private:
  // 节点标识
  const NodeID &self_node_id_;
  
  // 集群资源调度器
  ClusterResourceScheduler &cluster_resource_scheduler_;
  
  // 任务依赖管理器
  TaskDependencyManagerInterface &task_dependency_manager_;
  
  // 工作进程池
  WorkerPoolInterface &worker_pool_;
  
  // 等待分发的任务队列
  absl::flat_hash_map<SchedulingClass, std::deque<std::shared_ptr<internal::Work>>>
      tasks_to_dispatch_;
      
  // 等待任务队列
  std::list<std::shared_ptr<internal::Work>> waiting_task_queue_;

基本工作流程

// 1. 接收新任务（QueueAndScheduleTask）
// 2. 检查任务依赖
// 3. 等待依赖满足（waiting_task_queue_）
// 4. 准备分发任务（tasks_to_dispatch_）
// 5. 分配资源
// 6. 分配工作进程
// 7. 执行任务
// 8. 处理任务完成
// 9. 释放资源