[xray] Fix bug when counting a task's lineage size (#2600)

Author: stephanie-wang

src/ray/raylet/lineage_cache.cc

@@ -219,7 +219,12 @@ void LineageCache::AddReadyTask(const Task &task) {
   }
 }
 
-uint64_t LineageCache::CountUnsubscribedLineage(const TaskID &task_id) const {
+uint64_t LineageCache::CountUnsubscribedLineage(const TaskID &task_id,
+                                                std::unordered_set<TaskID> &seen) const {
+  if (seen.count(task_id) == 1) {
+    return 0;
+  }
+  seen.insert(task_id);
   if (subscribed_tasks_.count(task_id) == 1) {
     return 0;
   }
@@ -229,7 +234,7 @@ uint64_t LineageCache::CountUnsubscribedLineage(const TaskID &task_id) const {
   }
   uint64_t cnt = 1;
   for (const auto &parent_id : entry->GetParentTaskIds()) {
-    cnt += CountUnsubscribedLineage(parent_id);
+    cnt += CountUnsubscribedLineage(parent_id, seen);
   }
   return cnt;
 }
@@ -257,7 +262,9 @@ void LineageCache::RemoveWaitingTask(const TaskID &task_id) {
   // NOTE(swang): The number of entries in the uncommitted lineage also
   // includes local tasks that haven't been committed yet, not just remote
   // tasks, so this is an overestimate.
-  if (CountUnsubscribedLineage(task_id) > max_lineage_size_) {
+  std::unordered_set<TaskID> seen;
+  auto count = CountUnsubscribedLineage(task_id, seen);
+  if (count > max_lineage_size_) {
     // Since this task was in state WAITING, check that we were not
     // already subscribed to the task.
     RAY_CHECK(SubscribeTask(task_id));

src/ray/raylet/lineage_cache.h

@@ -248,8 +248,15 @@ class LineageCache {
   /// Unsubscribe from notifications for a task. Returns whether the operation
   /// was successful (whether we were subscribed).
   bool UnsubscribeTask(const TaskID &task_id);
-  /// Count the size of unsubscribed and uncommitted lineage
-  uint64_t CountUnsubscribedLineage(const TaskID &task_id) const;
+  /// Count the size of unsubscribed and uncommitted lineage of the given task
+  /// excluding the values that have already been visited.
+  ///
+  /// \param task_id The task whose lineage should be counted.
+  /// \param seen This set contains the keys of lineage entries counted so far,
+  /// so that we don't revisit those nodes.
+  /// \void The number of tasks that were counted.
+  uint64_t CountUnsubscribedLineage(const TaskID &task_id,
+                                    std::unordered_set<TaskID> &seen) const;
 
   /// The client ID, used to request notifications for specific tasks.
   /// TODO(swang): Move the ClientID into the generic Table implementation.

test/stress_tests.py

@@ -55,9 +55,6 @@ def f(x):
     assert ray.services.all_processes_alive()
 
 
-@pytest.mark.skipif(
-    os.environ.get("RAY_USE_XRAY") == "1",
-    reason="This test does not work with xray yet.")
 def test_dependencies(ray_start_combination):
     @ray.remote
     def f(x):
@@ -81,9 +78,6 @@ def g(*xs):
     assert ray.services.all_processes_alive()
 
 
-@pytest.mark.skipif(
-    os.environ.get("RAY_USE_XRAY") == "1",
-    reason="This test does not work with xray yet.")
 def test_submitting_many_tasks(ray_start_regular):
     @ray.remote
     def f(x):