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_task_cancellation(
                            comp_desc, self.components[comp_uid]["future"]
                        )
                    else:  # failed_deps
                        chained_exception = (
                            self._create_dependency_failure_exception(
                                comp_desc, failed_deps
                            )
                        )
                        self.handle_task_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:
                    explicit_files_to_stage = []

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

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

                        # 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_task_failure(
                        comp_desc, self.components[comp_uid]["future"], e
                    )
                    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,
                        backend=comp_desc.get("target_backend"),
                        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.append(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.remove(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"],
                        backend=t.get("target_backend"),
                        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():
            block_fut.set_result(result)
    except Exception as e:
        if not block_fut.done():
            block_fut.set_exception(e)
    finally:
        self._workflow_id_ctx.reset(token)

def handle_task_success(self, task: dict, task_fut: asyncio.Future) -> None:
    """Handles successful task completion and updates the associated future.

    Sets the result of the task's future based on whether the task was a function
    or a shell command. Raises an error if the future is already resolved.

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

    if not task_fut.done():
        if internal_task[FUNCTION]:
            task_fut.set_result(task["return_value"])
        else:
            task_fut.set_result(task["stdout"])
    else:
        logger.warning(
            f'Attempted to handle an already finished task "{task["uid"]}"'
        )

def handle_task_failure(
    self,
    task: dict,
    task_fut: asyncio.Future,
    override_error_message: Union[str, Exception] = None,
) -> None:
    """Handles task failure and sets the appropriate exception on the future.

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

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

    Returns:
        None
    """
    if task_fut.done():
        logger.warning(
            f'Attempted to handle an already failed task "{task["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 task's original exception or stderr
        original_error = (
            task.get("exception")
            or task.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))

    task_fut.set_exception(exception)

def handle_task_cancellation(self, task: dict, task_fut: asyncio.Future):
    """Handle task cancellation."""
    if task_fut.done():
        logger.warning(
            f'Attempted to handle an already cancelled task "{task["uid"]}"'
        )
        return

    # Restore original cancel method
    task_fut.cancel = task_fut.original_cancel
    return task_fut.cancel()

@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_task_success
        - RUNNING: Marks future as running
        - CANCELED: Cancels the future
        - FAILED: Calls handle_task_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_task_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.
        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_task_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_task_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.
        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"]
        comp_desc = comp["description"]
        if not future.done():
            self.handle_task_cancellation(comp_desc, future)

    # 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.")