Core Worker

CoreWorker 是 Ray 中最核心的组件之一,它封装了分布式系统的复杂性,为上层应用提供简单统一的编程模型。它提供了任务执行、对象管理、Actor调用等核心功能，每个组件都有明确的职责划分，通过良好的接口设计实现了功能解耦和高效协作。

核心文件功能总结

核心组件交互

• 内部时序图：

• 外部时序图

核心代码分析

CoreWorker

1. 核心类结构

//CoreWorker 核心作用总结：
//- 管理任务的执行和调度
//- 维护 Actor 的生命周期
//- 管理分布式对象的存储和访问
//- 处理进程间通信和 RPC 调用
//- 提供对象引用计数和内存管理
//- 支持任务重试和错误处理
//- 与 Raylet 和 GCS 交互
//- 提供性能监控和统计
//- 支持资源管理和调度
//- 实现分布式系统的容错机制

class CoreWorker {
 private:
  // === 核心组件 ===
  // 配置选项
  const CoreWorkerOptions options_;
  
  // 工作进程上下文
  WorkerContext worker_context_;
  
  // 核心服务组件
  std::unique_ptr<raylet::RayletClient> raylet_client_;        // Raylet客户端
  std::unique_ptr<CoreWorkerMemoryStore> memory_store_;        // 内存存储
  std::unique_ptr<ReferenceCounter> reference_counter_;        // 引用计数
  std::unique_ptr<TaskManager> task_manager_;                  // 任务管理
  std::unique_ptr<ActorManager> actor_manager_;               // Actor管理
  std::shared_ptr<gcs::GcsClient> gcs_client_;               // GCS客户端
  
  // 任务执行相关
  std::unique_ptr<TaskReceiver> task_receiver_;               // 任务接收器
  std::unique_ptr<TaskSubmitter> direct_task_submitter_;      // 任务提交器
};

1. 核心功能接口:

class CoreWorker {
 public:
  // === 任务相关 ===
  // 提交任务
  Status SubmitTask(const TaskSpecification &task_spec);
  
  // 执行任务
  Status ExecuteTask(const TaskSpecification &task_spec,
                    const std::shared_ptr<TaskResourceIds> &resource_ids);
                    
  // 任务执行循环
  void RunTaskExecutionLoop();

  // === 对象管理 ===
  // 创建对象
  Status Put(const RayObject &object, ObjectID *object_id);
  
  // 获取对象
  Status Get(const std::vector<ObjectID> &ids,
             int64_t timeout_ms,
             std::vector<std::shared_ptr<RayObject>> *results);
             
  // === Actor相关 ===
  // 创建Actor
  Status CreateActor(const ActorCreationOptions &create_options,
                    ActorID *actor_id);
                    
  // 提交Actor任务
  Status SubmitActorTask(const ActorID &actor_id,
                        const TaskSpecification &task_spec);
};

1. 初始化流程:

CoreWorker::CoreWorker(const CoreWorkerOptions &options) 
    : options_(options) {
  // 1. 初始化工作进程上下文
  worker_context_ = WorkerContext(worker_id_, options.job_id);
  
  // 2. 创建核心组件
  memory_store_ = std::make_unique<CoreWorkerMemoryStore>();
  reference_counter_ = std::make_unique<ReferenceCounter>();
  task_manager_ = std::make_unique<TaskManager>();
  
  // 3. 连接外部服务
  ConnectToRaylet();
  ConnectToGcs();
  
  // 4. 初始化任务处理
  task_receiver_ = std::make_unique<TaskReceiver>();
  direct_task_submitter_ = std::make_unique<TaskSubmitter>();
}

1. 任务执行流程:

// 注意下列部分是伪代码
void CoreWorker::RunTaskExecutionLoop() {
  while (true) {
    // 1. 接收任务
    TaskSpecification task_spec = task_receiver_->ReceiveTask();
    
    // 2. 准备执行环境
    PrepareTaskExecution(task_spec);
    
    try {
      // 3. 执行任务
      auto results = options_.task_execution_callback(task_spec);
      
      // 4. 处理返回值
      HandleTaskResults(task_spec, results);
      
    } catch (const std::exception &e) {
      // 5. 错误处理
      HandleTaskError(task_spec, e);
    }
    
    // 6. 清理
    CleanupTaskExecution(task_spec);
  }
}

1. 对象管理实现:

Status CoreWorker::Put(const RayObject &object, ObjectID *object_id) {
  // 1. 生成对象ID
  *object_id = ObjectID::FromRandom();
  
  // 2. 添加到本地存储
  RAY_RETURN_NOT_OK(memory_store_->Put(object, *object_id));
  
  // 3. 更新引用计数
  reference_counter_->AddLocalReference(*object_id);
  
  // 4. 通知Raylet
  if (object.GetSize() > options_.inline_storage_threshold) {
    RAY_RETURN_NOT_OK(raylet_client_->ObjectReady(*object_id, object));
  }
  
  return Status::OK();
}

1. 核心功能总结：

TaskManager

核心设计思路：

• 任务生命周期管理
• 错误处理和重试机制
• 状态追踪和事件记录
• 流式生成器支持
• 任务重建和恢复
• 资源管理和清理

这些实现确保了Ray系统中任务的可靠执行和恢复。

1. TaskManager 核心结构:

class TaskManager {
 private:
  // === 核心数据结构 ===
  // 任务条目
  struct TaskEntry {
    // 任务规范
    TaskSpecification spec;
    // 任务状态
    TaskStatus status;
    // 重试次数
    uint64_t num_retries;
    // 返回对象引用
    std::vector<rpc::ObjectReference> return_refs;
    // 依赖对象
    std::vector<ObjectID> dependencies;
    // 回调函数
    std::function<void(const TaskSpecification &)> on_complete;
  };

  // 任务表
  absl::flat_hash_map<TaskID, std::shared_ptr<TaskEntry>> pending_tasks_;
  // 任务依赖图
  absl::flat_hash_map<ObjectID, std::vector<TaskID>> task_dependencies_;
};

1. 任务提交与管理:

class TaskManager {
 public:
  // 添加待处理任务
  std::vector<rpc::ObjectReference> AddPendingTask(
      const rpc::Address &caller_address,
      const TaskSpecification &spec,
      const std::string &call_site,
      int max_retries) {
    auto task_id = spec.TaskId();
    
    // 创建任务条目
    auto entry = std::make_shared<TaskEntry>();
    entry->spec = spec;
    entry->status = TaskStatus::PENDING;
    entry->num_retries = 0;
    
    // 生成返回对象引用
    std::vector<rpc::ObjectReference> return_refs;
    for (size_t i = 0; i < spec.NumReturns(); i++) {
      auto ref = CreateReturnObjectReference(spec.ReturnId(i), caller_address);
      return_refs.push_back(ref);
      entry->return_refs.push_back(ref);
    }
    
    // 记录依赖关系
    RecordTaskDependencies(task_id, spec.GetDependencies());
    
    // 添加到待处理表
    pending_tasks_.emplace(task_id, entry);
    
    return return_refs;
  }
};

1. ObjectRefStream - 对象流管理:

class ObjectRefStream {
  // 核心成员
  std::unordered_set<ObjectID> refs_written_to_stream_;  // 已写入流的对象
  std::unordered_set<ObjectID> temporarily_owned_refs_;   // 临时持有的对象
  int64_t next_index_ = 0;                               // 下一个读取位置
  int64_t end_of_stream_index_ = -1;                     // 流结束位置
  int64_t max_index_seen_ = -1;                          // 已见到的最大索引

