Docs discrepancy: dask.annotate(priority=...) has no effect with Prefect 3.x DaskTaskRunner

[zzstoatzz]

This discussion was created from a Slack thread conversation.

Original Thread: https://prefect-community.slack.com/archives/C079VLLH5D3/p1776339347743129

---

Summary
The Prefect docs page at https://docs.prefect.io/integrations/prefect-dask#with-prefect shows an example suggesting users can influence task execution order with dask.annotate(priority=...). In Prefect 3.x with the DaskTaskRunner, this annotation does not change scheduling order because Prefect’s Dask client passes an explicit priority=0 to distributed.Client.submit(...) for each task, which overrides any annotation context. The result is that annotated priorities do not take effect in Prefect-managed task submission.

Impact
Users trying to prioritize tasks (e.g., when m > n workers) by using dask.annotate(...) will see no effect and get confused, as reported by a user in Slack. This contradicts the current docs example.

Observed behavior

• Environment: Prefect 3.x, prefect-dask integration, local Dask cluster (n_workers > 1, threads_per_worker=1)
• User code (simplified):

from prefect import flow, task
from prefect_dask import DaskTaskRunner
import dask

@task
def bar(x):
    ...

@flow(task_runner=DaskTaskRunner(cluster_kwargs={"n_workers": 20,
                                                 "threads_per_worker": 1,
                                                 "memory_limit": "20GiB",
                                                 "silence_logs": False}),
      log_prints=True)
def foo(a):
    futures = []
    for x in a:
        with dask.annotate(priority=len(x)):
            futures.append(bar.submit(x=x))
    for f in futures:
        f.result()

• Expected (per docs): Tasks with higher priority run earlier.
• Actual: No prioritization effect observed; tasks run as workers free up.

Likely root cause (code-level)

• PrefectDaskClient.submit(...) forwards an explicit priority kwarg to distributed.Client.submit(...), defaulting to 0.
• In Dask, an explicit priority= argument overrides any annotate(...) context, including distributed.annotate.
• DaskTaskRunner.submit(...) does not expose a priority parameter, and Task.submit(...) does not either, so users cannot pass priority through Prefect at this time.

Requested doc correction

• On https://docs.prefect.io/integrations/prefect-dask#with-prefect, clarify that dask.annotate(priority=...) (and even distributed.annotate) will not affect Prefect-managed task submission because an explicit priority is set when submitting to Dask.
• Provide current workarounds:
  • Use dependency ordering via wait_for for strict sequencing.
  • Stage submissions (submit high-priority batch first, then others) if soft prioritization is sufficient.
• Optionally note a potential enhancement: add a first-class priority parameter to Prefect’s Task.submit(...) and/or DaskTaskRunner.submit(...) that passes through to the Dask client.

Links for context

• Docs page in question: https://docs.prefect.io/integrations/prefect-dask#with-prefect
• Dask annotations behavior: https://distributed.dask.org/en/stable/annotations.html

If maintainers confirm this behavior is intended, updating the docs would prevent confusion for users trying to control order with annotations under Prefect’s DaskTaskRunner.

---

This discussion was automatically created by the Marvin bot to preserve valuable community insights.

[pratyush618]

Thanks for filing this from the Slack thread — wanted to add a closer trace before we change the docs, because I think the root-cause attribution is slightly off.

PrefectDaskClient.submit does forward priority=0 to distributed.Client.submit, but in distributed.Client._graph_to_futures the priority kwarg is only injected into local annotations when it's truthy:

annotations = {}
if user_priority:
    annotations["priority"] = user_priority
...
annotations = merge(dask.get_annotations(), annotations)

With user_priority=0 no "priority" key is added locally, so the merge with dask.get_annotations() preserves whatever the active dask.annotate(...) context provides. So an explicit priority=0 shouldn't be overriding the annotation — it should just no-op.

The "no observable effect" symptom is almost certainly real, but the more likely explanations are:

1. fifo_timeout="100 ms" default. A tight loop submits everything inside a single 100 ms window, so the scheduler treats all of them as equal-priority for ordering purposes.
2. Priority doesn't preempt running tasks. With n_workers=20/threads_per_worker=1, the first 20 submissions get dispatched to workers immediately, so prioritization can only re-order whatever is still queued. The docs note right under the example already calls this out.
3. Thread-local context loss. dask.annotate is threading.local-based. If Prefect's task engine calls into the Dask client from a worker thread that wasn't inside the with block, dask.get_annotations() returns {} and the annotation is silently dropped. Worth a repro to confirm — if this is what's happening, it's the actual bug.

Workarounds today:

• wait_for= for hard ordering.
• Two-phase submission (high-priority batch first, await briefly, then the rest) for a soft hint.
• As a last resort, drop down to get_client() from prefect_dask and submit the wrapper task directly with a non-zero priority= (which will win over annotations, by the merge rule above).

Proposed follow-ups:

• Doc nit: the existing note under the example is good but could explicitly mention fifo_timeout (rapid submissions look like same-priority) and that thread-context can drop the annotation.
• Feature: plumb a priority= kwarg through DaskTaskRunner.submit (and optionally Task.submit for runner-agnostic call sites) so users don't need to reach for the raw client.
• Investigation: before changing anything, repro with a single-worker / single-thread cluster and fifo_timeout="0 ms" so prioritization is the only variable, and confirm whether the annotation actually reaches the scheduler. If it doesn't, the fix is on the threading side, not the kwarg side.

Minimal repro to disambiguate the three theories:

import dask
from prefect import flow, task
from prefect_dask import DaskTaskRunner

@task
def slow(i):
    import time; time.sleep(0.5)
    return i

@flow(task_runner=DaskTaskRunner(
    cluster_kwargs={"n_workers": 1, "threads_per_worker": 1},  # force a queue
))
def f():
    futs = []
    for i in range(10):
        with dask.annotate(priority=i):  # later submissions get HIGHER priority
            futs.append(slow.submit(i))
    return [f.result() for f in futs]

Task 0 starts immediately on the only worker; tasks 1–9 queue. If annotations reach the scheduler, the queued tasks should finish in order 9, 8, 7, …, 1 (highest priority first). If they finish in submission order (1, 2, …, 9) instead, the annotation is being dropped — next thing to check is whether the actual distributed.Client.submit call lands on a thread that didn't enter the with block. Watching the dashboard's task stream is the easiest way to see this directly.