前言

Ray的实现里面包含了大量grpc，好处就是能够充分解耦这套非常复杂的任务调度管理系统，这是我开始看代码前没有想到的，因此希望写一篇文章再梳理一下这几个服务。不知道是后来演进出来的CS架构还是一开始berkley的老哥们就已经按这个来设计了。趁天刚开始亮，开整。

rpc服务概览

GCS Server

Job管理服务

// 作业管理服务，处理Ray集群中所有job的生命周期
service JobInfoGcsService {
  // 添加新job到GCS，包含错误处理和状态跟踪
  // 常规场景: 正常提交job
  // Corner cases: 
  // 1. Job提交时GCS宕机 - 通过status返回错误
  // 2. 重复提交相同job - 通过job_id去重
  rpc AddJob(AddJobRequest) returns (AddJobReply);

  // 标记job完成，处理资源回收
  // Corner cases:
  // 1. Job已经被标记完成 - 幂等处理
  // 2. Job不存在 - 返回错误状态
  rpc MarkJobFinished(MarkJobFinishedRequest) returns (MarkJobFinishedReply);

  // 获取所有job信息，支持分页和过滤
  // 关键特性:
  // - 支持limit限制返回数量
  // - 支持按job_id或submission_id过滤
  // - 可选择跳过某些字段以优化性能
  rpc GetAllJobInfo(GetAllJobInfoRequest) returns (GetAllJobInfoReply);
}

Actor 管理服务

// Actor生命周期和状态管理服务
service ActorInfoGcsService {
  // 注册新Actor，处理依赖解析和资源分配
  // 错误处理:
  // 1. 依赖不可用 - 返回错误并清理
  // 2. 资源不足 - 进入等待队列
  rpc RegisterActor(RegisterActorRequest) returns (RegisterActorReply);

  // 重启Actor，处理故障恢复
  // 关键实现:
  // - 支持有限次数重试
  // - 保持Actor状态一致性
  // - 处理依赖Actor的连锁反应
  rpc RestartActor(RestartActorRequest) returns (RestartActorReply);

  // 获取Actor信息，支持按ID或名称查询
  // 优化设计:
  // - 支持命名空间隔离
  // - 缓存常用Actor信息
  rpc GetActorInfo(GetActorInfoRequest) returns (GetActorInfoReply);
}

节点资源管理服务

// 节点资源管理和监控服务
service NodeResourceInfoGcsService {
  // 获取所有节点可用资源
  // 设计特点:
  // - 支持实时资源视图
  // - 处理资源碎片
  rpc GetAllAvailableResources(GetAllAvailableResourcesRequest) 
      returns (GetAllAvailableResourcesReply);

  // 获取资源使用情况
  // 关键功能:
  // - 支持资源使用趋势分析
  // - 检测Actor创建瓶颈
  // - 辅助自动扩缩容决策
  rpc GetAllResourceUsage(GetAllResourceUsageRequest) 
      returns (GetAllResourceUsageReply);
}

关键特性总结

NodeManager

Worker 管理与调度

// Worker 调度和生命周期管理
service NodeManagerService {
  // 请求分配 Worker
  // 常规场景: 任务调度需要 worker 执行
  // Corner cases:
  // 1. 资源不足 - 通过 retry_at_raylet_address 重定向到其他节点
  // 2. Worker 启动失败 - 通过 failure_type 返回具体原因
  // 3. 任务取消 - 通过 canceled 标记避免资源浪费
  rpc RequestWorkerLease(RequestWorkerLeaseRequest) returns (RequestWorkerLeaseReply) {
    option (google.api.http) = {
      post: "/api/worker/lease"
      body: "*"
    };
  }

  // 预启动 Workers 以减少任务调度延迟
  // 关键特性:
  // - 支持语言特定 worker
  // - 支持 runtime env
  // - 控制 worker 存活时间
  rpc PrestartWorkers(PrestartWorkers Request) returns (PrestartWorkersReply);
}

节点状态管理

// 节点生命周期和健康检查
message DrainRayletRequest {
  // 节点下线原因和策略
  ray.rpc.autoscaler.DrainNodeReason reason = 1;
  // 优雅退出截止时间
  int64 deadline_timestamp_ms = 3;
}

service NodeManagerService {
  // 节点优雅下线
  // 关键实现:
  // 1. 支持多种下线原因
  // 2. 保证任务完成或迁移
  // 3. 处理长运行服务
  rpc DrainRaylet(DrainRayletRequest) returns (DrainRayletReply);

  // 获取节点统计信息
  // 设计特点:
  // - 可选是否包含内存详情
  // - 支持对象存储统计
  rpc GetNodeStats(GetNodeStatsRequest) returns (GetNodeStatsReply);
}

