text: import copy
  import itertools
  import multiprocessing
  import os
  import signal
  import subprocess
  import sys
  import threading
  import warnings
  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
  
  import prefect
  
  if TYPE_CHECKING:
      import prefect.engine.runner
      import prefect.engine.state
      from prefect.core.edge import Edge
      from prefect.core.task import Task
      from prefect.engine.state import State
  
  StateList = Union["State", List["State"]]
  
  
  def run_with_heartbeat(
      runner_method: Callable[..., "prefect.engine.state.State"]
  ) -> Callable[..., "prefect.engine.state.State"]:
      """
      Utility decorator for running class methods with a heartbeat.  The class should implement
      `self._heartbeat` with no arguments.
      """
  
      @wraps(runner_method)
      def inner(
          self: "prefect.engine.runner.Runner", *args: Any, **kwargs: Any
      ) -> "prefect.engine.state.State":
          try:
              p = None
              try:
                  if self._heartbeat():
                      # we use Popen + a prefect CLI for a few reasons:
                      # - using threads would interfere with the task; for example, a task
                      #   which does not release the GIL would prevent the heartbeat thread from
                      #   firing
                      # - using multiprocessing.Process would release the GIL but a subprocess
                      #   cannot be spawned from a deamonic subprocess, and Dask sometimes will
                      #   submit tasks to run within daemonic subprocesses
                      current_env = dict(os.environ).copy()
                      auth_token = prefect.context.config.cloud.get("auth_token")
                      api_url = prefect.context.config.cloud.get("api")
                      current_env.setdefault("PREFECT__CLOUD__AUTH_TOKEN", auth_token)
                      current_env.setdefault("PREFECT__CLOUD__API", api_url)
                      clean_env = {k: v for k, v in current_env.items() if v is not None}
                      p = subprocess.Popen(
                          self.heartbeat_cmd,
                          env=clean_env,
                          stdout=subprocess.DEVNULL,
                          stderr=subprocess.DEVNULL,
                      )
              except Exception:
                  self.logger.exception(
                      "Heartbeat failed to start.  This could result in a zombie run."
                  )
              return runner_method(self, *args, **kwargs)
          finally:
              if p is not None:
                  exit_code = p.poll()
                  if exit_code is not None:
                      msg = "Heartbeat process died with exit code {}".format(exit_code)
                      self.logger.error(msg)
                  p.kill()
  
      return inner
  
  
  def main_thread_timeout(
      fn: Callable, *args: Any, timeout: int = None, **kwargs: Any
  ) -> 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
      Returns:
          - the result of `f(*args, **kwargs)`
      Raises:
          - TimeoutError: if function execution exceeds the allowed timeout
          - ValueError: if run from outside the main thread
      """
  
      if timeout is None:
          return fn(*args, **kwargs)
  
      def error_handler(signum, frame):  # type: ignore
          raise TimeoutError("Execution timed out.")
  
      try:
          signal.signal(signal.SIGALRM, error_handler)
          signal.alarm(timeout)
          return fn(*args, **kwargs)
      finally:
          signal.alarm(0)
  
  
  def multiprocessing_safe_retrieve_value(
      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`
  
      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
      mismatches. Whilst this particular example only affects the `func` arg, any of the
      others could be affected by other pickle limitations as well.
  
      Args:
          - queue (multiprocessing.Queue): The queue to pass the resulting payload to
          - payload (bytes): A serialized dictionary containing the data required to run
              the function. Should be serialized with `cloudpickle.dumps`
              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
              - context (dict): The prefect context dictionary to use during execution
  
      Returns:
          - None
          Passes the serialized (with cloudpickle) return value or exception into the
          queue. Callers are expected to re-raise any exceptions.
      """
      request = cloudpickle.loads(payload)
  
      fn: Callable = request["fn"]
      context: dict = request.get("context", {})
      args: Sequence = request.get("args", [])
      kwargs: dict = request.get("kwargs", {})
  
      try:
          with prefect.context(context):
              return_val = fn(*args, **kwargs)
      except Exception as exc:
          return_val = exc
  
      queue.put(cloudpickle.dumps(return_val))
  
  
  def multiprocessing_timeout(
      fn: Callable, *args: Any, timeout: int = None, **kwargs: Any
  ) -> Any:
      """
      Helper function for implementing timeouts on function executions.
      Implemented by spawning a new multiprocess.Process() and joining with timeout.
  
      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
  
      Returns:
          - the result of `f(*args, **kwargs)`
  
      Raises:
          - AssertionError: if run from a daemonic process
          - TimeoutError: if function execution exceeds the allowed timeout
      """
  
      if timeout is None:
          return fn(*args, **kwargs)
  
      # Create a queue to pass the function return value back
      queue = multiprocessing.Queue()  # type: multiprocessing.Queue
  
      # Set internal kwargs for the helper function
      request = {
          "fn": fn,
          "args": args,
          "kwargs": kwargs,
          "context": prefect.context.to_dict(),
      }
      payload = cloudpickle.dumps(request)
  
      p = multiprocessing.Process(
          target=multiprocessing_safe_retrieve_value, args=(queue, payload)
      )
      p.start()
      p.join(timeout)
      p.terminate()
  
      # Handle the process result, if the queue is empty the function did not finish
      # before the timeout
      if not queue.empty():
          res = cloudpickle.loads(queue.get())
          if isinstance(res, Exception):
              raise res
          return res
      else:
          raise TimeoutError("Execution timed out.")
  
  
  def timeout_handler(
      fn: Callable, *args: Any, timeout: int = None, **kwargs: Any
  ) -> Any:
      """
      Helper function for implementing timeouts on function 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.
  
      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
  
      Returns:
          - the result of `f(*args, **kwargs)`
  
      Raises:
          - TimeoutError: if function execution exceeds the allowed timeout
      """
      # if no timeout, just run the function
      if timeout is None:
          return fn(*args, **kwargs)
  
      # 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)
          elif multiprocessing.current_process().daemon is False:
              return multiprocessing_timeout(fn, *args, timeout=timeout, **kwargs)
  
          soft_timeout_reason = "in a daemonic subprocess"
      else:
          soft_timeout_reason = "on Windows"
  
      msg = (
          f"This task is running {soft_timeout_reason}; "
          "consequently Prefect can only enforce a soft timeout limit, i.e., "
          "if your Task reaches its timeout limit it will enter a TimedOut state "
          "but continue running in the background."
      )
  
      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)
  
      fut = executor.submit(
          run_with_ctx, *args, _ctx_dict=prefect.context.to_dict(), **kwargs
      )
  
      try:
          return fut.result(timeout=timeout)
      except FutureTimeout as exc:
          raise TimeoutError(
              f"Execution timed out but was executed {soft_timeout_reason} and will "
              "continue to run in the background."
          ) from exc
  
  
  class RecursiveCall(Exception):
      def __init__(self, func: Callable, *args: Any, **kwargs: Any):
          self.func = func
          self.args = args
          self.kwargs = kwargs
  
  
  def tail_recursive(func: Callable) -> Callable:
      """
      Helper function to facilitate tail recursion of the wrapped function.
  
      This allows for recursion with unlimited depth since a stack is not allocated for
      each "nested" call. Note: instead of calling the target function in question, a
      `RecursiveCall` exception must be raised instead.
  
      Args:
          - fn (callable): the function to execute
  
      Returns:
          - the result of `f(*args, **kwargs)`
  
      Raises:
          - RecursionError: if a recursive "call" (raised exception) is made with a function that is
              not decorated with `tail_recursive` decorator.
      """
  
      @wraps(func)
      def wrapper(*args: Any, **kwargs: Any) -> Any:
          while True:
              try:
                  return func(*args, **kwargs)
              except RecursiveCall as exc:
                  try:
                      call_func = getattr(exc.func, "__wrapped_func__")
                  except AttributeError as attr_error:
                      raise RecursionError(
                          "function has not been wrapped to provide tail recursion (func={})".format(
                              exc.func
                          )
                      ) from attr_error
  
                  # there may be multiple nested recursive calls, we should only
                  # respond to calls for the wrapped function explicitly,
                  # otherwise allow the call to continue to propagate
                  if call_func != func:
                      raise exc
                  args = exc.args
                  kwargs = exc.kwargs
                  continue
  
      setattr(wrapper, "__wrapped_func__", func)
      return wrapper
  
  
  def prepare_upstream_states_for_mapping(
      state: "State",
      upstream_states: "Dict[Edge, State]",
      mapped_children: "Dict[Task, list]",
      executor: "prefect.engine.executors.Executor",
  ) -> list:
      """
      If the task is being mapped, submits children tasks for execution. Returns a `Mapped` state.
  
      Args:
          - state (State): the parent task's current state
          - upstream_states (Dict[Edge, State]): the upstream states to this task
          - mapped_children (Dict[Task, List[State]]): any mapped children upstream of this task
  
      Returns:
          - List: a restructured list of upstream states correponding to each new mapped child task
      """
  
      # if the current state is failed / skipped or otherwise
      # in a state that signifies we should not continue with mapping,
      # we return an empty list
      if state.is_pending() or state.is_failed() or state.is_skipped():
          return []
  
      map_upstream_states = []
  
      # copy the mapped children dict to avoid mutating the global one
      # when we flatten mapped children
      mapped_children = mapped_children.copy()
  
      # we don't know how long the iterables are, but we want to iterate until we reach
      # the end of the shortest one
      counter = itertools.count()
  
      # preprocessing
      for edge, upstream_state in upstream_states.items():
  
          # ensure we are working with populated result objects
          if edge.key in state.cached_inputs:
              upstream_state._result = state.cached_inputs[edge.key]
  
          # if the upstream was mapped and the edge is flattened, we need
          # to process the mapped children (which could be futures) into a
          # flat structure
          if upstream_state.is_mapped() and edge.flattened:
              mapped_children[edge.upstream_task] = flatten_mapped_children(
                  mapped_children=mapped_children[edge.upstream_task], executor=executor
              )
  
      # infinite loop, if upstream_states has any entries
      while True and upstream_states:
          i = next(counter)
          states = {}
          try:
  
              for edge, upstream_state in upstream_states.items():
  
                  # if the edge is not mapped over, then we take its state
                  if not edge.mapped:
                      states[edge] = upstream_state
  
                  # if the edge is mapped and the upstream state is Mapped, then we are mapping
                  # over a mapped task. In this case, we take the appropriately-indexed upstream
                  # state from the upstream tasks's `Mapped.map_states` array.
                  # Note that these "states" might actually be futures at this time; we aren't
                  # blocking until they finish.
                  elif edge.mapped and upstream_state.is_mapped():
                      states[edge] = mapped_children[edge.upstream_task][i]  # type: ignore
  
                  # Otherwise, we are mapping over the result of a "vanilla" task. In this
                  # case, we create a copy of the upstream state but set the result to the
                  # appropriately-indexed item from the upstream task's `State.result`
                  # array.
                  else:
                      states[edge] = copy.copy(upstream_state)
  
                      # if the current state is already Mapped, then we might be executing
                      # a re-run of the mapping pipeline. In that case, the upstream states
                      # might not have `result` attributes.
                      # Therefore, we only try to get a result if EITHER this task's
                      # state is not already mapped OR the upstream result is not None.
                      if (
                          not state.is_mapped()
                          or upstream_state._result != prefect.engine.result.NoResult
                      ):
                          # this line should never be hit due to a check
                          # in the TaskRunner when evaluating the mapped parent
                          if not hasattr(upstream_state.result, "__getitem__"):
                              raise TypeError(
                                  (
                                      "Cannot map over unsubscriptable object of type {t}: {val}..."
                                  ).format(
                                      t=type(upstream_state.result),
                                      val=repr(upstream_state.result)[:10],
                                  )
                              )
                          upstream_result = upstream_state._result.from_value(  # type: ignore
                              upstream_state.result[i]
                          )
                          states[edge].result = upstream_result
                      elif state.is_mapped():
                          if i >= len(state.map_states):  # type: ignore
                              raise IndexError()
  
              # only add this iteration if we made it through all iterables
              map_upstream_states.append(states)
  
          # index error means we reached the end of the shortest iterable
          except IndexError:
              break
  
      return map_upstream_states
  
  
  def _build_flattened_state(state: "State", index: int) -> "State":
      """Helper function for `flatten_upstream_state`"""
      new_state = copy.copy(state)
      new_state.result = state._result.from_value(state.result[index])  # type: ignore
      return new_state
  
  
  def flatten_upstream_state(upstream_state: "State") -> "State":
      """
      Given an upstream state, returns its result as a flattened list. If
      flattening fails, the object is returned unmodified.
      """
      try:
          # attempt to unnest
          flattened_result = [y for x in upstream_state.result for y in x]
      except TypeError:
          return upstream_state
  
      new_state = copy.copy(upstream_state)
      new_state.result = new_state._result.from_value(flattened_result)  # type: ignore
      return new_state
  
  
  def flatten_mapped_children(
      mapped_children: List["State"],
      executor: "prefect.engine.executors.Executor",
  ) -> List["State"]:
      counts = executor.wait(
          [executor.submit(lambda c: len(c._result.value), c) for c in mapped_children]
      )
      new_states = []
  
      for child, count in zip(mapped_children, counts):
          new_states.append(
              [executor.submit(_build_flattened_state, child, i) for i in range(count)]
          )
  
      flattened_states = [i for s in new_states for i in s]
      return flattened_states
  