Memory scheduling 2 DFS with Priority (ray-project#3)

Implemented basic DFS priority assignment and task blocking and preemption.

Currently segfaults due to object already existing during task retry.

Author: jaewan

benchmarks/single_node/test_single_node.py

@@ -19,6 +19,7 @@
 def assert_no_leaks():
     total = ray.cluster_resources()
     current = ray.available_resources()
+    print(total, current)
     total.pop("memory")
     total.pop("object_store_memory")
     current.pop("memory")
@@ -137,7 +138,8 @@ def test_large_object():
     assert big_obj[-1] == 0
 
 
-ray.init(address="auto")
+#ray.init(address="auto")
+ray.init()
 
 args_start = perf_counter()
 test_many_args()

cpp/src/ray/runtime/task/local_mode_task_submitter.cc

@@ -49,7 +49,7 @@ ObjectID LocalModeTaskSubmitter::Submit(InvocationSpec &invocation,
                             local_mode_ray_tuntime_.GetCurrentTaskId(), 0,
                             local_mode_ray_tuntime_.GetCurrentTaskId(), address, 1,
                             required_resources, required_placement_resources,
-                            std::make_pair(PlacementGroupID::Nil(), -1), true, "");
+                            std::make_pair(PlacementGroupID::Nil(), -1), true, "", Priority());
   if (invocation.task_type == TaskType::NORMAL_TASK) {
   } else if (invocation.task_type == TaskType::ACTOR_CREATION_TASK) {
     invocation.actor_id = local_mode_ray_tuntime_.GetNextActorID();

dashboard/client/package-lock.json

@@ -0,0 +1,67 @@
+import ray
+import time
+import sys
+import argparse
+import csv
+import numpy as np
+from time import perf_counter
+from time import perf_counter
+
+####################
+## Argument Parse ##
+####################
+parser = argparse.ArgumentParser()
+parser.add_argument('--WORKING_SET_RATIO', '-w', type=int, default=2)
+parser.add_argument('--OBJECT_STORE_SIZE', '-o', type=int, default=1_000_000_000)
+parser.add_argument('--OBJECT_SIZE', '-os', type=int, default=200_000_000)
+parser.add_argument('--NUM_STAGES', '-ns', type=int, default=1)
+args = parser.parse_args()
+params = vars(args)
+
+OBJECT_STORE_SIZE = params['OBJECT_STORE_SIZE'] 
+OBJECT_SIZE = params['OBJECT_SIZE'] 
+WORKING_SET_RATIO = params['WORKING_SET_RATIO']
+NUM_STAGES = params['NUM_STAGES']
+
+def test_ray_pipeline():
+    ray_pipeline_begin = perf_counter()
+
+    @ray.remote(num_cpus=1)
+    def consumer(obj_ref):
+        #args = ray.get(obj_ref)
+        return True
+
+    @ray.remote(num_cpus=1) 
+    def producer(): 
+        time.sleep(0.1)
+        for i in range(1000000):
+            pass
+        return np.zeros(OBJECT_SIZE // 8)
+        
+    num_fill_object_store = OBJECT_STORE_SIZE//OBJECT_SIZE 
+    produced_objs = [producer.remote()  for _ in range(WORKING_SET_RATIO*num_fill_object_store)]
+    refs = [[] for _ in range(NUM_STAGES)]
+
+    for obj in produced_objs:
+        refs[0].append(consumer.remote(obj))
+    '''
+    for stage in range(1, NUM_STAGES):
+        for r in refs[stage-1]:
+            refs[stage].append(consumer.remote(r))
+    '''
+    del produced_objs
+    #ray.get(refs[-1])
+    for r in refs[0]:
+        print(r)
+        ray.get(r)
+    '''
+    for ref in refs:
+        for r in ref:
+            ray.get(r)
+    '''
+    ray_pipeline_end = perf_counter()
+
+    return ray_pipeline_end - ray_pipeline_begin
+
+ray.init(object_store_memory=OBJECT_STORE_SIZE)
+print(test_ray_pipeline())

python/ray/tests/debug_block_tasks.py

@@ -374,7 +374,7 @@ RAY_CONFIG(int64_t, min_spilling_size, 100 * 1024 * 1024)
 
 /// If set to less than 1.0, Ray will start spilling objects when existing objects
 /// take more than this percentage of the available memory.
-RAY_CONFIG(float, object_spilling_threshold, 0.8)
+RAY_CONFIG(float, object_spilling_threshold, 1.0)
 
 /// Maximum number of objects that can be fused into a single file.
 RAY_CONFIG(int64_t, max_fused_object_count, 2000)

src/ray/common/ray_config_def.h

@@ -11,6 +11,16 @@ void Priority::extend(int64_t size) const {
   }
 }
 
+void Priority::SetFromParentPriority(Priority &parent, int s){
+  //param s id the last score to add
+  if(parent.score.size() == 1 && parent.score[0] == INT_MAX){
+	score[0] = s;
+  }else{
+	score = parent.score;
+    score.push_back(s);
+  }
+}
+
 bool Priority::operator<(const Priority &rhs) const {
   rhs.extend(score.size());
   extend(rhs.score.size());
@@ -25,6 +35,19 @@ bool Priority::operator<=(const Priority &rhs) const {
   return score <= rhs.score;
 }
 
+bool Priority::operator>(const Priority &rhs) const {
+  rhs.extend(score.size());
+  extend(rhs.score.size());
+
+  return score > rhs.score;
+}
+
+bool Priority::operator>=(const Priority &rhs) const {
+  rhs.extend(score.size());
+  extend(rhs.score.size());
+
+  return score >= rhs.score;
+}
 
 std::ostream &operator<<(std::ostream &os, const Priority &p) {
   os << "[ ";

src/ray/common/task/task_priority.cc

@@ -23,6 +23,8 @@ struct Priority {
 
   void extend(int64_t size) const;
 
+  void SetFromParentPriority(Priority &parent, int);
+
   bool operator==(const Priority &rhs) const {
     rhs.extend(score.size());
     extend(rhs.score.size());
@@ -37,6 +39,10 @@ struct Priority {
 
   bool operator<=(const Priority &rhs) const;
 
+  bool operator>(const Priority &rhs) const;
+
+  bool operator>=(const Priority &rhs) const;
+
   int GetScore(int64_t depth) const {
     extend(depth + 1);
     return score[depth];

src/ray/common/task/task_priority.h

@@ -105,6 +105,7 @@ class TaskSpecBuilder {
       const std::unordered_map<std::string, double> &required_placement_resources,
       const BundleID &bundle_id, bool placement_group_capture_child_tasks,
       const std::string &debugger_breakpoint,
+      const Priority &priority,
       const std::string &serialized_runtime_env = "{}",
       const std::vector<std::string> &runtime_env_uris = {},
       const std::string &concurrency_group_name = "") {
@@ -129,6 +130,12 @@ class TaskSpecBuilder {
         placement_group_capture_child_tasks);
     message_->set_debugger_breakpoint(debugger_breakpoint);
     message_->mutable_runtime_env()->set_serialized_runtime_env(serialized_runtime_env);
+
+    auto pri = message_->mutable_priority();
+    for (auto &s : priority.score) {
+      pri->Add(s);
+    }
+
     for (const std::string &uri : runtime_env_uris) {
       message_->mutable_runtime_env()->add_uris(uri);
     }

src/ray/common/task/task_util.h

@@ -40,15 +40,19 @@ void BuildCommonTaskSpec(
     const std::unordered_map<std::string, double> &required_resources,
     const std::unordered_map<std::string, double> &required_placement_resources,
     const BundleID &bundle_id, bool placement_group_capture_child_tasks,
-    const std::string debugger_breakpoint, const std::string &serialized_runtime_env,
+    const std::string debugger_breakpoint,
+    const Priority &priority,
+    const std::string &serialized_runtime_env,
     const std::vector<std::string> &runtime_env_uris,
     const std::string &concurrency_group_name = "") {
   // Build common task spec.
   builder.SetCommonTaskSpec(
       task_id, name, function.GetLanguage(), function.GetFunctionDescriptor(), job_id,
       current_task_id, task_index, caller_id, address, num_returns, required_resources,
       required_placement_resources, bundle_id, placement_group_capture_child_tasks,
-      debugger_breakpoint, serialized_runtime_env, runtime_env_uris,
+      debugger_breakpoint,
+      priority,
+      serialized_runtime_env, runtime_env_uris,
       concurrency_group_name);
   // Set task arguments.
   for (const auto &arg : args) {
@@ -681,7 +685,7 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
       std::move(lease_policy), memory_store_, task_manager_, local_raylet_id,
       RayConfig::instance().worker_lease_timeout_milliseconds(), actor_creator_,
       /*get_task_priority=*/[](const TaskSpecification &spec) {
-        return Priority();
+	    return spec.GetPriority();
       },
       RayConfig::instance().max_tasks_in_flight_per_worker(),
       boost::asio::steady_timer(io_service_));
@@ -1671,9 +1675,11 @@ std::vector<rpc::ObjectReference> CoreWorker::SubmitTask(
                       rpc_address_, function, args, task_options.num_returns,
                       constrained_resources, required_resources, placement_options,
                       placement_group_capture_child_tasks, debugger_breakpoint,
+                      Priority(),
                       task_options.serialized_runtime_env, task_options.runtime_env_uris);
   builder.SetNormalTaskSpec(max_retries, retry_exceptions);
   TaskSpecification task_spec = builder.Build();
+  //priority = task_manager_->GenerateTaskPriority(task_spec);
   RAY_LOG(DEBUG) << "Submit task " << task_spec.DebugString();
   std::vector<rpc::ObjectReference> returned_refs;
   if (options_.is_local_mode) {
@@ -1683,6 +1689,7 @@ std::vector<rpc::ObjectReference> CoreWorker::SubmitTask(
                                                   CurrentCallSite(), max_retries);
     io_service_.post(
         [this, task_spec]() {
+		//(Jae) This is the reason why tasks are not placed with priority
           RAY_UNUSED(direct_task_submitter_->SubmitTask(task_spec));
         },
         "CoreWorker.SubmitTask");
@@ -1727,6 +1734,7 @@ Status CoreWorker::CreateActor(const RayFunction &function,
                       new_placement_resources, actor_creation_options.placement_options,
                       actor_creation_options.placement_group_capture_child_tasks,
                       "", /* debugger_breakpoint */
+                      Priority(),
                       actor_creation_options.serialized_runtime_env,
                       actor_creation_options.runtime_env_uris);
 
@@ -1911,6 +1919,7 @@ std::vector<rpc::ObjectReference> CoreWorker::SubmitActorTask(
                       required_resources, std::make_pair(PlacementGroupID::Nil(), -1),
                       true, /* placement_group_capture_child_tasks */
                       "",   /* debugger_breakpoint */
+                      Priority(),
                       "{}", /* serialized_runtime_env */
                       {},   /* runtime_env_uris */
                       task_options.concurrency_group_name);
@@ -1923,6 +1932,7 @@ std::vector<rpc::ObjectReference> CoreWorker::SubmitActorTask(
 
   // Submit task.
   TaskSpecification task_spec = builder.Build();
+  std::vector<ObjectID> task_deps;
   std::vector<rpc::ObjectReference> returned_refs;
   if (options_.is_local_mode) {
     returned_refs = ExecuteTaskLocalMode(task_spec, actor_id);

src/ray/core_worker/core_worker.cc

@@ -262,6 +262,20 @@ void ReferenceCounter::RemoveLocalReference(const ObjectID &object_id,
   }
 }
 
+Priority& ReferenceCounter::GetObjectPriority(const ObjectID &object_id){
+  absl::MutexLock lock(&mutex_);
+  auto it = object_id_priority_.find(object_id);
+  RAY_CHECK(it != object_id_priority_.end()) << "Object priority not found " << object_id;
+  return it->second;
+}
+
+void ReferenceCounter::UpdateObjectPriority(
+		const ObjectID &object_id,
+		const Priority &priority){
+  absl::MutexLock lock(&mutex_);
+  object_id_priority_[object_id] =  priority;
+}
+
 void ReferenceCounter::UpdateSubmittedTaskReferences(
     const std::vector<ObjectID> &argument_ids_to_add,
     const std::vector<ObjectID> &argument_ids_to_remove, std::vector<ObjectID> *deleted) {