对象管理

// 对象存储和生命周期管理
message ObjectStoreStats {
  // 溢写统计
  double spill_time_total_s = 1;
  int64 spilled_bytes_total = 2;
  
  // 内存使用统计
  int64 object_store_bytes_used = 7;
  int64 object_store_bytes_avail = 8;
  
  // 对象引用计数
  int64 num_local_objects = 11;
}

service NodeManagerService {
  // 固定对象在内存中
  // 错误处理:
  // 1. 对象已被驱逐 - 返回失败标记
  // 2. 内存不足 - 触发溢写
  rpc PinObjectIDs(PinObjectIDsRequest) returns (PinObjectIDsReply);
}

关键特性总结

CoreWorker

任务执行管理

// 任务推送和执行管理
service CoreWorkerService {
  // 推送任务到 Worker 执行
  // 常规场景: Worker 正常执行任务并返回结果
  // Corner cases:
  // 1. Worker 崩溃 - 通过 worker_exiting 标记
  // 2. 任务执行失败 - 通过 is_application_error 区分应用错误
  // 3. 任务取消 - 通过 was_cancelled_before_running 标记
  rpc PushTask(PushTaskRequest) returns (PushTaskReply) {
    option (google.api.http) = {
      post: "/api/task/push"
      body: "*"
    };
  }

  // 任务执行结果包含:
  message PushTaskReply {
    // 返回对象列表
    repeated ReturnObject return_objects = 1;
    // 动态创建的对象
    repeated ReturnObject dynamic_return_objects = 2;
    // 借用的对象引用
    repeated ObjectReferenceCount borrowed_refs = 4;
    // 错误处理标记
    bool is_retryable_error = 5;
    bool is_application_error = 6;
  }
}

Actor 生命周期管理

// Actor 状态和引用管理
service CoreWorkerService {
  // 等待 Actor 引用被删除
  // 关键实现:
  // 1. 引用计数跟踪
  // 2. 级联清理
  rpc WaitForActorRefDeleted(WaitForActorRefDeletedRequest) 
      returns (WaitForActorRefDeletedReply);

  // Kill Actor 请求
  // 错误处理:
  // 1. 强制终止 - force_kill 参数
  // 2. Actor 已死亡 - 通过 death_cause 记录原因
  rpc KillActor(KillActorRequest) returns (KillActorReply);
}

// Actor 持久化状态
message ActorHandle {
  bytes actor_id = 1;
  bytes owner_id = 2;
  Address owner_address = 3;
  // 支持任务重试
  int64 max_task_retries = 9;
  // 支持命名 Actor
  string name = 10;
  string ray_namespace = 11;
  // 支持乱序执行
  bool execute_out_of_order = 12;
  // 限制等待任务数
  int32 max_pending_calls = 13;
}

对象生命周期管理

// 对象状态和位置管理
service CoreWorkerService {
  // 获取对象状态
  // 设计要点:
  // 1. 支持多种状态: CREATED/OUT_OF_SCOPE/FREED
  // 2. 追踪对象位置
  // 3. 处理嵌套对象
  rpc GetObjectStatus(GetObjectStatusRequest) returns (GetObjectStatusReply);

  // 更新对象位置信息
  // 关键特性:
  // 1. 批量更新提升性能
  // 2. 支持等离子存储和外部存储
  // 3. 处理动态创建的对象
  rpc UpdateObjectLocationBatch(UpdateObjectLocationBatchRequest)
      returns (UpdateObjectLocationBatchReply);
}

message ObjectLocationUpdate {
  bytes object_id = 1;
  optional ObjectPlasmaLocationUpdate plasma_location_update = 2;
  optional ObjectSpilledLocationUpdate spilled_location_update = 3;
  optional bytes generator_id = 4;
}

资源和内存管理

// 内存管理和垃圾回收
service CoreWorkerService {
  // 本地垃圾回收
  // 实现细节:
  // 1. 支持全局触发
  // 2. 处理循环引用
  rpc LocalGC(LocalGCRequest) returns (LocalGCReply);

  // 对象溢写管理
  // 关键特性:
  // 1. 批量操作
  // 2. 支持恢复
  // 3. 清理策略
  rpc SpillObjects(SpillObjectsRequest) returns (SpillObjectsReply);
  rpc RestoreSpilledObjects(RestoreSpilledObjectsRequest)
      returns (RestoreSpilledObjectsReply);
}

