Enhance AMM docstrings
@@ -189,9 +189,14 @@ class ActiveMemoryManagerExtension:
pending_repl: set[WorkerState],
) -> WorkerState | None:
"""Choose a worker to acquire a new replica of an in-memory task among a set of
- candidates. If candidates is None, default to all workers in the cluster that do
- not hold a replica yet. The worker with the lowest memory usage (downstream of
- pending replications and drops) will be returned.
+ candidates. If candidates is None, default to all workers in the cluster.
+ Regardless, workers that either already hold a replica or are scheduled to
+ receive one at the end of this AMM iteration are not considered.
+
+ Returns
+ -------
+ The worker with the lowest memory usage (downstream of pending replications and
+ drops), or None if no eligible candidates are available.
"""
if ts.state != "memory":
return None
@@ -210,9 +215,15 @@ class ActiveMemoryManagerExtension:
pending_drop: set[WorkerState],
) -> WorkerState | None:
"""Choose a worker to drop its replica of an in-memory task among a set of
- candidates. If candidates is None, default to all workers in the cluster that
- hold a replica. The worker with the highest memory usage (downstream of pending
- replications and drops) will be returned.
+ candidates. If candidates is None, default to all workers in the cluster.
+ Regardless, workers that either do not hold a replica or are already scheduled
+ to drop theirs at the end of this AMM iteration are not considered.
+ This method also ensures that a key will not lose its last replica.
+
+ Returns
+ -------
+ The worker with the highest memory usage (downstream of pending replications and
+ drops), or None if no eligible candidates are available.
"""
if len(ts.who_has) - len(pending_drop) < 2:
return None