  // 核心方法实现
  Status TryReadNextItem(ObjectID *object_id_out) {
    // 尝试读取下一个对象
    // 1. 如果对象已写入流,返回对象并更新计数
    // 2. 如果对象未写入,返回空对象让调用者重试
    *object_id_out = GetObjectRefAtIndex(next_index_);
    
    if (refs_written_to_stream_.find(*object_id_out) != refs_written_to_stream_.end()) {
      total_num_object_consumed_++;  // 更新消费计数
      next_index_++;                 // 移动读取位置
    } else {
      *object_id_out = ObjectID::Nil();  // 返回空对象
    }
    return Status::OK();
  }
};

4. 任务完成处理:

void TaskManager::CompletePendingTask(const TaskID &task_id, ...) {
  // 1. 处理动态返回对象
  if (reply.dynamic_return_objects_size() > 0) {
    // 对于第一次执行:
    // - 添加动态返回对象引用
    // - 记录需要存储在plasma的对象
    if (first_execution) {
      reference_counter_.AddDynamicReturn(object_id, generator_id);
      dynamic_return_ids.push_back(object_id);
    }
  }

  // 2. 处理流式生成器任务
  if (spec.IsStreamingGenerator()) {
    if (first_execution) {
      // 设置流式生成器返回数量
      spec.SetNumStreamingGeneratorReturns(num_streaming_generator_returns);
      
      // 记录需要在plasma存储的对象
      for (const auto &return_id_info : reply.streaming_generator_return_ids()) {
        if (return_id_info.is_plasma_object()) {
          it->second.reconstructable_return_ids.insert(...);
        }
      }
    }
  }

  // 3. 更新任务状态
  SetTaskStatus(it->second, 
    is_application_error ? rpc::TaskStatus::FAILED : rpc::TaskStatus::FINISHED);
}

1. 任务失败和重试处理:

bool TaskManager::FailOrRetryPendingTask(const TaskID &task_id, ...) {
  // 1. 尝试重试任务
  bool will_retry = false;
  if (!fail_immediately) {
    will_retry = RetryTaskIfPossible(task_id, error_info);
  }

  // 2. 如果不能重试,标记任务失败
  if (!will_retry && mark_task_object_failed) {
    FailPendingTask(task_id, error_type, status, ray_error_info);
  }

  // 3. 检查是否需要关闭
  ShutdownIfNeeded();
  return will_retry;
}

1. 任务状态管理

void TaskManager::SetTaskStatus(TaskEntry &task_entry, 
                              rpc::TaskStatus status,
                              const absl::optional<const rpc::RayErrorInfo> &error_info) {
  // 1. 更新任务状态
  task_entry.SetStatus(status);
  
  // 2. 记录任务状态事件
  task_event_buffer_.RecordTaskStatusEventIfNeeded(
      task_entry.spec.TaskId(),
      task_entry.spec.JobId(),
      task_entry.spec.AttemptNumber(),
      task_entry.spec,
      status,
      /* include_task_info */ false,
      worker::TaskStatusEvent::TaskStateUpdate(error_info));
}

1. 重建任务管理:

std::unordered_map<rpc::LineageReconstructionTask, uint64_t>
TaskManager::GetOngoingLineageReconstructionTasks(...) {
  // 1. 遍历所有任务
  for (const auto &task_it : submissible_tasks_) {
    const auto &task_entry = task_it.second;
    
    // 2. 只关注待处理且已成功执行过的任务
    if (!task_entry.IsPending() || task_entry.num_successful_executions == 0) {
      continue;
    }

    // 3. 根据任务类型收集重建信息
    if (task_entry.spec.IsNormalTask()) {
      // 普通任务 - 使用任务标签
      task.mutable_labels()->insert(...);
    } else if (task_entry.spec.IsActorTask()) {
      // Actor任务 - 使用Actor标签
      const auto &labels = actor_handle->GetLabels();
      task.mutable_labels()->insert(labels.begin(), labels.end());
    }
  }
}

ActorManager

• 设计亮点

• 核心定义

class ActorManager {
  // 核心成员:
  std::shared_ptr<gcs::GcsClient> gcs_client_; // GCS客户端,用于与全局控制存储交互
  ActorTaskSubmitterInterface &actor_task_submitter_; // 任务提交接口
  ReferenceCounterInterface &reference_counter_; // 引用计数管理
  
