Typing: Make TaskRunner generic parameter covariant so @flow(task_runner=RayTaskRunner()) type-checks

[zzstoatzz]

This discussion was created from a Slack thread conversation.

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

---

Summary
A typing variance issue occurs when passing a RayTaskRunner to the @flow decorator in Prefect 3.x. Type checkers (e.g., pyright/mypy) raise an error because TaskRunner’s generic parameter F is invariant, and RayTaskRunner is parameterized over a subtype of PrefectFuture (PrefectRayFuture[T]). At runtime everything works; this is a static typing-only mismatch.

Reproduction

from typing import Any
from prefect import flow
from prefect_ray.task_runners import RayTaskRunner

@flow(task_runner=RayTaskRunner())
def my_flow() -> None:
    pass

Type checker error (example):

Argument of type "RayTaskRunner[Unknown]" cannot be assigned to parameter "task_runner" of type "TaskRunner[PrefectFuture[Any]] | None" in function "__call__"
  Type "RayTaskRunner[Unknown]" is not assignable to type "TaskRunner[PrefectFuture[Any]] | None"
    "RayTaskRunner[Unknown]" is not assignable to "TaskRunner[PrefectFuture[Any]]"
      Type parameter "F@TaskRunner" is invariant, but "PrefectRayFuture[Unknown]" is not the same as "PrefectFuture[Any]"

Root cause

• The @flow decorator parameter is typed as task_runner: Optional[TaskRunner[PrefectFuture[Any]]].
• RayTaskRunner (in prefect-ray) is TaskRunner[PrefectRayFuture[T]] where PrefectRayFuture[T] subclasses PrefectFuture[T].
• In prefect/src/prefect/task_runners.py, F is defined invariant:

  F = TypeVar("F", bound=PrefectFuture[Any], default=PrefectConcurrentFuture[Any])

  Because F is invariant, TaskRunner[PrefectRayFuture[Any]] is not a subtype of TaskRunner[PrefectFuture[Any]].

Why covariance is correct
A scan of TaskRunner shows F is only used in covariant (return) positions: return types of submit()/map()/duplicate(). F is not used in parameter positions. Making F covariant is therefore semantically correct and safe.

Proposed change (one line)
In src/prefect/task_runners.py, update F to be covariant:

- F = TypeVar("F", bound=PrefectFuture[Any], default=PrefectConcurrentFuture[Any])
+ F = TypeVar("F", bound=PrefectFuture[Any], default=PrefectConcurrentFuture[Any], covariant=True)

Impact

• Fixes type-checking for @flow(task_runner=RayTaskRunner()) and other third-party task runners returning a PrefectFuture subtype.
• No runtime behavior changes.
• Preserves useful type info while aligning with how F is used in the API.

Workarounds for users today

• Cast when passing the runner:

  from typing import Any, cast
  from prefect.futures import PrefectFuture
  from prefect.task_runners import TaskRunner
  from prefect_ray.task_runners import RayTaskRunner
  
  runner = cast(TaskRunner[PrefectFuture[Any]], RayTaskRunner())
  @flow(task_runner=runner)
  def my_flow():
      ...

• Or add a narrow type ignore on the assignment/argument (# type: ignore[arg-type]).

References

• TaskRunner and flow decorator: https://github.com/PrefectHQ/prefect/blob/main/src/prefect/task_runners.py, https://github.com/PrefectHQ/prefect/blob/main/src/prefect/flows.py
• Futures: https://github.com/PrefectHQ/prefect/blob/main/src/prefect/futures.py
• Ray runner: https://github.com/PrefectHQ/prefect-ray/blob/main/src/prefect_ray/task_runners.py

Request

• Please confirm that making F covariant is acceptable. If so, I’m happy to submit a PR with the change and typing tests. If there are additional constraints (e.g., other TaskRunner implementations relying on invariance), let’s discuss the preferred direction.

---

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

[Aliipou]

This is a classic variance mismatch. TaskRunner[F] is currently invariant in F, but it should be covariant for read-only use as a @flow parameter — you only ever consume futures from a task runner, never produce them into it from the outside.

The fix at the type-stub level:

from typing import TypeVar, Generic
from prefect.futures import PrefectFuture

F_co = TypeVar("F_co", bound=PrefectFuture, covariant=True)

class TaskRunner(Generic[F_co]):
    ...

Making F covariant (F_co) tells pyright/mypy that TaskRunner[PrefectRayFuture[T]] is a valid subtype of TaskRunner[PrefectFuture[T]], which matches runtime behaviour.

Workaround while waiting for the fix:

from typing import cast
from prefect.task_runners import TaskRunner

@flow(task_runner=cast(TaskRunner, RayTaskRunner()))
def my_flow():
    ...

Or use # type: ignore[arg-type] on the @flow line.

The root change needed is in prefect/task_runners.py — change TypeVar("F", bound=PrefectFuture) to TypeVar("F_co", bound=PrefectFuture, covariant=True) and update the class signature. Runtime behavior is unaffected since covariance is a type-checking-only constraint.