Repository: dask/distributed

Issue:
Remove `Scheduler.reschedule` from public API
I don't know if `Scheduler.reschedule` is even considered public API. There's no client method for it. It does show up in [Scheduler API docs](https://distributed.dask.org/en/stable/scheduling-state.html#distributed.scheduler.Scheduler.reschedule).

But `Scheduler.reschedule` presents misleading functionality. It looks like it can be used in general, in any circumstances, to reschedule any task: https://github.com/dask/distributed/blob/09837312ccb3a9f1a14ea068ba5825963fec82cc/distributed/scheduler.py#L7266-L7269

Reading the implementation, this function should actually only be used in one extremely specific circumstance: when the task has already been released in some way on the worker it's assigned to—either via cancellation or a `Reschedule` exception—and you are certain the worker will not send any further updates about the task to the scheduler.

Calling `reschedule` outside of that case can cause a task to be executing on multiple workers at once, and the scheduler only knows about it on one.

This inconsistency creates the need for "unexpected worker completed task" logic, and strange and un-tested transitions like `transition_waiting_memory`, `transition_no_worker_memory`, `transition_queued_memory`.

`reschedule` is really just a handler for the worker informing the scheduler of a `Reschedule` exception being raised. I'm not even 100% sure that work stealing's [use of it](https://github.com/dask/distributed/blob/09837312ccb3a9f1a14ea068ba5825963fec82cc/distributed/stealing.py#L363) is appropriate.

We should note in the docstring that this shouldn't be called directly.

cc @crusaderky @fjetter


Discussion:

Resolving patch:
diff --git a/distributed/scheduler.py b/distributed/scheduler.py
--- a/distributed/scheduler.py
+++ b/distributed/scheduler.py
@@ -3593,7 +3593,7 @@ def __init__(
             "release-worker-data": self.release_worker_data,
             "add-keys": self.add_keys,
             "long-running": self.handle_long_running,
-            "reschedule": self.reschedule,
+            "reschedule": self._reschedule,
             "keep-alive": lambda *args, **kwargs: None,
             "log-event": self.log_worker_event,
             "worker-status-change": self.handle_worker_status_change,
@@ -7263,13 +7263,15 @@ async def get_story(self, keys_or_stimuli: Iterable[str]) -> list[tuple]:
 
     transition_story = story
 
-    def reschedule(
+    def _reschedule(
         self, key: str, worker: str | None = None, *, stimulus_id: str
     ) -> None:
-        """Reschedule a task
+        """Reschedule a task.
 
-        Things may have shifted and this task may now be better suited to run
-        elsewhere
+        This function should only be used when the task has already been released in
+        some way on the worker it's assigned to — either via cancellation or a
+        Reschedule exception — and you are certain the worker will not send any further
+        updates about the task to the scheduler.
         """
         try:
             ts = self.tasks[key]
diff --git a/distributed/stealing.py b/distributed/stealing.py
--- a/distributed/stealing.py
+++ b/distributed/stealing.py
@@ -360,7 +360,7 @@ async def move_task_confirm(
                         *_log_msg,
                     )
                 )
-                self.scheduler.reschedule(key, stimulus_id=stimulus_id)
+                self.scheduler._reschedule(key, stimulus_id=stimulus_id)
             # Victim had already started execution
             elif state in _WORKER_STATE_REJECT:
                 self.log(("already-computing", *_log_msg))


Test patch:
diff --git a/distributed/tests/test_reschedule.py b/distributed/tests/test_reschedule.py
--- a/distributed/tests/test_reschedule.py
+++ b/distributed/tests/test_reschedule.py
@@ -1,4 +1,4 @@
-"""Tests for tasks raising the Reschedule exception and Scheduler.reschedule().
+"""Tests for tasks raising the Reschedule exception and Scheduler._reschedule().
 
 Note that this functionality is also used by work stealing;
 see test_steal.py for additional tests.
@@ -21,42 +21,10 @@
 )
 
 
-@gen_cluster(
-    client=True,
-    nthreads=[("", 1)] * 2,
-    config={
-        "distributed.scheduler.work-stealing": False,
-        # Difficult to get many tasks in processing with scheduler-side queuing
-        "distributed.scheduler.worker-saturation": float("inf"),
-    },
-)
-async def test_scheduler_reschedule(c, s, a, b):
-    xs = c.map(slowinc, range(100), key="x", delay=0.1)
-    while not a.state.tasks or not b.state.tasks:
-        await asyncio.sleep(0.01)
-    assert len(a.state.tasks) == len(b.state.tasks) == 50
-
-    ys = c.map(slowinc, range(100), key="y", delay=0.1, workers=[a.address])
-    while len(a.state.tasks) != 150:
-        await asyncio.sleep(0.01)
-
-    # Reschedule the 50 xs that are processing on a
-    for x in xs:
-        if s.tasks[x.key].processing_on is s.workers[a.address]:
-            s.reschedule(x.key, stimulus_id="test")
-
-    # Wait for at least some of the 50 xs that had been scheduled on a to move to b.
-    # This happens because you have 100 ys processing on a and 50 xs processing on b,
-    # so the scheduler will prefer b for the rescheduled tasks to obtain more equal
-    # balancing.
-    while len(a.state.tasks) == 150 or len(b.state.tasks) <= 50:
-        await asyncio.sleep(0.01)
-
-
 @gen_cluster()
 async def test_scheduler_reschedule_warns(s, a, b):
     with captured_logger("distributed.scheduler") as sched:
-        s.reschedule(key="__this-key-does-not-exist__", stimulus_id="test")
+        s._reschedule(key="__this-key-does-not-exist__", stimulus_id="test")
 
     assert "not found on the scheduler" in sched.getvalue()
     assert "Aborting reschedule" in sched.getvalue()
diff --git a/distributed/tests/test_steal.py b/distributed/tests/test_steal.py
--- a/distributed/tests/test_steal.py
+++ b/distributed/tests/test_steal.py
@@ -1186,7 +1186,7 @@ def block(x, event):
 
     steal.move_task_request(victim_ts, wsA, wsB)
 
-    s.reschedule(victim_key, stimulus_id="test")
+    s._reschedule(victim_key, stimulus_id="test")
     await event.set()
     await c.gather(futs1)
 
@@ -1321,7 +1321,7 @@ async def test_reschedule_concurrent_requests_deadlock(c, s, *workers):
     steal.move_task_request(victim_ts, wsA, wsB)
 
     s.set_restrictions(worker={victim_key: [wsB.address]})
-    s.reschedule(victim_key, stimulus_id="test")
+    s._reschedule(victim_key, stimulus_id="test")
     assert wsB == victim_ts.processing_on
     # move_task_request is not responsible for respecting worker restrictions
     steal.move_task_request(victim_ts, wsB, wsC)