Repository: PrefectHQ/prefect

Issue:
LocalDaskExecutor with threaded scheduler fails to run mapped tasks with timeout on linux
## Description
The following example: https://github.com/mianos/pbug demonstrates prefect/dask not scheduling a trivial task.
python slow.py --range=10 --sleep_time=3 > ubuntu2004run.txt

https://github.com/mianos/pbug/blob/main/ubuntu/ubuntu2004run.txt
This is on a 8 core local physical machine with no load.


## Expected Behavior
The messages '==== IN TASK XX Sleeping 3' should appear for each item in the range.
Note that the task runner is scheduling the tasks from the lines 'Task 'slow_task[3]': Starting task run...'
The result set should include results for all items.
The task runner hits the timeout (30 seconds) and everything gives up


## Reproduction
https://github.com/mianos/pbug/blob/main/slow.py
checkout, make a venv, pip install -r requirements.txt, run as above.

Sometimes this works but it seems the faster the machine the more it fails.

## Environment
I have tested it on redhat, CentOS Linux release 7.7.1908 (Core)

The results provided above are from Ununtu 20.04
Run and all versions are provided here:
https://github.com/mianos/pbug/tree/main/ubuntu

Fails on python 3.85 and python 3.6.8

{
  "config_overrides": {},
  "env_vars": [],
  "system_information": {
    "platform": "Linux-5.4.0-48-generic-x86_64-with-glibc2.29",
    "prefect_version": "0.13.3",
    "python_version": "3.8.5"
  }
}

I can provide SSH access to some machines demonstrating the failure if needed.


Discussion:
Linking to the community thread for reference https://prefect-community.slack.com/archives/CL09KU1K7/p1602728642497400

Resolving patch:
diff --git a/src/prefect/engine/task_runner.py b/src/prefect/engine/task_runner.py
--- a/src/prefect/engine/task_runner.py
+++ b/src/prefect/engine/task_runner.py
@@ -1,5 +1,6 @@
 from contextlib import redirect_stdout
 from dask.base import tokenize
+from contextlib import AbstractContextManager
 from typing import (
     Any,
     Callable,
@@ -40,6 +41,7 @@
     run_with_heartbeat,
     tail_recursive,
 )
+from prefect.utilities.compatibility import nullcontext
 
 
 TaskRunnerInitializeResult = NamedTuple(
@@ -825,17 +827,20 @@ def get_task_run_state(self, state: State, inputs: Dict[str, Result]) -> State:
                     name=prefect.context.get("task_full_name", self.task.name)
                 )
             )
