diff --git a/src/ray/raylet/lineage_cache.cc b/src/ray/raylet/lineage_cache.cc index 16f6da5a0..07193d5f9 100644 --- a/src/ray/raylet/lineage_cache.cc +++ b/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 &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 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)); diff --git a/src/ray/raylet/lineage_cache.h b/src/ray/raylet/lineage_cache.h index 0a104ac97..402d49a67 100644 --- a/src/ray/raylet/lineage_cache.h +++ b/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 &seen) const; /// The client ID, used to request notifications for specific tasks. /// TODO(swang): Move the ClientID into the generic Table implementation. diff --git a/test/stress_tests.py b/test/stress_tests.py index 2875d5b88..37ee465c6 100644 --- a/test/stress_tests.py +++ b/test/stress_tests.py @@ -55,9 +55,6 @@ def test_submitting_tasks(ray_start_combination): 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 test_dependencies(ray_start_combination): 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):