Ray 源码解析二：资源抽象和调度策略

class FractionalResourceQuantity {
public:
 // 默认构造函数：resource_quantity_ = 0
 FractionalResourceQuantity();
 // 基本构造函数：指定资源量
 FractionalResourceQuantity(double resource_quantity);
 /// 加减运算
 const FractionalResourceQuantity operator+(const FractionalResourceQuantity &rhs) const;
 const FractionalResourceQuantity operator-(const FractionalResourceQuantity &rhs) const;
 void operator+=(const FractionalResourceQuantity &rhs);
 void operator-=(const FractionalResourceQuantity &rhs);
 /// 布尔运算
 bool operator==(const FractionalResourceQuantity &rhs) const;
 bool operator!=(const FractionalResourceQuantity &rhs) const;
 bool operator<(const FractionalResourceQuantity &rhs) const;
 bool operator>(const FractionalResourceQuantity &rhs) const;
 bool operator<=(const FractionalResourceQuantity &rhs) const;
 bool operator>=(const FractionalResourceQuantity &rhs) const;
 /// 浮点型的实际值
 double ToDouble() const;
private:
 // 以 1/kResourceConversionFactor 为单位的资源量大小
 int64_t resource_quantity_;
};

// 两个基本操作
bool FractionalResourceQuantity::operator==(const FractionalResourceQuantity &rhs) const {
 return resource_quantity_ == rhs.resource_quantity_;
}
bool FractionalResourceQuantity::operator<(const FractionalResourceQuantity &rhs) const {
 return resource_quantity_ < rhs.resource_quantity_;
}
// 以下调用基本操作完成定义
bool FractionalResourceQuantity::operator!=(const FractionalResourceQuantity &rhs) const {
 return !(*this == rhs);
}
bool FractionalResourceQuantity::operator>(const FractionalResourceQuantity &rhs) const {
 return rhs < *this;
}
bool FractionalResourceQuantity::operator<=(const FractionalResourceQuantity &rhs) const {
 return !(*this > rhs);
}
bool FractionalResourceQuantity::operator>=(const FractionalResourceQuantity &rhs) const {
 return !(*this < rhs);
}

class ResourceSet {
public:
 ResourceSet();
 // 三个构造函数：
 // 根据一个字典或者一个键值对列表构建 label(string)->amount(FractionalResourceQuantity) 的字典。
 ResourceSet(
     const std::unordered_map<std::string, FractionalResourceQuantity> &resource_map);
 ResourceSet(const std::unordered_map<std::string, double> &resource_map);
 ResourceSet(const std::vector<std::string> &resource_labels,
             const std::vector<double> resource_capacity);
 // 析构函数
 ~ResourceSet();
// 判断资源集合 A 是否为 B 的子集(A 中所有 label 的 amount 都不大于 B 中对应 label 的 amount)
 // 以该操作作为基本操作，可以实现接下来三个操作
 bool IsSubset(const ResourceSet &other) const;
 // 下面函数中，前两个函数的实现一毛一样
 bool operator==(const ResourceSet &rhs) const;
 bool IsEqual(const ResourceSet &other) const;
 bool IsSuperset(const ResourceSet &other) const;
 // 类似于字典的一些增删改查操作，即对某个种类的资源数量进行增删改查
 void AddOrUpdateResource(const std::string &resource_name,
                          const FractionalResourceQuantity &capacity);
 bool DeleteResource(const std::string &resource_name);
 FractionalResourceQuantity GetResource(const std::string &resource_name) const;
 bool IsEmpty() const;
// 两个集合间的加减运算。需要注意的是，在增加的时候有时候不能超过一个上界：比如节点的资源总量大小，于是有了
 // AddResourcesCapacityConstrained; 在减小的时候，资源量不能减小为负值，于是有了
 // SubtractResourcesStrict。这是两个上下界保护的函数。
 void AddResourcesCapacityConstrained(const ResourceSet &other,
                                      const ResourceSet &total_resources);
 void AddResources(const ResourceSet &other);
 void SubtractResources(const ResourceSet &other);
 void SubtractResourcesStrict(const ResourceSet &other);
 // 由于所有节点必然存在 CPU，所有任务调度时也必然需要 CPU 资源，因此单独拿出来作为一个函数。
 // 其对应的 label name 为：kCPU_ResourceLabel = "CPU"
 const ResourceSet GetNumCpus() const;
 // 返回字典形式组织的资源列表和对应数量；一个是 double 形式表示，一个是 FractionalResourceQuantity
 // 形式表示
 const std::unordered_map<std::string, double> GetResourceMap() const;
 const std::unordered_map<std::string, FractionalResourceQuantity>
     &GetResourceAmountMap() const;
 const std::string ToString() const;
private:
 // 内部资源及其数量字典
 std::unordered_map<std::string, FractionalResourceQuantity> resource_capacity_;
};

