dask / dask/distributed

Reimplement `replicate()` using the Active Memory Manager

Open
#6,578 9 comments 0 reactions 1 assignee Claimed by @crusaderky View on GitHub
memory
Dominant language
Python
Stars
1.7k
Forks
778
Avg merge
2h 50m
Merged PRs (30d)
3

Description

This issue is very tightly related to #4906.

`replicate()`, as well as its wrapper `scatter(..., broadcast=True)`, have several issues:
- It's fragile on a moving cluster:
- #3024
- #3990
- #1930
- It uses transfer machinery that is used exclusively by replicate and rebalance (see #4906)
- If you later lose a replica, there's nothing recreating it
- If you lose a worker in the middle of it, it will hang until it's restarted (if ever)

# Proposed design
Reimplement replicate() on top of the Active Memory Manager.
The command will just start an AMM policy for the involved keys.
The policy will track the keys and, every two seconds, create new replicas if there aren't enough.
Once the keys cease to exist, or the client calls `replicate(n=1)` on the same keys, the policy detaches itself from the AMM.

The number of desired replicas will be tracked through a `replicate: ` annotation on the involved keys.

### Side effects
- `replicate()` becomes non-blocking. It won't wait for replication to complete and won't even wait for keys to become in-memory.
- You don't need to have all target workers online when you invoke replicate. As a matter of fact, `replicate(n=inf)` will become the default option; it means that if at any point in the future a new worker joins the cluster, the key will be immediately replicated onto it.

### Optional additional feature
Alternatively, the client may also annotate the keys directly when building the graph:
```python
a = ...
with dask.annotate(replicate=2):
b = f(a)
b.compute(optimize_graph=False)
```
For this to work, when a new key lands on the scheduler, there must be machinery that parses the annotations, detects the `replicate` tag, and invokes Scheduler.replicate() under the hood.
This can be neatly implemented by a Scheduler plugin.

### Proposed contentious breaking chances
- Do not offer an option to make replicate blocking again
- Do not offer an option to specify a subset of workers to replicate to (this would make the annotation much simpler)
- The `replicate()` command will fail if the AMM is not enabled.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.