-            timeout_handler = prefect.utilities.executors.timeout_handler
-            if getattr(self.task, "log_stdout", False):
-                with redirect_stdout(  # type: ignore
-                    prefect.utilities.logging.RedirectToLog(self.logger)
-                ):
-                    value = timeout_handler(
-                        self.task.run, timeout=self.task.timeout, **raw_inputs
-                    )
-            else:
-                value = timeout_handler(
-                    self.task.run, timeout=self.task.timeout, **raw_inputs
+
+            # Create a stdout redirect if the task has log_stdout enabled
+            log_context = (
+                redirect_stdout(prefect.utilities.logging.RedirectToLog(self.logger))
+                if getattr(self.task, "log_stdout", False)
+                else nullcontext()
+            )  # type: AbstractContextManager
+
+            with log_context:
+                value = prefect.utilities.executors.run_task_with_timeout(
+                    task=self.task,
+                    args=(),
+                    kwargs=raw_inputs,
+                    logger=self.logger,
                 )
 
         # inform user of timeout
diff --git a/src/prefect/utilities/compatibility.py b/src/prefect/utilities/compatibility.py
new file mode 100644
--- /dev/null
+++ b/src/prefect/utilities/compatibility.py
@@ -0,0 +1,23 @@
+"""
+Functions that exist for compatibility with different python versions. These should
+only be used internally.
+"""
+
+import sys
+
+
+# Provide `nullcontext`
+
+if sys.version_info < (3, 7):
+
+    from contextlib import contextmanager
+    from typing import Iterator
+
+    @contextmanager
+    def nullcontext() -> Iterator[None]:
+        yield
+
+
+else:
+
+    from contextlib import nullcontext  # noqa: F401
diff --git a/src/prefect/utilities/executors.py b/src/prefect/utilities/executors.py
--- a/src/prefect/utilities/executors.py
+++ b/src/prefect/utilities/executors.py
@@ -1,4 +1,5 @@
 import copy
+import cloudpickle
 import itertools
 import multiprocessing
 import os
@@ -10,11 +11,11 @@
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import TimeoutError as FutureTimeout
 from functools import wraps
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Union, Sequence
-
-import cloudpickle
+from logging import Logger
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Union, Sequence, Mapping
 
 import prefect
+from prefect.utilities.logging import get_logger
 
 if TYPE_CHECKING:
     import prefect.engine.runner
@@ -23,6 +24,7 @@
     from prefect.core.task import Task
     from prefect.engine.state import State
 
+
 StateList = Union["State", List["State"]]
 
 
@@ -77,24 +79,40 @@ def inner(
     return inner
 
 
-def main_thread_timeout(
-    fn: Callable, *args: Any, timeout: int = None, **kwargs: Any
+def run_with_thread_timeout(
+    fn: Callable,
+    args: Sequence = (),
+    kwargs: Mapping = None,
+    timeout: int = None,
+    logger: Logger = None,
+    name: str = None,
 ) -> Any:
     """
     Helper function for implementing timeouts on function executions.
     Implemented by setting a `signal` alarm on a timer. Must be run in the main thread.
+
     Args:
         - fn (callable): the function to execute
-        - *args (Any): arguments to pass to the function
-        - timeout (int): the length of time to allow for
-            execution before raising a `TimeoutError`, represented as an integer in seconds
-        - **kwargs (Any): keyword arguments to pass to the function
+        - args (Sequence): arguments to pass to the function
+        - kwargs (Mapping): keyword arguments to pass to the function
+        - timeout (int): the length of time to allow for execution before raising a
+            `TimeoutError`, represented as an integer in seconds
+        - logger (Logger): an optional logger to use. If not passed, a logger for the
+            `prefect.executors.run_with_thread_timeout` namespace will be created.
+        - name (str): an optional name to attach to logs for this function run, defaults
+            to the name of the given function. Provides an interface for passing task
+            names for logs.
+
     Returns:
-        - the result of `f(*args, **kwargs)`
+        - the result of `fn(*args, **kwargs)`
+
     Raises:
         - TimeoutError: if function execution exceeds the allowed timeout
         - ValueError: if run from outside the main thread
     """
+    logger = logger or get_logger()
+    name = name or f"Function '{fn.__name__}'"
+    kwargs = kwargs or {}
 
     if timeout is None:
         return fn(*args, **kwargs)
@@ -103,21 +121,25 @@ def error_handler(signum, frame):  # type: ignore
         raise TimeoutError("Execution timed out.")
 
     try:
+        # Set the signal handler for alarms
         signal.signal(signal.SIGALRM, error_handler)
+        # Raise the alarm if `timeout` seconds pass
+        logger.debug(f"{name}: Sending alarm with {timeout}s timeout...")
         signal.alarm(timeout)
+        logger.debug(f"{name}: Executing function in main thread...")
         return fn(*args, **kwargs)
     finally:
         signal.alarm(0)
 
 
-def multiprocessing_safe_retrieve_value(
+def multiprocessing_safe_run_and_retrieve(
     queue: multiprocessing.Queue,
     payload: bytes,
 ) -> None:
     """
     Gets the return value from a function and puts it in a multiprocessing-safe
-    container. Helper function for `multiprocessing_timeout`, must be defined top-level
-    so it can be pickled and sent to `multiprocessing.Process`
+    container. Helper function for `run_with_multiprocess_timeout`, must be defined
+    top-level so it can be pickled and sent to `multiprocessing.Process`
 
     Passing the payload serialized allows us to escape the limitations of the python
     native pickler which will fail on tasks defined in scripts because of name
@@ -131,8 +153,11 @@ def multiprocessing_safe_retrieve_value(
             Expects the following keys:
             - fn (Callable): The function to call
             - args (list): Positional argument values to call the function with
-            - kwargs (dict): Keyword arguments to call the function with
+            - kwargs (Mapping): Keyword arguments to call the function with
             - context (dict): The prefect context dictionary to use during execution
+            - name (str): an optional name to attach to logs for this function run,
+                defaults to the name of the given function. Provides an interface for
+                passing task names for logs.
 
     Returns:
         - None
@@ -141,22 +166,32 @@ def multiprocessing_safe_retrieve_value(
     """
     request = cloudpickle.loads(payload)
 
+    logger = get_logger()
+
     fn: Callable = request["fn"]
     context: dict = request.get("context", {})
     args: Sequence = request.get("args", [])
-    kwargs: dict = request.get("kwargs", {})
+    kwargs: Mapping = request.get("kwargs", {})
+    name: str = request.get("name", f"Function '{fn.__name__}'")
 
     try:
         with prefect.context(context):
+            logger.info(f"{name}: Executing...")
             return_val = fn(*args, **kwargs)
     except Exception as exc:
         return_val = exc
 
+    logger.info(f"{name}: Passing result back to main process...")
     queue.put(cloudpickle.dumps(return_val))
 
 
-def multiprocessing_timeout(
-    fn: Callable, *args: Any, timeout: int = None, **kwargs: Any
+def run_with_multiprocess_timeout(
+    fn: Callable,
+    args: Sequence = (),
+    kwargs: Mapping = None,
+    timeout: int = None,
+    logger: Logger = None,
+    name: str = None,
 ) -> Any:
     """
     Helper function for implementing timeouts on function executions.
@@ -164,10 +199,15 @@ def multiprocessing_timeout(
 
     Args:
         - fn (callable): the function to execute
-        - *args (Any): arguments to pass to the function
+        - args (Sequence): arguments to pass to the function
+        - kwargs (Mapping): keyword arguments to pass to the function
         - timeout (int): the length of time to allow for execution before raising a
             `TimeoutError`, represented as an integer in seconds
-        - **kwargs (Any): keyword arguments to pass to the function
+        - logger (Logger): an optional logger to use. If not passed, a logger for the
+            `prefect.` namespace will be created.
+        - name (str): an optional name to attach to logs for this function run, defaults
+            to the name of the given function. Provides an interface for passing task
+            names for logs.
 
     Returns:
         - the result of `f(*args, **kwargs)`
@@ -176,12 +216,17 @@ def multiprocessing_timeout(
         - AssertionError: if run from a daemonic process
         - TimeoutError: if function execution exceeds the allowed timeout
     """
+    logger = logger or get_logger()
+    name = name or f"Function '{fn.__name__}'"
+    kwargs = kwargs or {}
 
     if timeout is None:
         return fn(*args, **kwargs)
 
+    spawn_mp = multiprocessing.get_context("spawn")
+
     # Create a queue to pass the function return value back
-    queue = multiprocessing.Queue()  # type: multiprocessing.Queue
+    queue = spawn_mp.Queue()  # type: multiprocessing.Queue
 
     # Set internal kwargs for the helper function
     request = {
@@ -189,46 +234,58 @@ def multiprocessing_timeout(
         "args": args,
         "kwargs": kwargs,
         "context": prefect.context.to_dict(),
+        "name": name,
     }
     payload = cloudpickle.dumps(request)
 
-    p = multiprocessing.Process(
-        target=multiprocessing_safe_retrieve_value, args=(queue, payload)
+    run_process = spawn_mp.Process(
+        target=multiprocessing_safe_run_and_retrieve, args=(queue, payload)
     )
-    p.start()
-    p.join(timeout)
-    p.terminate()
+    logger.debug(f"{name}: Sending execution to a new process...")
+    run_process.start()
+    logger.debug(f"{name}: Waiting for process to return with {timeout}s timeout...")
+    run_process.join(timeout)
+    run_process.terminate()
 
     # Handle the process result, if the queue is empty the function did not finish
     # before the timeout
+    logger.debug(f"{name}: Execution process closed, collecting result...")
     if not queue.empty():
-        res = cloudpickle.loads(queue.get())
-        if isinstance(res, Exception):
-            raise res
-        return res
+        result = cloudpickle.loads(queue.get())
+        if isinstance(result, Exception):
+            raise result
+        return result
     else:
-        raise TimeoutError("Execution timed out.")
+        raise TimeoutError(f"Execution timed out for {name}.")
 
 
-def timeout_handler(
-    fn: Callable, *args: Any, timeout: int = None, **kwargs: Any
+def run_task_with_timeout(
+    task: "Task",
+    args: Sequence = (),
+    kwargs: Mapping = None,
+    logger: Logger = None,
 ) -> Any:
     """
-    Helper function for implementing timeouts on function executions.
+    Helper function for implementing timeouts on task executions.
 
     The exact implementation varies depending on whether this function is being
     run in the main thread or a non-daemonic subprocess.  If this is run from a
-    daemonic subprocess or on Windows, the task is run in a
-    `ThreadPoolExecutor` and only a soft timeout is enforced, meaning a
-    `TimeoutError` is raised at the appropriate time but the task continues
-    running in the background.
+    daemonic subprocess or on Windows, the task is run in a `ThreadPoolExecutor`
+    and only a soft timeout is enforced, meaning a `TimeoutError` is raised at the
+    appropriate time but the task continues running in the background.
+
+    The task is passed instead of a function so we can give better logs and messages.
+    If you need to run generic functions with timeout handlers,
+    `run_with_thread_timeout` or `run_with_multiprocess_timeout` can be called directly
 
     Args:
-        - fn (callable): the function to execute
-        - *args (Any): arguments to pass to the function
-        - timeout (int): the length of time to allow for execution before raising a
-            `TimeoutError`, represented as an integer in seconds
-        - **kwargs (Any): keyword arguments to pass to the function
+        - task (Task): the task to execute
+            `task.timeout` specifies the number of seconds to allow `task.run` to run
+            for before terminating
+        - args (Sequence): arguments to pass to the function
+        - kwargs (Mapping): keyword arguments to pass to the function
+        - logger (Logger): an optional logger to use. If not passed, a logger for the
+            `prefect.run_task_with_timeout_handler` namespace will be created.
 
     Returns:
         - the result of `f(*args, **kwargs)`
@@ -236,21 +293,52 @@ def timeout_handler(
     Raises:
         - TimeoutError: if function execution exceeds the allowed timeout
     """
+    logger = logger or get_logger()
+    name = prefect.context.get("task_full_name", task.name)
+    kwargs = kwargs or {}
+
     # if no timeout, just run the function
-    if timeout is None:
-        return fn(*args, **kwargs)
+    if task.timeout is None:
+        return task.run(*args, **kwargs)  # type: ignore
 
     # if we are running the main thread, use a signal to stop execution at the
     # appropriate time; else if we are running in a non-daemonic process, spawn
     # a subprocess to kill at the appropriate time
     if not sys.platform.startswith("win"):
+
         if threading.current_thread() is threading.main_thread():
-            return main_thread_timeout(fn, *args, timeout=timeout, **kwargs)
+            # This case is typically encountered when using a non-parallel or local
+            # multiprocess scheduler because then each worker is in the main
+            # thread
+            logger.debug(f"Task '{name}': Attaching thread based timeout handler...")
+            return run_with_thread_timeout(
+                task.run,
+                args,
+                kwargs,
+                timeout=task.timeout,
+                logger=logger,
+                name=f"Task '{name}'",
+            )
+
         elif multiprocessing.current_process().daemon is False:
-            return multiprocessing_timeout(fn, *args, timeout=timeout, **kwargs)
+            # This case is typically encountered when using a multithread distributed
+            # executor
+            logger.debug(f"Task '{name}': Attaching process based timeout handler...")
+            return run_with_multiprocess_timeout(
+                task.run,
+                args,
+                kwargs,
+                timeout=task.timeout,
+                logger=logger,
+                name=f"Task '{name}'",
+            )
 
+        # We are in a daemonic process and cannot enforce a timeout
+        # This case is typically encountered when using a multiprocess distributed
+        # executor
         soft_timeout_reason = "in a daemonic subprocess"
     else:
+        # We are in windows and cannot enforce a timeout
         soft_timeout_reason = "on Windows"
 
     msg = (
@@ -260,19 +348,22 @@ def timeout_handler(
         "but continue running in the background."
     )
 
+    logger.debug(
+        f"Task '{name}': Falling back to daemonic soft limit timeout handler because "
+        f"we are running {soft_timeout_reason}."
+    )
     warnings.warn(msg, stacklevel=2)
     executor = ThreadPoolExecutor()
 
-    def run_with_ctx(*args: Any, _ctx_dict: dict, **kwargs: Any) -> Any:
-        with prefect.context(_ctx_dict):
-            return fn(*args, **kwargs)
+    def run_with_ctx(context: dict) -> Any:
+        with prefect.context(context):
+            return task.run(*args, **kwargs)  # type: ignore
 
-    fut = executor.submit(
-        run_with_ctx, *args, _ctx_dict=prefect.context.to_dict(), **kwargs
-    )
+    # Run the function in the background and then retrieve its result with a timeout
+    fut = executor.submit(run_with_ctx, prefect.context.to_dict())
 
     try:
-        return fut.result(timeout=timeout)
+        return fut.result(timeout=task.timeout)
     except FutureTimeout as exc:
         raise TimeoutError(
             f"Execution timed out but was executed {soft_timeout_reason} and will "


Test patch:
diff --git a/tests/engine/test_task_runner.py b/tests/engine/test_task_runner.py
--- a/tests/engine/test_task_runner.py
+++ b/tests/engine/test_task_runner.py
@@ -307,6 +307,9 @@ def test_task_runner_accepts_dictionary_of_edges():
     sys.platform == "win32", reason="Windows doesn't support any timeout logic"
 )
 def test_timeout_actually_stops_execution():
+    # Note this replicates test coverage from
+    # `tests.core.test_flow.test_timeout_actually_stops_execution`
+
     # Note: this is a potentially brittle test! In some cases (local and sync) signal.alarm
     # is used as the mechanism for timing out a task. This passes off the job of measuring
     # the time for the timeout to the OS, which uses the "wallclock" as reference (the real
diff --git a/tests/utilities/test_executors.py b/tests/utilities/test_executors.py
--- a/tests/utilities/test_executors.py
+++ b/tests/utilities/test_executors.py
@@ -10,113 +10,180 @@
 import prefect
 from prefect.utilities.configuration import set_temporary_config
 from prefect.utilities.executors import (
-    timeout_handler,
+    run_with_thread_timeout,
+    run_with_multiprocess_timeout,
     tail_recursive,
     RecursiveCall,
 )
 
 
-def test_timeout_handler_times_out():
+# We will test the low-level timeout handlers here and `run_task_with_timeout`
+# is covered in `tests.core.test_flow.test_timeout_actually_stops_execution`
+# and `tests.engine.test_task_runner.test_timeout_actually_stops_execution`
+TIMEOUT_HANDLERS = [run_with_thread_timeout, run_with_multiprocess_timeout]
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+@pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
+def test_timeout_handler_times_out(timeout_handler):
     slow_fn = lambda: time.sleep(2)
     with pytest.raises(TimeoutError):
-        timeout_handler(slow_fn, timeout=1)
+        timeout_handler(slow_fn, timeout=2)
 
 
 @pytest.mark.skipif(
     sys.platform == "win32", reason="Windows doesn't support any timeout logic"
 )
-def test_timeout_handler_actually_stops_execution():
+@pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
+def test_timeout_handler_actually_stops_execution(timeout_handler):
     with tempfile.TemporaryDirectory() as call_dir:
         FILE = os.path.join(call_dir, "test.txt")
+        # Increase timeout for multiprocess because its startup time is slower and
+        # the file was not always created
+        TIMEOUT = 1 if timeout_handler is run_with_thread_timeout else 3
+        WRITES = 6
 
         def slow_fn():
-            "Runs for 1.5 seconds, writes to file 6 times"
-            iters = 0
-            while iters < 6:
-                time.sleep(0.26)
+            "Runs for TIMEOUT * 2 seconds, writes to file WRITES times"
+            for _ in range(WRITES):
                 with open(FILE, "a") as f:
                     f.write("called\n")
-                iters += 1
+                time.sleep((TIMEOUT * 2) / WRITES)
 
         with pytest.raises(TimeoutError):
-            # allow for at most 3 writes
-            timeout_handler(slow_fn, timeout=1)
+            # We should get less than WRITES lines -- roughly half expected
+            timeout_handler(slow_fn, timeout=TIMEOUT)
 
         time.sleep(0.5)
         with open(FILE, "r") as g:
             contents = g.read()
 
-    assert len(contents.split("\n")) <= 4
+    assert len(contents.split("\n")) < WRITES
 
 
-def test_timeout_handler_passes_args_and_kwargs_and_returns():
-    def do_nothing(x, y=None):
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+@pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
+def test_timeout_handler_passes_args_and_kwargs_and_returns(timeout_handler):
+    def just_return(x, y=None):
         return x, y
 
-    assert timeout_handler(do_nothing, 5, timeout=1, y="yellow") == (5, "yellow")
+    assert timeout_handler(
+        just_return, args=[5], kwargs=dict(y="yellow"), timeout=2
+    ) == (
+        5,
+        "yellow",
+    )
 
 
-def test_timeout_handler_doesnt_swallow_bad_args():
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+@pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
+def test_timeout_handler_doesnt_swallow_bad_args(timeout_handler):
     def do_nothing(x, y=None):
         return x, y
 
     with pytest.raises(TypeError):
-        timeout_handler(do_nothing, timeout=1)
+        timeout_handler(do_nothing, timeout=2)
 
     with pytest.raises(TypeError):
-        timeout_handler(do_nothing, 5, timeout=1, z=10)
+        timeout_handler(do_nothing, args=[5], kwargs=dict(z=10), timeout=2)
 
     with pytest.raises(TypeError):
-        timeout_handler(do_nothing, 5, timeout=1, y="s", z=10)
+        timeout_handler(do_nothing, args=[5], kwargs=dict(y="s", z=10), timeout=2)
 
 
-def test_timeout_handler_reraises():
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+@pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
+def test_timeout_handler_reraises(timeout_handler):
     def do_something():
         raise ValueError("test")
 
     with pytest.raises(ValueError, match="test"):
-        timeout_handler(do_something, timeout=1)
+        timeout_handler(do_something, timeout=2)
 
 
-@pytest.mark.skipif(sys.platform == "win32", reason="Test fails on Windows")
-def test_timeout_handler_allows_function_to_spawn_new_process():
+# Define a top-level helper function for a null-op process target, must be defined
+# as a non-local for the python native pickler used within `my_process`
+def do_nothing():
+    return None
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+@pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
+def test_timeout_handler_allows_function_to_spawn_new_process(timeout_handler):
     def my_process():
-        p = multiprocessing.Process(target=lambda: 5)
+        p = multiprocessing.Process(target=do_nothing())
         p.start()
         p.join()
         p.terminate()
 
-    assert timeout_handler(my_process, timeout=1) is None
+    assert timeout_handler(my_process, timeout=3) is None
 
 
-@pytest.mark.skipif(sys.platform == "win32", reason="Test fails on Windows")
-def test_timeout_handler_allows_function_to_spawn_new_thread():
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+@pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
+def test_timeout_handler_allows_function_to_spawn_new_thread(timeout_handler):
     def my_thread():
         t = threading.Thread(target=lambda: 5)
         t.start()
         t.join()
 
-    assert timeout_handler(my_thread, timeout=1) is None
+    assert timeout_handler(my_thread, timeout=3) is None
 
 
-def test_timeout_handler_doesnt_do_anything_if_no_timeout(monkeypatch):
-    assert timeout_handler(lambda: 4, timeout=1) == 4
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+@pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
+def test_timeout_handler_doesnt_do_anything_if_no_timeout(timeout_handler):
+    assert timeout_handler(lambda: 4, timeout=2) == 4
     assert timeout_handler(lambda: 4) == 4
 
 
-def test_timeout_handler_preserves_context():
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+@pytest.mark.parametrize("timeout_handler", TIMEOUT_HANDLERS)
+def test_timeout_handler_preserves_context(timeout_handler):
     def my_fun(x, **kwargs):
         return prefect.context.get("test_key")
 
     with prefect.context(test_key=42):
-        res = timeout_handler(my_fun, 2, timeout=1)
+        res = timeout_handler(my_fun, args=[2], timeout=2)
 
     assert res == 42
 
 
-def test_timeout_handler_preserves_logging(caplog):
-    timeout_handler(prefect.Flow("logs").run, timeout=2)
-    assert len(caplog.records) >= 2  # 1 INFO to start, 1 INFO to end
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+def test_run_with_thread_timeout_preserves_logging(caplog):
+    run_with_thread_timeout(prefect.Flow("logs").run, timeout=2)
+    assert len(caplog.messages) >= 2  # 1 INFO to start, 1 INFO to end
+
+
+@pytest.mark.skipif(
+    sys.platform == "win32", reason="Windows doesn't support any timeout logic"
+)
+def test_run_with_multiprocess_timeout_preserves_logging(capfd):
+    """
+    Requires fd capturing because the subprocess output won't be captured by caplog
+    """
+    run_with_multiprocess_timeout(prefect.Flow("logs").run, timeout=2)
+    stdout = capfd.readouterr().out
+    assert "Beginning Flow run" in stdout
+    assert "Flow run SUCCESS" in stdout
 
 
 def test_recursion_go_case():