class ResourceIds {
public:
 ResourceIds();
 // 通过给定份数的资源构建 ResourceIds。resource_quantity 要么是个整数，要么是纯小数
 explicit ResourceIds(double resource_quantity);
 // 在内部实现上，整数份的资源列表用 vector<int> 来表示所有对应的 ID，每个具有一整份资源。
 // 对于小数份资源，使用一个 vector<pair> 来表示某个 ID 和其对应的资源量值
 // 因此有以下三个构造函数。
 explicit ResourceIds(const std::vector<int64_t> &whole_ids);
 explicit ResourceIds(
     const std::vector<std::pair<int64_t, FractionalResourceQuantity>> &fractional_ids);
 ResourceIds(
     const std::vector<int64_t> &whole_ids,
     const std::vector<std::pair<int64_t, FractionalResourceQuantity>> &fractional_ids);
 // 看是否有足够的要求的 resource_quantity 份资源。如果 resource_quantity 是整数，则只需看所有
 // 整数份资源列表。需要注意的是如果 resource_quantity 是小数，那么必须要有单个 ID 的资源量大于
 // resource_quantity 才可以(或者有不小于一份的整数资源，或者有大于resource_quantity的小数）
 // 而不能将两个小数凑在一块去大于 resource_quantity。因为分属于两个资源 ID 的量不能合到一块
 bool Contains(const FractionalResourceQuantity &resource_quantity) const;
 // 根据上面的原则切下来一块资源 或者 回收一块资源。
 // 分配资源时，适当地进行拆分
 // 回收资源时，适当地进行合并。
 ResourceIds Acquire(const FractionalResourceQuantity &resource_quantity);
 void Release(const ResourceIds &resource_ids);
 // 虽然语义不一样，但是和 Release 实现一样：将两个 resource_ids 加和
 ResourceIds Plus(const ResourceIds &resource_ids) const;
 // 获取整数份/小数份资源 ID 列表
 const std::vector<int64_t> &WholeIds() const;
 const std::vector<std::pair<int64_t, FractionalResourceQuantity>> &FractionalIds()
     const;
 // 看是不是该ID集合中没有任何资源
 bool TotalQuantityIsZero() const;
 // 所有资源加和以 FractionalResourceQuantity 形式返回
 FractionalResourceQuantity TotalQuantity() const;
 std::string ToString() const;
 // 通过 IncreaseCapacity 和 DecreaseCapacity 更新到指定资源量；这个是为了满足用户对
 // 对自定义资源动态调整而做的。
 void UpdateCapacity(int64_t new_capacity);
private:
 // 判断 resource_quantity 是不是一个整数
 bool IsWhole(double resource_quantity) const;
 void IncreaseCapacity(int64_t increment_quantity);
 void DecreaseCapacity(int64_t decrement_quantity);
 // 两个列表
 std::vector<int64_t> whole_ids_;
 std::vector<std::pair<int64_t, FractionalResourceQuantity>> fractional_ids_;
// 追踪总量，总量即用 FractionalResourceQuantity 表示，也说明了该类是 FractionalResourceQuantity
 // 的细化
 FractionalResourceQuantity total_capacity_;
 // 暂时性记下超额资源请求
 int64_t decrement_backlog_;
};

class ResourceIdSet {
public:
 // 各种构造函数，就是构建 unordered_map
 ResourceIdSet();
 ResourceIdSet(const ResourceSet &resource_set);
 ResourceIdSet(const std::unordered_map<std::string, ResourceIds> &available_resources);
 // 是否包含，索要和放回。和 ResourceIds 对应操作语义相同，只不过由单类资源变为了一组总资源。
 bool Contains(const ResourceSet &resource_set) const;
 ResourceIdSet Acquire(const ResourceSet &resource_set);
 void Release(const ResourceIdSet &resource_id_set);
 void ReleaseConstrained(const ResourceIdSet &resource_id_set,
                         const ResourceSet &resources_total);
 // 与 Release 实现一样
 ResourceIdSet Plus(const ResourceIdSet &resource_id_set) const;
 /// 对于某类资源数量的增删查
 // 给某类资源增加指定数量
 void AddOrUpdateResource(const std::string &resource_name, int64_t capacity);
 // 删除某类资源
 void DeleteResource(const std::string &resource_name);
 // 清空所有资源
 void Clear();
 // 获取所有可用资源，就是返回内部的散列表
 const std::unordered_map<std::string, ResourceIds> &AvailableResources() const;
 // cpu 所有任务都得用，所以单独拿出来
 ResourceIdSet GetCpuResources() const;
 // 将带标号的资源集合转变为只有数量描述的资源集合
 ResourceSet ToResourceSet() const;
 // 打印和序列化
 std::string ToString() const;
 std::vector<rpc::ResourceIdSetInfo> ToProtobuf() const;
private:
 // 从资源种类到带标号的资源集合的映射
 std::unordered_map<std::string, ResourceIds> available_resources_;
};

