radical.asyncflow.workflow_manager¶

async def start_telemetry(
    self,
    resource_poll_interval: float = 5.0,
    checkpoint_interval: float | None = None,
    checkpoint_path: str | None = None,
    span_processors: list | None = None,
    metric_readers: list | None = None,
    resource: object | None = None,
) -> Any:
    """Create and start a RHAPSODY TelemetryManager for this workflow engine.

    Mirrors Session.start_telemetry(): creates the manager, wires all
    registered backends (silently skipping LocalExecutionBackend and
    NoopExecutionBackend which have no adapter), starts the async dispatch
    loop, and returns the manager.

    Args:
        resource_poll_interval: Seconds between resource metric polls (default: 5.0).
        checkpoint_interval:    Seconds between periodic metric+span flushes to disk.
                                None = no periodic flush (file still written at stop).
        checkpoint_path:        Directory for the JSONL checkpoint file.
                                None = no file output.
        span_processors:        Optional list of OTel SpanProcessor instances added to
                                RHAPSODY's TracerProvider alongside the internal
                                SpanBuffer. Callers own exporter construction.
        metric_readers:         Optional list of OTel MetricReader instances added to
                                RHAPSODY's MeterProvider alongside InMemoryMetricReader.
        resource:               Optional ``opentelemetry.sdk.resources.Resource``.
                                When None, ``Resource.create()`` reads
                                ``OTEL_SERVICE_NAME`` from the environment automatically.

    Returns:
        The active TelemetryManager instance.
    """
    from rhapsody.telemetry.manager import TelemetryManager  # noqa: PLC0415

    self._telemetry = TelemetryManager(
        session_id=self.uid,
        checkpoint_interval=checkpoint_interval,
        checkpoint_path=checkpoint_path,
        span_processors=span_processors,
        metric_readers=metric_readers,
        resource=resource,
    )

    for backend in self._backends.values():
        self._telemetry.attach_backend(
            backend,
            session_id=self.uid,
            backend_name=backend.name,
            interval=resource_poll_interval,
        )

    await self._telemetry.start()

    self._telemetry.register_span_enricher(
        lambda event: (
            {"asyncflow.workflow_id": event.attributes["asyncflow.workflow_id"]}
            if isinstance(event.attributes, dict)
            and "asyncflow.workflow_id" in event.attributes
            else {}
        )
    )

    from rhapsody.telemetry.events import (  # noqa: PLC0415
        TaskCanceled,
        TaskCompleted,
        TaskCreated,
        TaskFailed,
        TaskQueued,
        TaskStarted,
        TaskSubmitted,
        define_event,
        make_event,
    )

    # asyncflow.TaskResolved is DAG-specific — not a RHAPSODY core event.
    # Defined here with define_event so RHAPSODY stays unaware of DAG semantics.
    TaskResolved = define_event("asyncflow.TaskResolved")

    self._tel_make_event = make_event
    self._tel_events = {
        "TaskCreated": TaskCreated,
        "TaskQueued": TaskQueued,
        "asyncflow.TaskResolved": TaskResolved,
        "TaskSubmitted": TaskSubmitted,
        "TaskStarted": TaskStarted,
        "TaskCompleted": TaskCompleted,
        "TaskFailed": TaskFailed,
        "TaskCanceled": TaskCanceled,
    }

    return self._telemetry

@asynccontextmanager
async def workflow_scope(
    self, workflow_id: str | None = None
) -> AsyncIterator[str]:
    """Tag all tasks registered inside this scope with a shared workflow_id.

    Also creates an OTel workflow span (when telemetry is active) so that task spans
    registered inside become structural children in the trace hierarchy.

    Args:
        workflow_id: Explicit id for this workflow instance.
            If omitted, a short UUID is generated automatically.

    Yields:
        The workflow_id string in use.
    """
    wid = workflow_id or f"wf-{uuid.uuid4().hex[:8]}"
    token = self._workflow_id_ctx.set(wid)
    try:
        scope = (
            self._telemetry.span_scope("workflow", {"asyncflow.workflow_id": wid})
            if self._telemetry is not None
            else nullcontext()
        )
        with scope:
            yield wid
    finally:
        self._workflow_id_ctx.reset(token)