  // Actor句柄缓存,保存ActorID到ActorHandle的映射
  absl::flat_hash_map<ActorID, std::shared_ptr<ActorHandle>> actor_handles_;
  
  // 命名Actor缓存,避免频繁访问GCS
  absl::flat_hash_map<std::string, ActorID> cached_actor_name_to_ids_;
  
  // 记录Actor存活状态
  absl::flat_hash_map<ActorID, bool> subscribed_actors_;
}

• 核心实现

// 1. Actor Handle 注册与管理
ActorID ActorManager::RegisterActorHandle(...) {
  // 核心实现:
  // - 为 Actor 创建唯一标识符
  // - 添加 Actor Handle 到内部管理表
  // - 设置引用计数追踪
  const ActorID actor_id = actor_handle->GetActorID();
  RAY_UNUSED(AddActorHandle(...));
  // 为 Actor Handle 添加借用对象引用计数
  reference_counter_.AddBorrowedObject(actor_handle_id, outer_object_id, owner_address);
  return actor_id;
}

// 2. Named Actor 的查找与缓存机制
std::pair<std::shared_ptr<const ActorHandle>, Status> 
ActorManager::GetNamedActorHandle(...) {
  // 优先从本地缓存查找
  ActorID actor_id = GetCachedNamedActorID(name);
  if (!actor_id.IsNil()) {
    return std::make_pair(GetActorHandle(actor_id), Status::OK());
  }

  // 缓存未命中则从 GCS 同步获取
  rpc::ActorTableData actor_table_data;
  const auto status = gcs_client_->Actors().SyncGetByName(...);
  
  // 获取成功后更新本地缓存
  if (status.ok()) {
    auto actor_handle = std::make_unique<ActorHandle>(actor_table_data, task_spec);
    AddNewActorHandle(...);
  }
}

// 3. Actor 状态管理与通知机制
void ActorManager::HandleActorStateNotification(...) {
  // 处理 Actor 的不同状态转换
  if (actor_data.state() == rpc::ActorTableData::RESTARTING) {
    // 处理重启状态
    actor_task_submitter_.DisconnectActor(..., /*is_dead=*/false, ...);
  } else if (actor_data.state() == rpc::ActorTableData::DEAD) {
    // 处理死亡状态
    OnActorKilled(actor_id);
    actor_task_submitter_.DisconnectActor(..., /*is_dead=*/true, ...);
  } else if (actor_data.state() == rpc::ActorTableData::ALIVE) {
    // 处理活跃状态
    actor_task_submitter_.ConnectActor(...);
  }
}

MemoryStore

• 设计亮点

• 核心数据结构

class CoreWorkerMemoryStore {
private:
  // 对象存储的核心数据结构
  absl::flat_hash_map<ObjectID, std::shared_ptr<RayObject>> objects_;
  
  // 跟踪Get请求的数据结构
  absl::flat_hash_map<ObjectID, std::vector<std::shared_ptr<GetRequest>>> object_get_requests_;
  
  // 跟踪异步Get请求的回调函数
  absl::flat_hash_map<ObjectID, std::vector<std::function<void(std::shared_ptr<RayObject>)>>> 
    object_async_get_requests_;
}

• 核心实现-Put