关键特性总结

ObjectManager

对象传输服务

// 对象传输服务定义
service ObjectManagerService {
  // 对象推送服务
  // 常规场景: 节点间对象传输
  // Corner cases:
  // 1. 传输失败 - 通过 push_id 重试
  // 2. 目标节点故障 - 返回错误重新选择节点
  // 3. 大对象分块 - 通过 chunk_index 重组
  rpc Push(PushRequest) returns (PushReply);

  // 对象拉取服务
  // 常规场景: 按需获取远程对象
  // Corner cases:
  // 1. 对象不存在 - 返回错误触发重建
  // 2. 源节点故障 - 从其他副本获取
  rpc Pull(PullRequest) returns (PullReply);
}

// 对象推送请求
message PushRequest {
  // 推送会话ID，用于重试去重
  bytes push_id = 1;
  
  // 对象元数据
  bytes object_id = 2;
  bytes node_id = 3;
  Address owner_address = 4;
  
  // 分块传输支持
  uint32 chunk_index = 5;
  uint64 data_size = 6;
  uint64 metadata_size = 7;
  bytes data = 8;
}

内存管理

// 对象释放管理
service ObjectManagerService {
  // 释放对象服务
  // 设计要点:
  // 1. 批量释放提升性能
  // 2. 引用计数检查
  // 3. 级联清理
  rpc FreeObjects(FreeObjectsRequest) returns (FreeObjectsReply);
}

message FreeObjectsRequest {
  // 支持批量释放
  repeated bytes object_ids = 1;
}

// 对象大小管理
message PushRequest {
  // 对象大小跟踪
  uint64 data_size = 6;    // 数据大小
  uint64 metadata_size = 7; // 元数据大小
  
  // 分块传输优化
  uint32 chunk_index = 5;   // 当前块索引
  bytes data = 8;          // 块数据
}

错误处理和恢复

service ObjectManagerService {
  // Push 服务的错误处理机制
  // 1. 使用 push_id 标识传输会话
  // 2. 支持重试和故障转移
  // 3. 完整性校验
  rpc Push(PushRequest) returns (PushReply);
}

message PushRequest {
  // 传输会话追踪
  bytes push_id = 1;      // 用于重试去重
  bytes node_id = 3;      // 源节点标识
  Address owner_address = 4; // 所有者地址
}

分块传输机制

message PushRequest {
  // 分块传输支持
  uint32 chunk_index = 5;    // 块索引
  uint64 data_size = 6;      // 总大小
  uint64 metadata_size = 7;   // 元数据大小
  bytes data = 8;            // 块数据
  
  // 对象标识
  bytes object_id = 2;       // 对象ID
  bytes push_id = 1;         // 传输会话ID
}

关键特性总结

AutoScaler

设计逻辑

• GCS (C++) 负责状态管理和协调
• Python层负责具体的扩缩容逻辑
• 提供者(Provider)层处理具体云平台操作

AutoscalingState

// 表示自动扩缩容的状态
message AutoscalingState {
  // 记录上次看到的集群资源状态版本，确保状态的一致性和可追溯性。
  int64 last_seen_cluster_resource_state_version = 1;  
  // 标识当前自动扩缩容状态的版本，便于跟踪状态变化。
  int64 autoscaler_state_version = 2; 
  // 存储待处理的实例请求，确保所有请求都能被及时处理。
  repeated PendingInstanceRequest pending_instance_requests = 3; 
  // 记录不可行的资源请求，帮助用户了解哪些请求无法满足。
  repeated ResourceRequest infeasible_resource_requests = 4;  
  // 类似于不可行的资源请求，但专门针对群组资源请求。
  repeated GangResourceRequest infeasible_gang_resource_requests = 5;
  // 不可行的集群资源约束
  repeated ClusterResourceConstraint infeasible_cluster_resource_constraints = 6; 
  // 跟踪已启动但尚未开始的实例，确保这些实例能够被正确管理。
  repeated PendingInstance pending_instances = 7;
  // 记录失败的实例请求，提供详细的故障信息，便于后续的故障排查和处理。
  repeated FailedInstanceRequest failed_instance_requests = 8;  
}

关键文件

• 关键python文件：
  • python/ray/autoscaler/autoscaler.py - 主要的自动扩缩容逻辑
  • python/ray/autoscaler/node_provider.py - 节点供应商抽象接口
  • python/ray/autoscaler/sdk.py - 用户侧SDK接口
  • python/ray/autoscaler/tags.py - 节点标签管理
