Transition queued->memory causes AssertionError

[crusaderky]

While trying to reproduce #7063, I came across a different error, this one with queueing enabled.
The below reproducer is NOT minimal - there is likely quite a bit of simplification possible.

@gen_cluster(client=True, nthreads=[("", 1)], config={"distributed.scheduler.worker-saturation": 1.5})
async def test_steal_rootish_while_retiring(c, s, a):
    """https://github.com/dask/distributed/issues/7063

    Note that this applies to both tasks that raise Reschedule as well as work stealing.
    """
    ev = Event()

    # Put a task in memory on a, which will be retired, and prevent b from acquiring
    # a replica. This will cause a to be stuck in closing_gracefully state until we
    # set b.block_gather_dep.
    m = c.submit(inc, 1, key="m", workers=[a.address])
    await wait(m)

    async with BlockedGatherDep(s.address, nthreads=1) as b:
        # Large number of tasks to make sure they're rootish
        futures = c.map(
            lambda i, ev: ev.wait(), range(10), ev=ev, key=[f"x-{i}" for i in range(10)]
        )

        while a.state.executing_count != 1 or b.state.executing_count != 1:
            await asyncio.sleep(0.01)

        assert s.is_rootish(s.tasks[futures[0].key])

        retire_task = asyncio.create_task(c.retire_workers([a.address]))
        # Wait until AMM sends AcquireReplicasEvent to b to move away m
        await b.in_gather_dep.wait()
        assert s.workers[a.address].status == Status.closing_gracefully

        # Steal any of the tasks on a
        steal_key = next(iter(a.state.executing)).key
        s.reschedule(steal_key, stimulus_id="steal")
        await ev.set()

        # The stolen task can now complete on the other worker
        await wait_for_state(steal_key, "memory", b)
        await wait_for_state(steal_key, "memory", s)

        # Let graceful retirement of a complete.
        # This in turn reschedules whatever tasks were still processing on a to b.
        b.block_gather_dep.set()
        await retire_task
        await wait(futures)

