Ray源码解析之调度部分

src/ray/raylet/scheduling_policy.h
---------------------------------------------

namespace ray {

    /// \brief Implements a scheduling policy for the node manager.
    class SchedulingPolicy {
     public:
      /// \param scheduling_queue: reference to a scheduler queues object for access to
      ///        tasks.
      SchedulingPolicy(const SchedulingQueue &scheduling_queue);
    
      /// Perform a scheduling operation, given a set of cluster resources and
      /// producing a mapping of tasks to node managers.
      ///
      ///  \param cluster_resources: a set of cluster resources representing
      ///         configured and current resource capacity on each node.
      /// \return Scheduling decision, mapping tasks to node managers for placement.
      std::unordered_map<TaskID, ClientID, UniqueIDHasher> Schedule(
          const std::unordered_map<ClientID, SchedulingResources, UniqueIDHasher>
              &cluster_resources);
    
      virtual ~SchedulingPolicy();
    
     private:
      /// An immutable reference to the scheduling task queues.
      const SchedulingQueue &scheduling_queue_;
    };
}

scheduling_policy.cc
----------------------------------

namespace ray {
....
    std::unordered_map<TaskID, ClientID, UniqueIDHasher> SchedulingPolicy::Schedule(
        const std::unordered_map<ClientID, SchedulingResources, UniqueIDHasher>
            &cluster_resources) {
      static ClientID local_node_id = ClientID::nil();
      std::unordered_map<TaskID, ClientID, UniqueIDHasher> decision;
      // TODO(atumanov): consider all cluster resources.
      SchedulingResources resource_supply = cluster_resources.at(local_node_id);
      const auto &resource_supply_set = resource_supply.GetAvailableResources();
    
      // Iterate over running tasks, get their resource demand and try to schedule.
      for (const auto &t : scheduling_queue_.GetReadyTasks()) {
        // Get task's resource demand
        const auto &resource_demand = t.GetTaskSpecification().GetRequiredResources();
        bool task_feasible = resource_demand.IsSubset(resource_supply_set);
        if (task_feasible) {
          const TaskID &task_id = t.GetTaskSpecification().TaskId();
          decision[task_id] = local_node_id;
        }
      }
      return decision;
    }
....
}

scheduling_queue.h
--------------------------------

namespace ray {

class SchedulingQueue {
 public:
  SchedulingQueue() {}

  virtual ~SchedulingQueue() {}

  const std::list<Task> &GetWaitingTasks() const;
  const std::list<Task> &GetReadyTasks() const;
  const std::list<Task> &GetReadyMethods() const;
  const std::list<Task> &GetScheduledTasks() const;
  const std::list<Task> &GetRunningTasks() const;

  std::vector<Task> RemoveTasks(std::unordered_set<TaskID, UniqueIDHasher> tasks);


  void QueueWaitingTasks(const std::vector<Task> &tasks);
  void QueueReadyTasks(const std::vector<Task> &tasks);
  void QueueScheduledTasks(const std::vector<Task> &tasks);
  void QueueRunningTasks(const std::vector<Task> &tasks);

  /// Register an actor.
  ///
  /// \param actor_id The ID of the actor to register.
  /// \param actor_information Information about the actor.
  bool RegisterActor(ActorID actor_id, const ActorInformation &actor_information);

 private:
  std::list<Task> waiting_tasks_;
  std::list<Task> ready_tasks_;
  std::list<Task> scheduled_tasks_;
  std::list<Task> running_tasks_;
  /// The registry of known actors.
  std::unordered_map<ActorID, ActorInformation, UniqueIDHasher> actor_registry_;
};
}  // namespace ray

// Helper function to remove tasks in the given set of task_ids from a
// queue, and append them to the given vector removed_tasks.
void removeTasksFromQueue(std::list<Task> &queue,
                          std::unordered_set<TaskID, UniqueIDHasher> &task_ids,
                          std::vector<Task> &removed_tasks) {
  for (auto it = queue.begin(); it != queue.end();) {
    auto task_id = task_ids.find(it->GetTaskSpecification().TaskId());
    if (task_id != task_ids.end()) {
      task_ids.erase(task_id);
      removed_tasks.push_back(std::move(*it));
      it = queue.erase(it);
    } else {
      it++;
    }
  }
}

scheduling_resources.h
------------------------------------

namespace ray {

/// Resource availability status reports whether the resource requirement is
/// (1) infeasible, (2) feasible but currently unavailable, or (3) available.
typedef enum {
  kInfeasible,            ///< Cannot ever satisfy resource requirements.
  kResourcesUnavailable,  ///< Feasible, but not currently available.
  kFeasible               ///< Feasible and currently available.
} ResourceAvailabilityStatus;

class ResourceSet {
 public:
  ResourceSet();
  ResourceSet(const std::unordered_map<std::string, double> &resource_map);
  ~ResourceSet();

  bool operator==(const ResourceSet &rhs) const;

  bool IsEqual(const ResourceSet &other) const;

  bool IsSubset(const ResourceSet &other) const;

  bool IsSuperset(const ResourceSet &other) const;

  bool AddResource(const std::string &resource_name, double capacity);

  bool RemoveResource(const std::string &resource_name);

  bool AddResources(const ResourceSet &other);

  bool SubtractResources(const ResourceSet &other);

  // 返回指定资源的容量值（赋给value指向的值），如果资源有大于0的数量，且value不为空指针，则返回true
  bool GetResource(const std::string &resource_name, double *value) const;

 private:
  // 资源容量map
  std::unordered_map<std::string, double> resource_capacity_;
};

/// \class SchedulingResources
/// SchedulingResources 封装资源的状态和资源的计数。资源包括配置资源束容量和GPU分配图。
class SchedulingResources {
 public:
  SchedulingResources();
  SchedulingResources(const ResourceSet &total);
  ~SchedulingResources();
  // 检查一个资源集是否能被满足，有几种状态，(1) infeasible, (2) feasible but currently unavailable, or (3) available.
  ResourceAvailabilityStatus CheckResourcesSatisfied(ResourceSet &set) const;

  // 请求现在可用的资源集，返回一个不可变的资源集合
  const ResourceSet &GetAvailableResources() const;

  bool Release(const ResourceSet &resources);

  bool Acquire(const ResourceSet &resources);

 private:
  /// 静态资源配置
  ResourceSet resources_total_;
  /// 动态资源容量
  ResourceSet resources_available_;
  /// gpu_map - replace with ResourceMap (for generality).
};

}  // namespace ray

bool ResourceSet::operator==(const ResourceSet &rhs) const {
  return (this->IsSubset(rhs) && rhs.IsSubset(*this));
}

bool ResourceSet::IsSubset(const ResourceSet &other) const {
  // Check to make sure all keys of this are in other.
  for (const auto &resource_pair : resource_capacity_) {
    const auto &resource_name = resource_pair.first;
    const double lhs_quantity = resource_pair.second;
    double rhs_quantity = 0;
    if (!other.GetResource(resource_name, &rhs_quantity)) {
      // Resource not found in rhs, therefore lhs is not a subset of rhs.
      return false;
    }
    if (lhs_quantity > rhs_quantity) {
      // Resource found in rhs, but lhs capacity exceeds rhs capacity.
      return false;
    }
  }
  return true;
}