bool CoreWorkerMemoryStore::Put(const RayObject &object, const ObjectID &object_id) {
  // 1. 创建对象副本或使用自定义分配器
  std::shared_ptr<RayObject> object_entry = nullptr;
  if (object_allocator_) {
    object_entry = object_allocator_(object, object_id);
  } else {
    object_entry = std::make_shared<RayObject>(...);
  }

  {
    absl::MutexLock lock(&mu_);
    
    // 2. 检查对象是否已存在
    if (objects_.find(object_id) != objects_.end()) {
      return true;
    }

    // 3. 处理异步Get请求的回调
    auto async_callback_it = object_async_get_requests_.find(object_id);
    if (async_callback_it != object_async_get_requests_.end()) {
      async_callbacks = std::move(callbacks);
      object_async_get_requests_.erase(async_callback_it);
    }

    // 4. 处理同步Get请求
    bool should_add_entry = true;
    auto object_request_iter = object_get_requests_.find(object_id);
    if (object_request_iter != object_get_requests_.end()) {
      // 通知等待的Get请求
      for (auto &get_request : get_requests) {
        get_request->Set(object_id, object_entry);
        if (get_request->ShouldRemoveObjects() && ref_counter_ == nullptr) {
          should_add_entry = false;
        }
      }
    }

    // 5. 根据引用计数决定是否存储
    if (ref_counter_ && !ref_counter_->HasReference(object_id)) {
      should_add_entry = false;
    }

    // 6. 存储对象或更新统计信息
    if (should_add_entry) {
      EmplaceObjectAndUpdateStats(object_id, object_entry);
    } else {
      OnDelete(object_entry);
    }
  }
}

• 核心实现-Get

Status CoreWorkerMemoryStore::GetImpl(...) {
  // 1. 首先检查已有对象
  {
    absl::MutexLock lock(&mu_);
    for (size_t i = 0; i < object_ids.size() && count < num_objects; i++) {
      auto iter = objects_.find(object_id);
      if (iter != objects_.end()) {
        // 找到对象，更新访问状态
        iter->second->SetAccessed();
        (*results)[i] = iter->second;
        count++;
      } else {
        remaining_ids.insert(object_id);
      }
    }
  }

  // 2. 创建等待请求
  get_request = std::make_shared<GetRequest>(
    std::move(remaining_ids),
    required_objects,
    remove_after_get,
    abort_if_any_object_is_exception
  );

  // 3. 等待对象就绪或超时
  while (!timed_out && signal_status.ok() && 
         !(done = get_request->Wait(iteration_timeout))) {
    // 检查信号和超时
    if (check_signals_) {
      signal_status = check_signals_();
    }
    // 更新超时时间
    if (remaining_timeout >= 0) {
      remaining_timeout -= iteration_timeout;
      timed_out = remaining_timeout <= 0;
    }
  }
}

• 核心实现-错误机制处理

inline bool IsUnhandledError(const std::shared_ptr<RayObject> &obj) {
  rpc::ErrorType error_type;
  return obj->IsException(&error_type) &&
         // 只对任务失败发出警告
         (error_type == rpc::ErrorType::WORKER_DIED ||
          error_type == rpc::ErrorType::TASK_EXECUTION_EXCEPTION) &&
         !obj->WasAccessed();
}

void CoreWorkerMemoryStore::NotifyUnhandledErrors() {
  // 定期扫描未处理的错误
  int64_t threshold = absl::GetCurrentTimeNanos() - kUnhandledErrorGracePeriodNanos;
  auto it = objects_.begin();
  int count = 0;
  while (it != objects_.end() && count < kMaxUnhandledErrorScanItems) {
    if (IsUnhandledError(obj) && obj->CreationTimeNanos() < threshold) {
      unhandled_exception_handler_(*obj);
    }
    count++;
  }
}

ReferenceCounter

• 设计亮点和特点：

• 基础引用技术管理

// 这是最基础的本地引用计数增加实现,主要处理:
// 1. 空对象ID的检查
// 2. 新对象引用的创建
// 3. 引用计数增加
// 4. 嵌套对象引用状态维护

void ReferenceCounter::AddLocalReference(const ObjectID &object_id, 
                                       const std::string &call_site) {
  if (object_id.IsNil()) {
    return;
  }
  absl::MutexLock lock(&mutex_);
  auto it = object_id_refs_.find(object_id);
  if (it == object_id_refs_.end()) {
    // 如果对象不存在,创建新的引用记录
    it = object_id_refs_.emplace(object_id, Reference(call_site, -1)).first;
  }
  bool was_in_use = it->second.RefCount() > 0;
  it->second.local_ref_count++;
  
  // 如果对象从未使用变为使用中,需要递归设置嵌套引用状态
  if (!was_in_use && it->second.RefCount() > 0) {
    SetNestedRefInUseRecursive(it);
  }
}