class SchedulingResources {
public:
 // 默认构造函数和指定总量的构造函数
 SchedulingResources();
 SchedulingResources(const ResourceSet &total);
 ~SchedulingResources();
 // 总量，负载量和可用量的 getter
 const ResourceSet &GetAvailableResources() const;
 const ResourceSet &GetLoadResources() const;
 const ResourceSet &GetTotalResources() const;
 // 负载量和可用量的 setter，总量在构造时确定
 void SetLoadResources(ResourceSet &&newset);
 void SetAvailableResources(ResourceSet &&newset);
 // 获取或者释放一组资源：在可用资源量上做增删
 void Release(const ResourceSet &resources);
 void Acquire(const ResourceSet &resources);
 // 动态调整节点资源总量：更新某种种类的资源总量 or 删除某种种类资源
 void UpdateResource(const std::string &resource_name, int64_t capacity);
 void DeleteResource(const std::string &resource_name);
 std::string DebugString() const;
private:
 ResourceSet resources_total_;
 ResourceSet resources_available_;
 ResourceSet resources_load_;
};

void SchedulingResources::UpdateResource(const std::string &resource_name,
                                        int64_t capacity) {
 const FractionalResourceQuantity new_capacity = FractionalResourceQuantity(capacity);
 const FractionalResourceQuantity &current_capacity =
     resources_total_.GetResource(resource_name);
 if (current_capacity > 0) {
   // 如果该类资源存在，则将其总容量以及可用量进行相应更新
   const FractionalResourceQuantity capacity_difference =
       new_capacity - current_capacity;
   const FractionalResourceQuantity &current_available_capacity =
       resources_available_.GetResource(resource_name);
   FractionalResourceQuantity new_available_capacity =
       current_available_capacity + capacity_difference;
   if (new_available_capacity < 0) {
     new_available_capacity = 0;
  }
   resources_total_.AddOrUpdateResource(resource_name, new_capacity);
   resources_available_.AddOrUpdateResource(resource_name, new_available_capacity);
} else {
   // 如果不存在，则直接添加
   resources_total_.AddOrUpdateResource(resource_name, new_capacity);
   resources_available_.AddOrUpdateResource(resource_name, new_capacity);
}
}

class SchedulingPolicy {
public:
 // 构造函数：拿到本节点所有任务的引用
 SchedulingPolicy(const SchedulingQueue &scheduling_queue);
 // 根据集群资源分布以及任务资源需求做调度决策，返回任务与其所调度到的节点的集合
 std::unordered_map<TaskID, ClientID> Schedule(
     std::unordered_map<ClientID, SchedulingResources> &cluster_resources,
     const ClientID &local_client_id);
 // 在感知到新节点加入或者本地资源动态调整之后，对原先不可放置的任务进行尝试，并且匀出
 // 至多一个 READY 的任务到新节点。
 std::vector<TaskID> SpillOver(SchedulingResources &remote_scheduling_resources) const;
 virtual ~SchedulingPolicy();
private:
 // 任务队列的引用
 const SchedulingQueue &scheduling_queue_;
 /// 一个随机种子
 std::mt19937_64 gen_;
};

for task in placeable_tasks_list:
 clients = find_all_available_resources_statisfied_clients() # available not include load
 if cliens.is_not_empty():
   decision[task] = random_select_one(clients)
 else:
   clients = find_all_total_resources_statisfied_clients() # node whole resource
   if clients.is_not_empty():
     decision[task] = random_select_one(clients)
return decision

for (const auto &t : scheduling_queue_.GetTasks(TaskState::PLACEABLE)) {
   const auto &spec = t.GetTaskSpecification();
   const auto &resource_demand = spec.GetRequiredPlacementResources();
   const TaskID &task_id = spec.TaskId();
   std::vector<ClientID> client_keys;
   for (const auto &client_resource_pair : cluster_resources) {
     ClientID node_client_id = client_resource_pair.first;
     const auto &node_resources = client_resource_pair.second;
     ResourceSet available_node_resources = ResourceSet(node_resources.GetAvailableResources());
     // 1. 获取节点真正可用资源（resources_available_ - resources_load_）。
     available_node_resources.SubtractResources(node_resources.GetLoadResources());
     // 检测资源约束
     if (resource_demand.IsSubset(available_node_resources)) {
       client_keys.push_back(node_client_id);
    }
  }
   if (!client_keys.empty()) {
     // 随机选一个 index
     std::uniform_int_distribution<int> distribution(0, client_keys.size() - 1);
     int client_key_index = distribution(gen_);
     const ClientID &dst_client_id = client_keys[client_key_index];
     decision[task_id] = dst_client_id;
     // 2. 更新对应节点负载资源
     ResourceSet new_load(cluster_resources[dst_client_id].GetLoadResources());
     new_load.AddResources(resource_demand);
     cluster_resources[dst_client_id].SetLoadResources(std::move(new_load));
  }
...

def spill_over(remote_scheduling_resources):
 decision = []
 new_load = ResourceSet()
 # 尝试原先不可放置的任务
for task in infeasible_task_list:
   if task.required_res.is_subset_of(remote_scheduling_resources):
     decision.append(task.id)
     new_load.add(task.required_res)
 # 匀出至多一个 ready 任务
 for task in ready_task_list:
   if task.required_res.is_subset_of(remote_scheduling_resources):
     decision.append(task.id)
     new_load.add(task.required_res)
     break
 # 设置其节点资源负载
 remote_scheduling_resources.set_load(new_load)
 return decision