[core] Periodically GC metadata for streaming generators #43772

stephanie-wang · 2024-03-07T06:59:37Z

Why are these changes needed?

Less invasive fix than #43584. Changes:

When generator goes out of scope, add it to a list of streaming generator tasks that we scan periodically
For each generator task, check if the task and streaming metadata can be removed. It can be removed if the generator task and all generated return refs have gone out of scope.
Fixes an existing potential leak where task completes after the generator ref and returned refs have gone out of scope, by deleting the task metadata with the stream metadata.

#43584 is probably better long-term as it refactors stream metadata to be GCed through the normal ref counting path, and lineage can be GCed eagerly. However, this fix is safer.

Related issue number

Closes #39151.

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run scripts/format.sh to lint the changes in this PR.
I've included any doc changes needed for https://docs.ray.io/en/master/.
- I've added any new APIs to the API Reference. For example, if I added a
  method in Tune, I've added it in doc/source/tune/api/ under the
  corresponding .rst file.
I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
Testing Strategy
- Unit tests
- Release tests
- This PR is not tested :(

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

rkooo567 · 2024-03-08T01:08:26Z

src/ray/core_worker/task_manager.cc

@@ -238,7 +238,7 @@ std::vector<rpc::ObjectReference> TaskManager::AddPendingTask(
  // The language frontend is responsible for calling DeleteObjectRefStream.
  if (spec.IsStreamingGenerator()) {
    const auto generator_id = spec.ReturnId(0);
-    RAY_LOG(DEBUG) << "Create an object ref stream of an id " << generator_id;
+    RAY_LOG(INFO) << "Create an object ref stream of an id " << generator_id;


is this debug -> info intended? we should revert?

Yes, it has been reverted.

rkooo567 · 2024-03-08T01:11:51Z

src/ray/core_worker/task_manager.h

@@ -410,7 +423,9 @@ class TaskManager : public TaskFinisherInterface, public TaskResubmissionInterfa
      ABSL_LOCKS_EXCLUDED(mu_);

  /// Return True if there's no more object to read. False otherwise.
-  bool IsFinished(const ObjectID &generator_id) const ABSL_LOCKS_EXCLUDED(mu_);
+  bool StreamingGeneratorIsFinished(const ObjectID &generator_id,


Can you add docstring for num_objects_generated?

src/ray/core_worker/reference_count.cc

rkooo567 · 2024-03-08T01:20:10Z

src/ray/core_worker/task_manager.cc

          for (const auto &return_id_info : reply.streaming_generator_return_ids()) {
            if (return_id_info.is_plasma_object()) {
+              // NOTE(swang): It is possible that the dynamically returned refs


is this a separate issue?

Existing issue that I added a fix for in this PR.

rkooo567 · 2024-03-08T01:22:57Z

src/ray/core_worker/core_worker.cc

+      continue;
+    }
+
+    bool can_gc = reference_counter_->CheckGeneratorRefsOutOfScope(generator_id,


can we remove all unconsumed objects here (and add a test)? This seems like a regression (where unconsumed objects are not deleted if lineage is alive).

E.g., when delete hapapens, all unconsumed objects are deleted (same behavior as now). And only the stream metadata is cleaned up after a delay.

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

rkooo567

LGTM!

rkooo567 · 2024-03-08T07:40:33Z

python/ray/tests/test_streaming_generator_2.py

+    assert_no_leak()
+
+    # Test that when the generator task stays in the in-scope lineage, we still
+    # clean up the unconsumed objects' values. The lineage (task and stream


nice tests!

src/ray/core_worker/core_worker.cc

rkooo567 · 2024-03-08T07:45:06Z

src/ray/core_worker/core_worker.cc

@@ -641,6 +641,11 @@ CoreWorker::CoreWorker(const CoreWorkerOptions &options, const WorkerID &worker_
      RayConfig::instance().metrics_report_interval_ms() / 2,
      "CoreWorker.RecordMetrics");

+  periodical_runner_.RunFnPeriodically(
+      [this] { TryDeleteObjectRefStreams(); },
+      RayConfig::instance().local_gc_min_interval_s() * 1000,


Looks like the wrong value? (it is already a second, but you multiply 1000).

Besides, do you think 10 seconds is good enough? Feel like it needs to be a little more frequent

It's to convert to ms?

Should be fine because we trigger it once immediately now.

rkooo567 · 2024-03-08T07:47:15Z

src/ray/core_worker/task_manager.cc

@@ -57,6 +57,38 @@ absl::flat_hash_set<ObjectID> ObjectRefStream::GetItemsUnconsumed() const {
  return result;
 }

+std::vector<ObjectID> ObjectRefStream::PopUnconsumedItems() {


is it same as GetItemsUnconsumed? If so can we remove other APi?

This method pops the items, not just get.

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

stephanie-wang added 3 commits March 6, 2024 21:26

Delete object ref streams periodically and after task completes

19efd2b

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

fix tests

00dacb3

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

lint

6f25782

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

stephanie-wang assigned jjyao and rkooo567 Mar 7, 2024

stephanie-wang added 2 commits March 7, 2024 10:37

test

12e2725

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

restore logging

a9d1d62

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

rkooo567 reviewed Mar 8, 2024

View reviewed changes

stephanie-wang added 6 commits March 7, 2024 21:39

test

b1be839

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

GC unconsumed refs

abdd963

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

unit test

cd8dafb

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

fix?

78b463c

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

simpler

f2901f5

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

lint

b4c9512

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

rkooo567 approved these changes Mar 8, 2024

View reviewed changes

clean

1fc4e72

Signed-off-by: Stephanie Wang <swang@cs.berkeley.edu>

stephanie-wang merged commit 9440cb0 into ray-project:master Mar 8, 2024
9 checks passed

stephanie-wang deleted the fix-39151-2 branch March 8, 2024 21:50

stephanie-wang mentioned this pull request Mar 8, 2024

[core] Release streaming generator task metadata based on ref counter #43584

Closed

8 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[core] Periodically GC metadata for streaming generators #43772

[core] Periodically GC metadata for streaming generators #43772

stephanie-wang commented Mar 7, 2024 •

edited

rkooo567 Mar 8, 2024

stephanie-wang Mar 8, 2024

rkooo567 Mar 8, 2024

rkooo567 Mar 8, 2024

stephanie-wang Mar 8, 2024

rkooo567 Mar 8, 2024

rkooo567 left a comment

rkooo567 Mar 8, 2024

rkooo567 Mar 8, 2024

stephanie-wang Mar 8, 2024

rkooo567 Mar 8, 2024

stephanie-wang Mar 8, 2024

[core] Periodically GC metadata for streaming generators #43772

[core] Periodically GC metadata for streaming generators #43772

Conversation

stephanie-wang commented Mar 7, 2024 • edited

Why are these changes needed?

Related issue number

Checks

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

rkooo567 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stephanie-wang commented Mar 7, 2024 •

edited