@classmethod
async def create(
    cls,
    backend: Any = None,
    dry_run: bool = False,
    implicit_data: bool = True,
    uid: Optional[str] = None,
    work_dir: Optional[str] = None,
) -> "WorkflowEngine":
    """Factory method to create and initialize a WorkflowEngine.

    Args:
        backend: Execution backend. If None and dry_run=True,
                 uses NoopExecutionBackend
        dry_run: Whether to run in dry-run mode
        implicit_data: Whether to enable implicit data dependency linking
        uid: Optional unique identifier for this engine. Defaults to
            ``asyncflow.session.{hex8}`` (UUID-based).
        work_dir: Working directory for output files (default: cwd)

    Returns:
        WorkflowEngine: Fully initialized workflow engine

    Example:
        engine = await WorkflowEngine.create(dry_run=True)
    """
    # Setup and validate backend first
    validated_backend = await cls._setup_execution_backend(backend, dry_run)

    # Create instance with validated backend
    instance = cls(
        backend=validated_backend,
        dry_run=dry_run,
        implicit_data=implicit_data,
        uid=uid,
        work_dir=work_dir,
    )

    # Initialize async components
    await instance._start_async_components()

    return instance

async def run(self):
    """Manages asynchronous execution of workflow components.

    Continuously monitors and manages workflow components, handling their
    dependencies and execution states. Performs dependency resolution and
    prepares components for execution when their dependencies are satisfied.

    Workflow Process:
        1. Monitors unresolved components
        2. Checks dependency resolution status
        3. Prepares resolved components for execution
        4. Handles data staging between components
        5. Submits ready components to execution

    Raises:
        asyncio.CancelledError: If the coroutine is cancelled during execution

    State Management:
        - unresolved (set): Component UIDs with pending dependencies
        - resolved (set): Component UIDs with satisfied dependencies
        - running (list): Currently executing component UIDs
        - dependencies (dict): Maps component UIDs to dependency info
        - components (dict): Maps UIDs to component descriptions and futures
        - queue (asyncio.Queue): Execution queue for ready components

    Note:
        - Runs indefinitely until cancelled or shutdown is signaled
        - Uses sleep intervals to prevent busy-waiting
        - Handles both implicit and explicit data dependencies
        - Trigger internal shutdown on loop failure
    """
    while not self._shutdown_event.is_set():
        try:
            to_submit = []

            # Process ready components
            while self._ready_queue and not self._shutdown_event.is_set():
                comp_uid = self._ready_queue.popleft()

                # Skip if already processed
                if comp_uid in self.resolved or comp_uid in self.running:
                    continue

                # Check if future is already done (could be cancelled/failed)
                if self.components[comp_uid]["future"].done():
                    self.resolved.add(comp_uid)
                    self._notify_dependents(comp_uid)
                    continue

                # Verify dependencies are still met
                dependencies = self.dependencies[comp_uid]
                dep_futures = [
                    self.components[dep["uid"]]["future"] for dep in dependencies
                ]

                failed_deps = []
                cancelled_deps = []

                for fut in dep_futures:
                    if fut.cancelled():
                        cancelled_deps.append(fut)
                    elif fut.exception() is not None:
                        failed_deps.append(fut.exception())

                # Handle dependency issues
                if cancelled_deps or failed_deps:
                    comp_desc = self.components[comp_uid]["description"]

                    if cancelled_deps:
                        logger.info(
                            f"Cancelling {comp_desc['name']} "
                            "due to cancelled dependencies"
                        )
                        self.handle_component_cancellation(
                            comp_desc, self.components[comp_uid]["future"]
                        )
                    else:  # failed_deps
                        chained_exception = (
                            self._create_dependency_failure_exception(
                                comp_desc, failed_deps
                            )
                        )
                        self.handle_component_failure(
                            comp_desc,
                            self.components[comp_uid]["future"],
                            chained_exception,
                        )

                    # Common cleanup
                    self.resolved.add(comp_uid)
                    self._notify_dependents(comp_uid)
                    continue

                # Handle data dependencies for tasks
                comp_desc = self.components[comp_uid]["description"]
                if self.components[comp_uid]["type"] == TASK:
                    # Data linking runs only on backends that declare it;
                    # implicit and explicit linking are independent.
                    if self._links_implicit_data and self.implicit_data_mode:
                        for dep in dependencies:
                            dep_desc = self.components[dep["uid"]]["description"]

                            # Link implicit data dependencies
                            if not dep_desc["metadata"].get("output_files"):
                                logger.debug(
                                    "Linking implicit file(s): from "
                                    f"{dep_desc['name']} to {comp_desc['name']}"
                                )
                                self.backend.link_implicit_data_deps(
                                    dep_desc, comp_desc
                                )

                    if self._links_explicit_data:
                        explicit_files_to_stage = []

                        for dep in dependencies:
                            dep_desc = self.components[dep["uid"]]["description"]

                            # Link explicit data dependencies
                            for output_file in dep_desc["metadata"]["output_files"]:
                                if (
                                    output_file
                                    in comp_desc["metadata"]["input_files"]
                                ):
                                    logger.debug(
                                        f"Linking explicit file ({output_file}) "
                                        f"from {dep_desc['name']} "
                                        f"to {comp_desc['name']}"
                                    )
                                    data_dep = self.backend.link_explicit_data_deps(
                                        src_task=dep_desc,
                                        dst_task=comp_desc,
                                        file_name=output_file,
                                    )
                                    explicit_files_to_stage.append(data_dep)

                        # Input staging data dependencies
                        dependency_output_files = self._get_dependency_output_files(
                            dependencies
                        )
                        staged_targets = {
                            Path(item["target"]).name
                            for item in explicit_files_to_stage
                        }

                        for input_file in comp_desc["metadata"]["input_files"]:
                            input_basename = Path(input_file).name
                            if (
                                input_basename not in staged_targets
                                and input_basename not in dependency_output_files
                            ):
                                logger.debug(
                                    f"Staging {input_file} "
                                    f"to {comp_desc['name']} work dir"
                                )
                                data_dep = self.backend.link_explicit_data_deps(
                                    src_task=None,
                                    dst_task=comp_desc,
                                    file_name=input_basename,
                                    file_path=input_file,
                                )
                                explicit_files_to_stage.append(data_dep)

                try:
                    # Update the component description with resolved values
                    (
                        comp_desc["args"],
                        comp_desc["kwargs"],
                    ) = await self._extract_dependency_values(comp_desc)
                except Exception as e:
                    logger.error(
                        f"Failed to resolve future for {comp_desc['name']}: {e}"
                    )
                    self.handle_component_failure(
                        comp_desc, self.components[comp_uid]["future"], e
                    )
                    self.resolved.add(comp_uid)
                    self._notify_dependents(comp_uid)
                    continue
                except asyncio.CancelledError:
                    # Only a cancelled future can raise this here:
                    # _extract_dependency_values never suspends, so the run
                    # task itself cannot be cancelled inside it. Cancel the
                    # dependent instead of letting it stop the run loop.
                    logger.info(f"Cancelling {comp_desc['name']}")
                    self.handle_component_cancellation(
                        comp_desc, self.components[comp_uid]["future"]
                    )
                    self.resolved.add(comp_uid)
                    self._notify_dependents(comp_uid)
                    continue

                to_submit.append(comp_desc)
                res_deps = [dep["name"] for dep in dependencies]
                msg = f"Ready to submit: {comp_desc['name']}"
                msg += f" with resolved dependencies: {res_deps}"
                logger.debug(msg)

                if self.components[comp_uid]["type"] == TASK:
                    self._emit(
                        "TaskQueued",
                        task_id=comp_uid,
                        workflow_id=comp_desc.get("workflow_id"),
                        attributes={
                            "executable": comp_desc["_tel_executable"],
                            "task_type": comp_desc["_tel_task_type"],
                        },
                    )

            # Submit ready components
            if to_submit:
                await self.submit(to_submit)
                for comp_desc in to_submit:
                    comp_uid = comp_desc["uid"]
                    self.running.add(comp_uid)
                    self.resolved.add(comp_uid)

            # Check for completed components and update dependency tracking
            completed_components = []
            for comp_uid in list(self.running):
                if self.components[comp_uid]["future"].done():
                    completed_components.append(comp_uid)
                    self.running.discard(comp_uid)

            # Notify dependents of completed components
            for comp_uid in completed_components:
                self._notify_dependents(comp_uid)

            # Signal changes
            if completed_components:
                self._component_change_event.set()

            # If nothing is ready/running, wait for changes or shutdown
            if not self._ready_queue and not to_submit and not completed_components:
                try:
                    # Create tasks for event waiting
                    event_task = asyncio.create_task(
                        self._component_change_event.wait(),
                        name="component-event-task",
                    )
                    shutdown_task = asyncio.create_task(
                        self._shutdown_event.wait(), name="shutdown-event-task"
                    )

                    done, pending = await asyncio.wait(
                        [event_task, shutdown_task],
                        return_when=asyncio.FIRST_COMPLETED,
                        timeout=1.0,
                    )
                    # Cancel any pending tasks to clean up
                    for task in pending:
                        task.cancel()
                    # Clear component change event if it was set
                    if event_task in done:
                        self._component_change_event.clear()
                except asyncio.CancelledError:
                    # If we get cancelled, make sure to clean up our tasks
                    for task in [event_task, shutdown_task]:
                        if not task.done():
                            task.cancel()
                    raise
            else:
                await asyncio.sleep(0.01)

        except asyncio.CancelledError:
            logger.debug("Run component stopped")
            break
        except Exception as e:
            logger.exception(f"Error in run loop: {e}")
            await self._handle_shutdown_signal(signal.SIGUSR1, source="internal")
            break