The test is green; however I read in the log:

  File "/home/crusaderky/github/distributed/distributed/scheduler.py", line 5284, in handle_task_finished
    r: tuple = self.stimulus_task_finished(
  File "/home/crusaderky/github/distributed/distributed/scheduler.py", line 4649, in stimulus_task_finished
    r: tuple = self._transition(
  File "/home/crusaderky/github/distributed/distributed/scheduler.py", line 1813, in _transition
    assert not args and not kwargs, (args, kwargs, start, finish)
AssertionError: ((), {'worker': 'tcp://127.0.0.1:45929', 'nbytes': 28, 'type': b'\x80\x04\x95\x15\x00\x00\x00\x00\x00\x00\x00\x8c\x08builtins\x94\x8c\x04bool\x94\x93\x94.', 'typename': 'bool', 'metadata': {}, 'thread': 139862053221952, 'startstops': ({'action': 'compute', 'start': 1666802403.9580944, 'stop': 1666802403.9590282},), 'status': 'OK'}, 'queued', 'memory')

What is happening:

1. steal_key is processing on a
2. steal_key is rescheduled, which causes the scheduler to send a free-keys message to a and put the task back in queue
3. before the free-keys message can reach a, steal_key finishes on a
4. steal_key transitions to memory on a, sending a TaskFinishedMsg to the scheduler.
5. a queued->memory transition happens which, I suspect, is otherwise untested.

This is timing-sensitive; if free-keys reached a before the task end, then steal_key would be cancelled and transition to forgotten without any messaging when it ends.

[gjoseph92]

I somewhat intentionally didn't add a queued->memory transition since it seemed like a messed-up case to me, and I didn't want to cargo-cult it in just because other transitions of the same form were there. Seems it is in fact a messed-up case, but a possible one.

So I guess we should just add the transition, which should be straightforward. Minimizing this test, on the other hand, might not be.

[crusaderky]

I can work on minimising the test

[crusaderky]

I'm working on this.
I think we have a much wider problem.

• transition_processing_memory (with unexpected worker),
• transition_waiting_memory, and
• transition_no_worker_memory

All implement the same logic to accept the task output from wherever it comes. However, the only case I can think of when this happens is when a cancellation request is already in flight but hasn't yet reached the worker, so the worker is going to lose the task immediately afterwards. I think they all have to be fixed to just disregard the transition.

CC @fjetter

[fjetter]

    # Steal any of the tasks on a
   steal_key = next(iter(a.state.executing)).key
   s.reschedule(steal_key, stimulus_id="steal")
   await ev.set()

This is not how stealing works. Stealing will be rejected if a key is already executing and nothing will be transitioned.

I'm not sure how such a reschedule would actually occur. If it did we'd intentionally calculate a key on two workers simultaneously.

[fjetter]

I'm currently struggling to come up with a case where this transition would actually occur. I believe we eradicated these type of transitions when removing the worker reconnect. However, I'm not entirely sure.

I do not consider it a valid API use to call Scheduler.reschedule. This is internally used when

• Stealing made a mistake and needs to correct because the key is not running anywhere
• The user raises a Reschedule exception. in this case, by definition the key is already done

[gjoseph92]

Won't it occur just by calling Scheduler.reschedule directly in the circumstance Guido has created, where the task is root-ish and all workers are full?

The implementation of Scheduler.reschedule seems odd to me. In work stealing, when we want to reschedule a task in processing, we first talk to the worker, get confirmation it's cancelled the task, and only then change state on the scheduler.

Scheduler.reschedule just blindly releases the state on the scheduler, even though the task may be happily executing on a worker.

It seems like the implementation of reschedule is targeted very much towards being the handler for Reschedule exceptions. That's a very specific use-case, where we can safely make the assumption that:

1. the task is no longer actually processing on the worker we thought it was, and that worker has already released the task internally
2. we are not going to receive any subsequent messages from that worker about that task (it's not going to go to memory immediately after)

That specific case doesn't seem like something that should be a public API. I would guess that the only reason it works to call it in work stealing is that the situation is equivalent, because move_task_request has already cancelled the task on the original worker, so we know no more updates about the task will come from there.

If Scheduler.reschedule were actually a generic, public API, I would expect it to implement the same move_task_request, move_task_confirm pattern that work-stealing does. Otherwise it leaves the door open for this inconsistency where a task is executing on two workers simultaneously.

I am leaning towards #7205 (comment) that reschedule should not be a public API, and we should add a comment to note the specific case in which it can be used internally.

[crusaderky]

So if I understand correctly, if I replace the contents of

• transition_processing_memory (with unexpected worker)
• transition_waiting_memory
• transition_no_worker_memory

with breakpoint(), I should catch exclusively use cases that none of us know about?

[crusaderky]

Yikes.
I did the above (crusaderky@d71804a) and I got ONE waiting->memory in test_stress_creation_and_deletion on localhost, and ZERO on CI 😨
These transitions are virtually untested.

story of test_stress_creation_and_deletion (breakpoint in transition_waiting_memory; edited for readability):

[
    ("('sum-partial-0', 9, 7)", 'released', 'waiting', {}, 'update-graph-1666879902.9497185', 1666879902.9718037),
    ("('sum-0', 19, 14)", 'processing', 'memory', {"('sum-partial-0', 9, 7)": 'processing', "('random_sample-0', 19, 14)": 'released', "('random_sample-0', 14, 19)": 'released'}, 'task-finished-1666879904.9479938', 1666879904.956455)
    ("('sum-partial-0', 9, 7)", 'waiting', 'processing', {}, 'task-finished-1666879904.9479938', 1666879904.9565284
    ("('sum-0', 19, 14)", 'memory', 'released', {"('sum-0', 19, 14)": 'waiting', "('sum-partial-0', 9, 7)": 'waiting'}, 'handle-worker-cleanup-1666879904.976327', 1666879904.9796247)
    ("('sum-partial-0', 9, 7)", 'processing', 'released', {"('sum-partial-0', 9, 7)": 'waiting'}, 'handle-worker-cleanup-1666879904.976327', 1666879904.9796507)
    ("('sum-partial-0', 9, 7)", 'released', 'waiting', {"('sum-partial-0', 9, 7)": 'waiting', "('sum-0', 19, 14)": 'waiting'}, 'handle-worker-cleanup-1666879904.976327', 1666879904.979665)
    ("('sum-partial-0', 9, 7)", 'waiting', 'memory', <breakpoint>
]

What I think I'm reading here:

1. sum-partial-0 is processing on a worker
2. the worker abruptly leaves the cluster (handle-worker-cleanup)
3. this causes the task to transition from processing to waiting (because we also just lost its dependency sum-0)
4. after that happens, a worker (which I inspected in the transition to have status=running) tells the scheduler that it completed the task. I do not know if it's the same worker that left the cluster before or not.

[crusaderky]

My opinion: none of this is new.
I'd suggest to merge a PR now (WIP) that implements transition_queued_memory, as a copy-paste of the other equivalent transitions, without any tests - as the others don't have any either - and then open a follow-up issue for investigation and test coverage.