[Core] Locality-aware leasing: Milestone 1 - Owned refs, pinned location

[clarkzinzow]

Implements the first milestone of locality-aware leasing, where the locality cost only includes owned refs, and only considers pinned object locations.

Related issue number

Closes #12813

Checks

• 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 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 :(

TODOs:

• Validate design and implementation.
• Ensure proper concurrency control.
• Finish updating existing tests.
• Add tests for LeasePolicy.
• Add a feature flag that gates the locality-aware leasing functionality.
• Benchmark.

[clarkzinzow]

This clunkiness around the raylet_address API/semantics will be gone once locality-aware spillback is implemented.

[stephanie-wang]

This looks great! I left some minor suggestions for cleanups.

Also, could we add tests in C++ and Python to check the data locality behavior? For C++, it'd be great to check both the LessorPicker behavior as well as the CoreWorkerDirectTaskSubmitter (that it only uses the locality policy on the first lease request). Let me know if you need any pointers on that.

[clarkzinzow]

@stephanie-wang Thanks for the feedback!

Also, could we add tests in C++ and Python to check the data locality behavior?

Yep, I was waiting to validate the general design and implementation before writing tests, I'll get on that next!

[clarkzinzow]

Microbenchmark results:

No obvious performance regressions, but we should watch out for regressions on the microbenchmarks that contain a high expected number of lease requests, e.g. the sync task microbenchmarks (I got 26% lower throughput on the single client tasks sync per second microbenchmark in a different run). If those start to dip, we might want to provide extra logic/heuristics to short-circuit the locality consideration for lightweight tasks.

It also might be worth it to (eventually) add a set of data-intensive microbenchmarks, e.g. tasks with a few MiBs of task dependencies.

Before this PR:

single client get calls per second 34568.79 +- 433.35
single client put calls per second 26084.66 +- 444.76
multi client put calls per second 47940.87 +- 9045.39
single client get calls (Plasma Store) per second 6076.26 +- 227.36
single client put calls (Plasma Store) per second 2063.29 +- 12.54
multi client put calls (Plasma Store) per second 3220.9 +- 49.3
single client put gigabytes per second 21.01 +- 11.51
multi client put gigabytes per second 6.33 +- 0.52
single client tasks sync per second 694.72 +- 6.14
single client tasks async per second 5963.91 +- 107.73
multi client tasks async per second 6740.55 +- 123.5
1:1 actor calls sync per second 1121.29 +- 17.09
1:1 actor calls async per second 4098.41 +- 135.86
1:1 actor calls concurrent per second 3876.48 +- 68.63
1:n actor calls async per second 7819.1 +- 72.62
n:n actor calls async per second 8260.84 +- 111.55
n:n actor calls with arg async per second 3218.98 +- 155.9
1:1 async-actor calls sync per second 785.03 +- 130.74
1:1 async-actor calls async per second 2773.56 +- 27.59
1:1 async-actor calls with args async per second 1865.92 +- 16.58
1:n async-actor calls async per second 5297.05 +- 31.62
n:n async-actor calls async per second 5651.63 +- 91.19

After this PR:

single client get calls per second 37885.74 +- 755.56
single client put calls per second 25117.27 +- 431.47
multi client put calls per second 55474.79 +- 963.0
single client get calls (Plasma Store) per second 6165.02 +- 127.57
single client put calls (Plasma Store) per second 2043.24 +- 25.41
multi client put calls (Plasma Store) per second 3118.0 +- 4.0
single client put gigabytes per second 19.89 +- 11.1
multi client put gigabytes per second 6.31 +- 0.47
single client tasks sync per second 644.91 +- 9.36
single client tasks async per second 5972.95 +- 155.72
multi client tasks async per second 6560.8 +- 335.35
1:1 actor calls sync per second 1091.83 +- 133.3
1:1 actor calls async per second 4300.17 +- 267.23
1:1 actor calls concurrent per second 3417.22 +- 222.95
1:n actor calls async per second 7861.75 +- 213.67
n:n actor calls async per second 8116.69 +- 181.7
n:n actor calls with arg async per second 3353.72 +- 53.48
1:1 async-actor calls sync per second 864.3 +- 10.82
1:1 async-actor calls async per second 2753.3 +- 38.87
1:1 async-actor calls with args async per second 1838.37 +- 25.53
1:n async-actor calls async per second 5240.66 +- 96.54
n:n async-actor calls async per second 5467.0 +- 50.04

[rkooo567]

Do you mind doing this in this PR? I'd like to see the difference.

It also might be worth it to (eventually) add a set of data-intensive microbenchmarks, e.g. tasks with a few MiBs of task dependencies.

[clarkzinzow]

Do you mind doing this in this PR? I'd like to see the difference.

It also might be worth it to (eventually) add a set of data-intensive microbenchmarks, e.g. tasks with a few MiBs of task dependencies.

@rkooo567 I wouldn't expect to see a difference unless the microbenchmarks were running on a multi-node cluster, where object transfer between nodes across the network starts to become a performance factor. I can add a few data-intensive microbenchmarks, but I don't expect them to be very illuminating for locality-aware leasing in the local setting!

[rkooo567]

@clarkzinzow You can simulate the multi node cluster with the cluster_utils module. I guess we can create large objects and do some random scheduling to see the difference? (maybe it won't be big difference, but I expect to see some performance improvement)

[clarkzinzow]

@rkooo567 Oh cool, I didn't know that ray clusterbenchmark was a thing ah, looks like it's no longer exposed.

[clarkzinzow]

Cluster microbenchmark results:

Slightly faster on sync data-intensive tasks over 3 workers (head + 2 non-head workers), although these are some very naive cluster microbenchmarks. Each task adds two 40 x 1024 x 1024 ndarrays of 64-bit floats, i.e. two 320 MiB arguments and a 320 MiB return value. I'll try to think of some better data-intensive microbenchmarks that are sure to (1) not reuse the same worker lease, and (2) schedule tasks onto multiple workers.

Locality-aware leasing turned off:

single client data-intensive tasks sync mean latency 2.91 +- 0.92
single client data-intensive tasks async mean latency 29.77 +- 3.58

Locality-aware leasing turned on:

single client data-intensive tasks sync mean latency 2.49 +- 0.89
single client data-intensive tasks async mean latency 31.32 +- 0.45

[stephanie-wang]

Looks great, thanks!

[rkooo567]

Btw, what's going to happen if the locality data is nullopt?

[clarkzinzow]

This inner for loop will only be reached if the prior if (...) condition is true, and that will only happen if locality_data contains a value. I've seen this optional pattern used and promoted elsewhere, but do you think this is too implicit? We could always opt for the more verbose

auto locality_data = locality_data_provider_->GetLocalityData(object_id);
if (locality_data.has_value()) {
    // ....
}

[rkooo567]

Also one quick question. When's the case that the owned refs are pinned on a non-local node?

[clarkzinzow]

Also one quick question. When's the case that the owned refs are pinned on a non-local node?

An example would be the return object of a task A that was submitted on node 1 but ultimately executed on node 2. I.e., task A is initially submitted to the local raylet (node 1), the local raylet responds by telling the submitter to retry the scheduling request at node 2, the raylet on node 2 accepts said scheduling request and executes the task on a node 2 worker, and when the node 2 worker completes task A the output is pinned on node 2 since its the first copy of said output object that has been created. But given that a worker on node 1 submitted the task, node 1 is seen as the owner of the return object ref. At least that's my current understanding of the protocol, lmk if I'm mistaken!

The appendix of the whitepaper actually walks through this example.