async def submit(self, objects: list) -> None:
    """Manages asynchronous submission of tasks/blocks to the execution backend.

    Retrieves and submit ready tasks and blocks for execution. Separates incoming
    objects into tasks and blocks based on their UID pattern, and dispatches each
    to the appropriate backend method.

    Submission Process:
        1. Receive batch of objects
        2. Filters objects into tasks and blocks
        3. Submits tasks via `backend.submit_tasks`
        4. Submits blocks via `_submit_blocks` asynchronously

    Args:
        objects: Tasks and blocks are identified by `uid` field content

    Returns:
        None

    Raises:
        Exception: If an unexpected error occurs during submission
    """
    try:
        # Separate tasks and blocks
        tasks = [t for t in objects if t and BLOCK not in t["uid"]]
        blocks = [b for b in objects if b and TASK not in b["uid"]]

        logger.info(f"Submitting {len(objects)} tasks/blocks for execution")
        logger.debug(f"Submitting: {[b['name'] for b in objects]}")

        if tasks:
            if self._telemetry is not None:
                now = time.time()
                for t in tasks:
                    self._task_submit_times[t["uid"]] = now
                    self._emit(
                        "TaskSubmitted",
                        task_id=t["uid"],
                        workflow_id=t.get("workflow_id"),
                        event_time=now,
                        attributes={
                            "executable": t["_tel_executable"],
                            "task_type": t["_tel_task_type"],
                        },
                    )

            if not any(t.get("target_backend") for t in tasks):
                # Fast path: all tasks go to the default backend
                await self.backend.submit_tasks(tasks)
            else:
                by_backend: dict = {}
                for t in tasks:
                    by_backend.setdefault(
                        t.get("target_backend") or self._default_backend_name, []
                    ).append(t)
                # Submit to all target backends in parallel
                await asyncio.gather(
                    *(
                        self._backends[b_name].submit_tasks(b_tasks)
                        for b_name, b_tasks in by_backend.items()
                    )
                )
        if blocks:
            await self._submit_blocks(blocks)
    except Exception as e:
        logger.exception(f"Error in submit component: {e}")
        raise