• 核心服务实现：
  • src/ray/gcs/gcs_server/autoscaler_state_service.h
  • src/ray/gcs/gcs_server/autoscaler_state_service.cc

关键特性总结

自动扩缩容条件

1. 资源请求：当集群中的资源请求（如 CPU、内存等）增加，且当前可用资源不足以满足这些请求时，自动扩缩容机制会被触发。
2. 节点状态变化：如果某个节点的状态变为 DEAD 或 DRAINING，系统会评估是否需要启动新的节点来替代这些节点。
3. 用户请求：用户通过 API 调用（如 ray.autoscaler.sdk.request_resources）请求额外的资源时，系统会根据当前的集群状态和资源约束来决定是否进行扩容。
4. 集群约束：如果集群资源约束（如最小 / 最大节点数）被触发，系统会根据这些约束来调整集群规模。
5. 负载均衡：在某些情况下，系统会根据负载均衡策略自动调整节点数量，以确保工作负载的均匀分布。

RunTimeEnvAgent

message GetOrCreateRuntimeEnvRequest {
  /// Json serialized [RuntimeEnv].
  /// For details, checkout "//python/ray/runtime_env/runtime_env.py".
  string serialized_runtime_env = 1; // 请求创建或获取运行环境的序列化信息

  RuntimeEnvConfig runtime_env_config = 2; // 运行环境的配置
  /// Hex format of [JobId].
  /// For details, checkout "//src/ray/common/id.h".
  bytes job_id = 3; // 关联的作业ID，确保运行环境与作业的关联
  string source_process = 4; // 发起请求的进程信息
}

message GetOrCreateRuntimeEnvReply {
  AgentRpcStatus status = 1; // 请求的状态，成功或失败
  string error_message = 2; // 错误信息，若有
  /// Json serialized [RuntimeEnvContext].
  /// For details, checkeout "//python/ray/_private/runtime_env/context.py".
  string serialized_runtime_env_context = 3; // 返回的运行环境上下文
}

message DeleteRuntimeEnvIfPossibleRequest {
  /// Json serialized [RuntimeEnv].
  string serialized_runtime_env = 1; // 请求删除的运行环境的序列化信息
  string source_process = 2; // 发起请求的进程信息
}

message DeleteRuntimeEnvIfPossibleReply {
  AgentRpcStatus status = 1; // 请求的状态
  string error_message = 2; // 错误信息，若有
}

message GetRuntimeEnvsInfoRequest {
  // Maximum number of entries to return.
  // If not specified, return the whole entries without truncation.
  optional int64 limit = 1; // 可选的返回条目限制
}

message GetRuntimeEnvsInfoReply {
  repeated RuntimeEnvState runtime_env_states = 1; // 返回的运行环境状态列表
  // Length of the corresponding resource without truncation.
  int64 total = 2; // 返回的总条目数
}

这个RuntimeEnvAgent对于集群的管理比较关键，核心代码放在了*ray/_private/runtime_env*下面，这里不再展开了，下一篇文章详细讲解一下整体架构和实现。 **

最后

所有的rpc服务都通过grpc_server.cc的 GrpcServer::RegisterService 进行注册，由GcsServer统一进行初始化和管理。写到这里，收获还是蛮多的，特别是对于C++如何做服务发现，如何实现InternalGrpcService。Enjoy!

class GcsServer {
	protected:
	/// ....
  /// Initialize gcs node manager.
  void InitGcsNodeManager(const GcsInitData &gcs_init_data);
  /// Initialize gcs health check manager.
  void InitGcsHealthCheckManager(const GcsInitData &gcs_init_data);
  /// Initialize gcs resource manager.
  void InitGcsResourceManager(const GcsInitData &gcs_init_data);
  /// Initialize synchronization service
  void InitRaySyncer(const GcsInitData &gcs_init_data);
  /// Initialize cluster resource scheduler.
  void InitClusterResourceScheduler();
  /// Initialize cluster task manager.
  void InitClusterTaskManager();
  /// Initialize gcs job manager.
  void InitGcsJobManager(const GcsInitData &gcs_init_data);
  /// Initialize gcs actor manager.
  void InitGcsActorManager(const GcsInitData &gcs_init_data);
  /// Initialize gcs placement group manager.
  void InitGcsPlacementGroupManager(const GcsInitData &gcs_init_data);
  /// Initialize gcs worker manager.
  void InitGcsWorkerManager();
  /// Initialize gcs task manager.
  void InitGcsTaskManager();
  /// Initialize gcs autoscaling manager.
  void InitGcsAutoscalerStateManager(const GcsInitData &gcs_init_data);
	/// ...
 }