• 分布式引用追踪

// 这个函数处理分布式场景下的引用借用:
// 1. 验证对象所有权
// 2. 记录借用者信息
// 3. 设置引用释放等待

void ReferenceCounter::AddBorrowerAddress(const ObjectID &object_id,
                                        const rpc::Address &borrower_address) {
  absl::MutexLock lock(&mutex_);
  auto it = object_id_refs_.find(object_id);
  RAY_CHECK(it != object_id_refs_.end());
  RAY_CHECK(it->second.owned_by_us)
      << "AddBorrowerAddress should only be used for owner references.";

  // 添加借用者地址并等待其引用释放
  auto inserted = it->second.mutable_borrow()->borrowers.insert(borrower_address).second;
  if (inserted) {
    WaitForRefRemoved(it, borrower_address);
  }
}

• 对象位置跟踪

// 这个函数处理对象溢出场景:
// 1.记录溢出状态
// 2.检查节点存活性
// 3.更新位置信息
// 4.处理节点失效恢复

bool ReferenceCounter::HandleObjectSpilled(const ObjectID &object_id,
                                         const std::string spilled_url,
                                         const NodeID &spilled_node_id) {
  absl::MutexLock lock(&mutex_);
  auto it = object_id_refs_.find(object_id);
  if (it == object_id_refs_.end()) {
    return false;
  }

  it->second.spilled = true;
  it->second.did_spill = true;
  
  // 检查溢出节点是否存活
  bool spilled_location_alive = 
    spilled_node_id.IsNil() || check_node_alive_(spilled_node_id);
    
  if (spilled_location_alive) {
    // 更新溢出位置信息
    if (spilled_url != "") {
      it->second.spilled_url = spilled_url;
    }
    if (!spilled_node_id.IsNil()) {
      it->second.spilled_node_id = spilled_node_id; 
    }
    PushToLocationSubscribers(it);
  } else {
    // 节点失效,需要恢复对象
    UnsetObjectPrimaryCopy(it);
    objects_to_recover_.push_back(object_id);
  }
  return true;
}

• 引用清理与GC

// 这是引用清理的核心实现:
// 执行引用移除回调
// 递归处理嵌套引用
// 处理对象生命周期结束
// 清理引用记录

void ReferenceCounter::DeleteReferenceInternal(ReferenceTable::iterator it,
                                             std::vector<ObjectID> *deleted) {
  const ObjectID id = it->first;
  
  // 执行引用移除回调
  if (it->second.RefCount() == 0 && it->second.on_ref_removed) {
    it->second.on_ref_removed(id);
    it->second.on_ref_removed = nullptr;
  }

  // 处理嵌套对象引用
  if (it->second.OutOfScope(lineage_pinning_enabled_)) {
    for (const auto &inner_id : it->second.nested().contains) {
      auto inner_it = object_id_refs_.find(inner_id);
      if (inner_it != object_id_refs_.end()) {
        if (it->second.owned_by_us) {
          inner_it->second.mutable_nested()->contained_in_owned.erase(id);
        } else {
          inner_it->second.mutable_nested()->contained_in_borrowed_ids.erase(id);
        }
        DeleteReferenceInternal(inner_it, deleted);
      }
    }
    OnObjectOutOfScopeOrFreed(it);
    if (deleted) {
      deleted->push_back(id);
    }
  }

  // 清理引用
  if (it->second.ShouldDelete(lineage_pinning_enabled_)) {
    ReleaseLineageReferences(it);
    EraseReference(it);
  }
}

最后

作为ray的三大核心组件之一，CoreWorker的实现确实很solid，但分析起来也真不容易，前后依赖以及背后的交互逻辑需要抽丝剥茧。不过，分析到这里已经基本把最最核心的代码分析完了，希望你也跟我一样有满满的收获吧。See you!