async def execute_block(
    self, block_fut: asyncio.Future, func: Callable, *args: Any, **kwargs: Any
) -> None:
    """Executes a block function and sets its result on the associated future.

    Calls the given function with provided args, awaiting it if it's a coroutine,
    or running it in the executor otherwise. On completion, updates the `block_fut`
    with the result or exception.

    Args:
        block_fut (asyncio.Future): Future to store the result or exception.
        func (Callable): Function or coroutine function to execute.
        *args: Positional arguments to pass to the function.
        **kwargs: Keyword arguments to pass to the function.

    Returns:
        None
    """

    block_uid = block_fut.block["uid"]
    token = self._workflow_id_ctx.set(block_uid)
    try:
        scope = (
            self._telemetry.span_scope(
                "block", {"asyncflow.workflow_id": block_uid}
            )
            if self._telemetry is not None
            else nullcontext()
        )
        with scope:
            result = await func(*args, **kwargs)
        if not block_fut.done():
            self.handle_component_success(
                {"uid": block_uid, "return_value": result}, block_fut
            )
    except Exception as e:
        if not block_fut.done():
            self.handle_component_failure(block_fut.block, block_fut, e)
    except asyncio.CancelledError:
        # The block body raised CancelledError, e.g. by awaiting a member task
        # that was cancelled. Cancel the block so it and its dependents don't
        # hang. When the engine itself cancelled the block (block_fut.cancel()
        # or shutdown), block_fut is already done and this is a no-op.
        self.handle_component_cancellation(block_fut.block, block_fut)
        raise
    finally:
        self._workflow_id_ctx.reset(token)

def handle_component_success(
    self, component: dict, component_fut: asyncio.Future
) -> None:
    """Handles successful completion of a component (task or block).

    Sets the result of the component's future: the function's return value for
    function tasks and blocks, or stdout for executable and prompt tasks. Logs a
    warning if the future is already resolved.

    Args:
        component (dict): Completed component descriptor containing
            - 'uid' (str): Unique component identifier
            - 'return_value' / 'stdout': Result of the execution
        component_fut (asyncio.Future): Future to set the result on.
    """
    internal_component = self.components[component["uid"]]["description"]

    if not component_fut.done():
        if internal_component[FUNCTION]:
            component_fut.set_result(component["return_value"])
        else:
            component_fut.set_result(component["stdout"])
        component_fut.state = "DONE"
    else:
        logger.warning(
            "Attempted to handle an already finished component "
            f'"{component["uid"]}"'
        )

def handle_component_failure(
    self,
    component: dict,
    component_fut: asyncio.Future,
    override_error_message: Union[str, Exception] = None,
) -> None:
    """Handles failure of a component (task or block) and sets the exception.

    Marks the component's future as failed by setting an exception derived from
    either a provided override error or the component's own recorded
    error/stderr. Logs a warning if the future is already resolved.

    Args:
        component (dict): Dictionary with component details, including:
            - 'uid' (str): Unique component identifier
            - 'exception' or 'stderr': Error information from execution
        component_fut (asyncio.Future): Future to mark as failed.
        override_error_message (Union[str, Exception], optional): Custom
            error message or exception to set instead of the recorded error.

    Returns:
        None
    """
    if component_fut.done():
        logger.warning(
            f'Attempted to handle an already failed component "{component["uid"]}"'
        )
        return

    # Determine the appropriate exception to set
    if override_error_message is not None:
        # If it's already an exception (like DependencyFailureError),
        # use it directly
        if isinstance(override_error_message, Exception):
            exception = override_error_message
        else:
            # If it's a string, wrap it in RuntimeError
            exception = RuntimeError(str(override_error_message))
    else:
        # Use the component's original exception or stderr
        original_error = (
            component.get("exception")
            or component.get("stderr")
            or "failed with unknown error"
        )

        # Ensure we have an Exception object
        if isinstance(original_error, Exception):
            exception = original_error
        else:
            # If it's a string (stderr) or any other type, wrap it in RuntimeError
            exception = RuntimeError(str(original_error))

    component_fut.set_exception(exception)
    component_fut.state = "FAILED"

def handle_component_cancellation(
    self, component: dict, component_fut: asyncio.Future
):
    """Handle cancellation of a component (task or block)."""
    if component_fut.done():
        return  # already resolved — idempotent, nothing to do

    # cancel() would send a task running on the backend back to the backend,
    # which reports here. A future that the engine did not create has only
    # cancel().
    cancel = getattr(component_fut, "_cancel_locally", component_fut.cancel)
    result = cancel()
    component_fut.state = "CANCELLED"
    return result

@typeguard.typechecked
def task_callbacks(
    self, task: Any, state: str, service_callback: Optional[Callable] = None
) -> None:
    """Processes task state changes and invokes appropriate handlers.

    Handles state transitions for tasks, updates their futures, and triggers
    relevant state-specific handlers. Supports optional service-specific
    callbacks for extended functionality.

    Args:
        task (Union[dict, object]): Task object or dictionary containing task
            state information.
        state (str): New state of the task.
        service_callback (Optional[Callable], optional): Callback function
            for service tasks. Must be daemon-threaded to avoid blocking.
            Defaults to None.

    Returns:
        None

    State Transitions:
        - DONE: Calls handle_component_success
        - RUNNING: Marks future as running
        - CANCELED: Cancels the future
        - FAILED: Calls handle_component_failure

    Logging:
        - Debug: Non-relevant state received
        - Info: Task state changes
        - Warning: Unknown task received

    Example:
        ::

            def service_ready_callback(future, task, state):
                def wait_and_set():
                    try:
                        # Synchronous operation
                        future.set_result(info)
                    except Exception as e:
                        future.set_exception(e)
                threading.Thread(target=wait_and_set, daemon=True).start()
    """
    if (
        state not in self.task_states_map.terminal_states
        and state != self.task_states_map.RUNNING
    ):
        logger.debug(f"Non-relevant task state received: {state}. Skipping state.")
        return

    task_obj = task
    if isinstance(task, dict):
        task_dct = task
    else:
        task_dct = task.as_dict()

    if task_dct["uid"] not in self.components:
        logger.warning(
            f"Received an unknown task and will skip it: {task_dct['uid']}"
        )
        return

    comp = self.components[task_dct["uid"]]
    task_fut = comp["future"]
    comp_desc = comp["description"]
    logger.info(f"{task_dct['uid']} is in {state} state")

    if service_callback:
        # service tasks are marked done by a backend specific
        # mechanism that are provided during the callbacks only
        service_callback(task_fut, task_obj, state)

    uid = task_dct["uid"]

    tel_attrs = {
        "executable": comp_desc["_tel_executable"],
        "task_type": comp_desc["_tel_task_type"],
    }

    if state == self.task_states_map.DONE:
        self.handle_component_success(task_dct, task_fut)
        if self._telemetry is not None:
            now = time.time()
            start = self._task_start_times.pop(
                uid, self._task_submit_times.get(uid, now)
            )
            self._task_submit_times.pop(uid, None)
            self._emit(
                "TaskCompleted",
                task_id=uid,
                workflow_id=task_dct.get("workflow_id"),
                event_time=now,
                duration_seconds=now - start,
                attributes=tel_attrs,
            )
    elif state == self.task_states_map.RUNNING:
        # NOTE: with asyncio future the running state is
        # implicit: when a coroutine that awaits the future
        # is scheduled and started by the event loop, that's
        # when the "work" is running.
        task_fut.state = "RUNNING"
        if self._telemetry is not None:
            now = time.time()
            self._task_start_times[uid] = now
            self._emit(
                "TaskStarted",
                task_id=uid,
                workflow_id=task_dct.get("workflow_id"),
                event_time=now,
                attributes=tel_attrs,
            )
    elif state == self.task_states_map.CANCELED:
        self.handle_component_cancellation(task_dct, task_fut)
        if self._telemetry is not None:
            now = time.time()
            start = self._task_start_times.pop(
                uid, self._task_submit_times.get(uid, now)
            )
            self._task_submit_times.pop(uid, None)
            self._emit(
                "TaskCanceled",
                task_id=uid,
                workflow_id=task_dct.get("workflow_id"),
                event_time=now,
                duration_seconds=now - start,
                attributes=tel_attrs,
            )
    elif state == self.task_states_map.FAILED:
        self.handle_component_failure(task_dct, task_fut)
        if self._telemetry is not None:
            now = time.time()
            start = self._task_start_times.pop(
                uid, self._task_submit_times.get(uid, now)
            )
            self._task_submit_times.pop(uid, None)
            self._emit(
                "TaskFailed",
                task_id=uid,
                workflow_id=task_dct.get("workflow_id"),
                event_time=now,
                duration_seconds=now - start,
                error_type="unknown",
                attributes=tel_attrs,
            )

async def shutdown(self, skip_execution_backend: bool = False) -> None:
    """Internal implementation of asynchronous shutdown for the workflow manager.

    This method performs the following steps:
        1. Sets the shutdown event to signal components to exit
        2. Cancels background tasks responsible for running and
        submitting workflows.
        3. Waits for the cancellation and completion of these tasks,
        with a timeout of 5 seconds.
        4. Cancel workflow tasks, including those that have not
        registered yet.
        5. Logs a warning if the tasks do not complete within the timeout
        period.
        6. Shuts down the backend using an executor to avoid blocking the
        event loop.

    Args:
        skip_execution_backend (bool): If True, skips the shutdown of the
            execution backend.

    Returns:
        None

    Raises:
        asyncio.TimeoutError: If the background tasks do not complete
            within the timeout period.
        asyncio.CancelledError: If the shutdown is cancelled before
            completion.
    """
    logger.info("Initiating shutdown")
    # Signal components to exit
    self._shutdown_event.set()

    # cancel workflow futures (tasks and blocks)
    for comp in self.components.values():
        future = comp["future"]
        if not future.done():
            self.handle_component_cancellation(comp["description"], future)

    # cancel the futures of components that have not registered yet: an
    # executable or prompt task whose body is still running, or a dependent
    # waiting for its dependencies to register
    for future in list(self._pending_registrations):
        registration = getattr(future, "_registration", None)
        if registration is not None:
            registration.cancel()
        if not future.done():
            self.handle_component_cancellation({"uid": future.uid}, future)
    # A registration cancelled before it started never runs its finally
    self._pending_registrations.clear()

    # Cancel internal components task
    if not self._run_task.done():
        logger.debug(f"Shutting down run component")
        self._run_task.cancel()

    # Wait for internal components shutdown to complete
    try:
        await asyncio.wait_for(self._run_task, timeout=5.0)
    except asyncio.TimeoutError:
        logger.warning("Timeout waiting for internal components to shutdown")
    except asyncio.CancelledError:
        logger.warning("Internal components shutdown cancelled")

    # Shutdown all registered execution backends
    if not skip_execution_backend and self._backends:
        await asyncio.gather(*[b.shutdown() for b in self._backends.values()])
        self._clear_internal_records()
        logger.debug("Shutting down execution backend completed")
    else:
        logger.warning("Skipping execution backend shutdown as requested")

    logger.info("Shutdown completed for all components.")