Task

Task definition and configuration for DAG nodes.

Task definitions for DAG execution, including sync, async, stream, and process tasks.

AsyncFunctionShutdownTask

Bases: AsyncFunctionTask, AsyncShutdownTask

Async function task with an async shutdown hook.

Source code in shutils/dag/task.py

class AsyncFunctionShutdownTask(AsyncFunctionTask, AsyncShutdownTask):
    """Async function task with an async shutdown hook."""

    def __init__(
        self, shutdown_callable: AsyncShutdownCallableProtocol, config: TaskConfig | None = None, name: str = ""
    ):
        """Initialize with an async-shutdown-callable function.

        Args:
            shutdown_callable: An async callable that also has a shutdown() method.
            config: Task execution configuration.
            name: Optional task name.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(shutdown_callable, config, name)

    async def shutdown(self):
        """Delegate shutdown to the wrapped async callable."""
        await self._func.shutdown()

__init__(shutdown_callable, config=None, name='')

Initialize with an async-shutdown-callable function.

Parameters:

Name	Type	Description	Default
shutdown_callable	AsyncShutdownCallableProtocol	An async callable that also has a shutdown() method.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Source code in shutils/dag/task.py

def __init__(
    self, shutdown_callable: AsyncShutdownCallableProtocol, config: TaskConfig | None = None, name: str = ""
):
    """Initialize with an async-shutdown-callable function.

    Args:
        shutdown_callable: An async callable that also has a shutdown() method.
        config: Task execution configuration.
        name: Optional task name.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(shutdown_callable, config, name)

shutdown() async

Delegate shutdown to the wrapped async callable.

Source code in shutils/dag/task.py

async def shutdown(self):
    """Delegate shutdown to the wrapped async callable."""
    await self._func.shutdown()

AsyncFunctionTask

Bases: AsyncTask

Async task that wraps a simple async function.

Source code in shutils/dag/task.py

class AsyncFunctionTask(AsyncTask):
    """Async task that wraps a simple async function."""

    def __init__(
        self,
        func: Callable[[AsyncContext], Coroutine[Any, Any, list[AsyncContext] | AsyncContext | None]],
        config: TaskConfig | None = None,
        name: str = "",
    ):
        """Initialize an async function task.

        Args:
            func: The async function to execute.
            config: Task execution configuration.
            name: Optional task name.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(func, config, name)
        self._func = func

    async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
        """Execute the wrapped async function and return output contexts."""
        ret = await self._func(async_ctx)
        if isinstance(ret, AsyncContext):
            return [ret]
        elif ret is None:
            return []
        return ret

    def __repr__(self):
        return f"AsyncFunctionTask(func={self._func}, {TaskBase.__repr__(self)})"

__init__(func, config=None, name='')

Initialize an async function task.

Parameters:

Name	Type	Description	Default
func	Callable[[AsyncContext], Coroutine[Any, Any, list[AsyncContext] | AsyncContext | None]]	The async function to execute.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Source code in shutils/dag/task.py

def __init__(
    self,
    func: Callable[[AsyncContext], Coroutine[Any, Any, list[AsyncContext] | AsyncContext | None]],
    config: TaskConfig | None = None,
    name: str = "",
):
    """Initialize an async function task.

    Args:
        func: The async function to execute.
        config: Task execution configuration.
        name: Optional task name.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(func, config, name)
    self._func = func

call(async_ctx, env) async

Execute the wrapped async function and return output contexts.

Source code in shutils/dag/task.py

async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
    """Execute the wrapped async function and return output contexts."""
    ret = await self._func(async_ctx)
    if isinstance(ret, AsyncContext):
        return [ret]
    elif ret is None:
        return []
    return ret

AsyncLoopTask

Bases: AsyncStreamTask

Async stream task that automatically re-enqueues until the generator exhausts.

Source code in shutils/dag/task.py

class AsyncLoopTask(AsyncStreamTask):
    """Async stream task that automatically re-enqueues until the generator exhausts."""

    async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
        """Execute one iteration and append a LoopContext to continue the loop."""
        need_create_loop_context = not isinstance(async_ctx.context, LoopContext)
        ret = await super().call(async_ctx, env)
        if ret:
            if need_create_loop_context:
                ret.append(LoopContext(async_ctx.context._runtime, self).async_context)
            else:
                ret.append(async_ctx)
        elif not need_create_loop_context:
            await async_ctx.destory()
        return ret

call(async_ctx, env) async

Execute one iteration and append a LoopContext to continue the loop.

Source code in shutils/dag/task.py

async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
    """Execute one iteration and append a LoopContext to continue the loop."""
    need_create_loop_context = not isinstance(async_ctx.context, LoopContext)
    ret = await super().call(async_ctx, env)
    if ret:
        if need_create_loop_context:
            ret.append(LoopContext(async_ctx.context._runtime, self).async_context)
        else:
            ret.append(async_ctx)
    elif not need_create_loop_context:
        await async_ctx.destory()
    return ret

AsyncRouterTask

Bases: AsyncTask

Async task that routes contexts to specific downstream tasks.

Source code in shutils/dag/task.py

class AsyncRouterTask(AsyncTask):
    """Async task that routes contexts to specific downstream tasks."""

    def __init__(
        self,
        func: Callable[[AsyncContext], Coroutine[Any, Any, str | TaskBase | list[str | TaskBase]]],
        config: TaskConfig | None = None,
        name: str = "",
    ):
        """Initialize an async router task.

        Args:
            func: An async function returning task name(s) or task object(s) to route to.
            config: Task execution configuration.
            name: Optional task name.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(None, config, name)
        self._func = func
        self._mask_task_cache: dict[tuple[TaskBase, ...], set[TaskBase]] = {}

    async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
        """Route the context to specified downstream tasks and mask bypass tasks."""
        route_tasks = await self._func(async_ctx)
        if not isinstance(route_tasks, list):
            route_tasks = [route_tasks]

        real_route_tasks: list[TaskBase] = []
        for task_name in route_tasks:
            if type(task_name) is str:
                task = env.dag.tasks.get(task_name, None)
                if not task:
                    logger.error(f"task {task} not found in dag tasks")
                    await async_ctx.destory()
                    return []
            else:
                task = task_name
            if task not in self.downstream_tasks:
                logger.error(f"task {task} not found in downstream tasks")
                await async_ctx.destory()
                return []

            real_route_tasks.append(task)
        route_task_set = set(real_route_tasks)
        route_task_tuple = tuple(real_route_tasks)
        if route_task_tuple in self._mask_task_cache:
            mask_task_set = self._mask_task_cache[route_task_tuple]
        else:
            route_downstream_tasks = env.dag._get_all_downstream_tasks(route_task_set, True)
            noroute_downstream_tasks = env.dag._get_all_downstream_tasks(self.downstream_tasks - route_task_set, True)
            mask_task_set = noroute_downstream_tasks - route_downstream_tasks

        for task in mask_task_set:
            await async_ctx.complete(task)

        return [async_ctx]

__init__(func, config=None, name='')

Initialize an async router task.

Parameters:

Name	Type	Description	Default
func	Callable[[AsyncContext], Coroutine[Any, Any, str | TaskBase | list[str | TaskBase]]]	An async function returning task name(s) or task object(s) to route to.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Source code in shutils/dag/task.py

def __init__(
    self,
    func: Callable[[AsyncContext], Coroutine[Any, Any, str | TaskBase | list[str | TaskBase]]],
    config: TaskConfig | None = None,
    name: str = "",
):
    """Initialize an async router task.

    Args:
        func: An async function returning task name(s) or task object(s) to route to.
        config: Task execution configuration.
        name: Optional task name.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(None, config, name)
    self._func = func
    self._mask_task_cache: dict[tuple[TaskBase, ...], set[TaskBase]] = {}

call(async_ctx, env) async

Route the context to specified downstream tasks and mask bypass tasks.

Source code in shutils/dag/task.py

async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
    """Route the context to specified downstream tasks and mask bypass tasks."""
    route_tasks = await self._func(async_ctx)
    if not isinstance(route_tasks, list):
        route_tasks = [route_tasks]

    real_route_tasks: list[TaskBase] = []
    for task_name in route_tasks:
        if type(task_name) is str:
            task = env.dag.tasks.get(task_name, None)
            if not task:
                logger.error(f"task {task} not found in dag tasks")
                await async_ctx.destory()
                return []
        else:
            task = task_name
        if task not in self.downstream_tasks:
            logger.error(f"task {task} not found in downstream tasks")
            await async_ctx.destory()
            return []

        real_route_tasks.append(task)
    route_task_set = set(real_route_tasks)
    route_task_tuple = tuple(real_route_tasks)
    if route_task_tuple in self._mask_task_cache:
        mask_task_set = self._mask_task_cache[route_task_tuple]
    else:
        route_downstream_tasks = env.dag._get_all_downstream_tasks(route_task_set, True)
        noroute_downstream_tasks = env.dag._get_all_downstream_tasks(self.downstream_tasks - route_task_set, True)
        mask_task_set = noroute_downstream_tasks - route_downstream_tasks

    for task in mask_task_set:
        await async_ctx.complete(task)

    return [async_ctx]

AsyncServiceTask

Bases: AsyncTask, LongRunningTask, AsyncShutdownTask

Async task backed by a long-running coroutine service.

Source code in shutils/dag/task.py

class AsyncServiceTask(AsyncTask, LongRunningTask, AsyncShutdownTask):
    """Async task backed by a long-running coroutine service."""

    def __init__(
        self,
        func: Callable[[asyncio.Queue[tuple[AsyncContext, asyncio.Future]]], Coroutine[Any, Any, None]],
        config: TaskConfig | None = None,
        name: str = "",
    ):
        """Initialize an async service task.

        Args:
            func: An async function that consumes from a queue of (context, future) pairs.
            config: Task execution configuration.
            name: Optional task name.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(func, config, name)
        self.__input_queue = asyncio.Queue()
        self.__task = None
        self._func = func

    def _task_down_callback(self, task: asyncio.Task):
        """Callback when the service task finishes unexpectedly."""
        if self.__future.done():
            return

        try:
            result = task.result()
            self.__future.set_result(result)
        except Exception as e:
            logger.error(f"[AsyncServiceTask]: task exit with exception: {e}")
            self.__future.set_exception(e)

    async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
        """Send context to the service coroutine and await the result."""
        if self.__task is None:
            self.__task = asyncio.create_task(self._func(self.__input_queue))
            self.__task.add_done_callback(self._task_down_callback)
        elif self.__task.done():
            logger.error("[AsyncServiceTask]: task exit unexpectedly")
            raise RuntimeError("task is not alive")
        self.__future = asyncio.Future()
        await self.__input_queue.put((async_ctx, self.__future))
        result = await self.__future

        if isinstance(result, AsyncContext):
            return [result]
        elif isinstance(result, list):
            return result
        return []

    def __repr__(self):
        return f"AsyncServiceTask(func={self._func}, {TaskBase.__repr__(self)})"

    async def shutdown(self):
        """Send a stop signal to the service coroutine and await its completion."""
        if self.__task and not self.__task.done():
            await self.__input_queue.put((StopContext().async_context, asyncio.Future()))
            await self.__task
            self.__task = None

__init__(func, config=None, name='')

Initialize an async service task.

Parameters:

Name	Type	Description	Default
func	Callable[[Queue[tuple[AsyncContext, Future]]], Coroutine[Any, Any, None]]	An async function that consumes from a queue of (context, future) pairs.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Source code in shutils/dag/task.py

def __init__(
    self,
    func: Callable[[asyncio.Queue[tuple[AsyncContext, asyncio.Future]]], Coroutine[Any, Any, None]],
    config: TaskConfig | None = None,
    name: str = "",
):
    """Initialize an async service task.

    Args:
        func: An async function that consumes from a queue of (context, future) pairs.
        config: Task execution configuration.
        name: Optional task name.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(func, config, name)
    self.__input_queue = asyncio.Queue()
    self.__task = None
    self._func = func

call(async_ctx, env) async

Send context to the service coroutine and await the result.

Source code in shutils/dag/task.py

async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
    """Send context to the service coroutine and await the result."""
    if self.__task is None:
        self.__task = asyncio.create_task(self._func(self.__input_queue))
        self.__task.add_done_callback(self._task_down_callback)
    elif self.__task.done():
        logger.error("[AsyncServiceTask]: task exit unexpectedly")
        raise RuntimeError("task is not alive")
    self.__future = asyncio.Future()
    await self.__input_queue.put((async_ctx, self.__future))
    result = await self.__future

    if isinstance(result, AsyncContext):
        return [result]
    elif isinstance(result, list):
        return result
    return []

shutdown() async

Send a stop signal to the service coroutine and await its completion.

Source code in shutils/dag/task.py

async def shutdown(self):
    """Send a stop signal to the service coroutine and await its completion."""
    if self.__task and not self.__task.done():
        await self.__input_queue.put((StopContext().async_context, asyncio.Future()))
        await self.__task
        self.__task = None

AsyncShutdownCallableProtocol

Bases: Protocol

Protocol for async callables that support a shutdown hook.

Source code in shutils/dag/task.py

class AsyncShutdownCallableProtocol(Protocol):
    """Protocol for async callables that support a shutdown hook."""

    async def __call__(self, context: AsyncContext) -> list[AsyncContext] | AsyncContext | None: ...
    async def shutdown(self) -> None: ...

AsyncShutdownTask

Bases: ABC

Abstract base for tasks that require an async shutdown hook.

Source code in shutils/dag/task.py

class AsyncShutdownTask(ABC):
    """Abstract base for tasks that require an async shutdown hook."""

    @abstractmethod
    async def shutdown(self):
        """Clean up resources asynchronously when the executor stops."""

shutdown() abstractmethod async

Clean up resources asynchronously when the executor stops.

Source code in shutils/dag/task.py

@abstractmethod
async def shutdown(self):
    """Clean up resources asynchronously when the executor stops."""

AsyncStreamTask

Bases: AsyncTask, AsyncShutdownTask

Async task backed by an async generator that streams multiple outputs.

Source code in shutils/dag/task.py

class AsyncStreamTask(AsyncTask, AsyncShutdownTask):
    """Async task backed by an async generator that streams multiple outputs."""

    def __init__(
        self,
        func: Callable[[], AsyncGenerator[AsyncContext | list[AsyncContext] | None, AsyncContext]],
        config: TaskConfig | None = None,
        name: str = "",
    ):
        """Initialize an async stream task.

        Args:
            func: An async generator function that yields output contexts.
            config: Task execution configuration.
            name: Optional task name.

        Raises:
            ValueError: If func does not return an async generator.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(func, config, name)
        self.__generator = func()
        if not isasyncgen(self.__generator):
            raise ValueError("func must be a async generator function")
        self.__activate_generator = False

    async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
        """Send the context to the async generator and return yielded outputs."""
        if not self.__activate_generator:
            self.__activate_generator = True
            await anext(self.__generator)
        try:
            ret = await self.__generator.asend(async_ctx)
            if ret is None:
                return []
            if isinstance(ret, AsyncContext):
                return [ret]
            return ret
        except StopAsyncIteration:
            return []

    def __repr__(self):
        return f"{self.__class__.__name__}(func={self.__generator}, {TaskBase.__repr__(self)})"

    async def shutdown(self):
        """Close the underlying async generator."""
        with contextlib.suppress(GeneratorExit):
            await self.__generator.aclose()

__init__(func, config=None, name='')

Initialize an async stream task.

Parameters:

Name	Type	Description	Default
func	Callable[[], AsyncGenerator[AsyncContext | list[AsyncContext] | None, AsyncContext]]	An async generator function that yields output contexts.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Raises:

Type	Description
ValueError	If func does not return an async generator.

Source code in shutils/dag/task.py

def __init__(
    self,
    func: Callable[[], AsyncGenerator[AsyncContext | list[AsyncContext] | None, AsyncContext]],
    config: TaskConfig | None = None,
    name: str = "",
):
    """Initialize an async stream task.

    Args:
        func: An async generator function that yields output contexts.
        config: Task execution configuration.
        name: Optional task name.

    Raises:
        ValueError: If func does not return an async generator.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(func, config, name)
    self.__generator = func()
    if not isasyncgen(self.__generator):
        raise ValueError("func must be a async generator function")
    self.__activate_generator = False

call(async_ctx, env) async

Send the context to the async generator and return yielded outputs.

Source code in shutils/dag/task.py

async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
    """Send the context to the async generator and return yielded outputs."""
    if not self.__activate_generator:
        self.__activate_generator = True
        await anext(self.__generator)
    try:
        ret = await self.__generator.asend(async_ctx)
        if ret is None:
            return []
        if isinstance(ret, AsyncContext):
            return [ret]
        return ret
    except StopAsyncIteration:
        return []

shutdown() async

Close the underlying async generator.

Source code in shutils/dag/task.py

async def shutdown(self):
    """Close the underlying async generator."""
    with contextlib.suppress(GeneratorExit):
        await self.__generator.aclose()

AsyncTask

Bases: TaskBase

Base class for asynchronous tasks.

Source code in shutils/dag/task.py

class AsyncTask(TaskBase):
    """Base class for asynchronous tasks."""

    @abstractmethod
    async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
        """Execute the task asynchronously.

        Args:
            async_ctx: The async context for this execution.
            env: The runtime environment.

        Returns:
            List of output async contexts.
        """

    async def __call__(self, context: Context, env: Environment) -> list[Context]:
        ret = self.call_before(context)
        if ret:
            return [ret]
        async_ret = await self.call(context.async_context, env)
        ret = [ctx.context for ctx in async_ret]
        self.call_after(ret)
        return ret

call(async_ctx, env) abstractmethod async

Execute the task asynchronously.

Parameters:

Name	Type	Description	Default
async_ctx	AsyncContext	The async context for this execution.	required
env	Environment	The runtime environment.	required

Returns:

Type	Description
list[AsyncContext]	List of output async contexts.

Source code in shutils/dag/task.py

@abstractmethod
async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
    """Execute the task asynchronously.

    Args:
        async_ctx: The async context for this execution.
        env: The runtime environment.

    Returns:
        List of output async contexts.
    """

Environment dataclass

Runtime environment passed to tasks during execution.

Attributes:

Name	Type	Description
runtime	Runtime	The runtime counter for tracking active contexts.
process_pool	ProcessPoolExecutor | None	Optional process pool for CPU-bound task execution.
dag	DAG	The DAG this environment belongs to.

Source code in shutils/dag/task.py

@dataclass
class Environment:
    """Runtime environment passed to tasks during execution.

    Attributes:
        runtime: The runtime counter for tracking active contexts.
        process_pool: Optional process pool for CPU-bound task execution.
        dag: The DAG this environment belongs to.
    """
    runtime: Runtime
    process_pool: ProcessPoolExecutor | None
    dag: "DAG"

ForegroundSyncFunctionTask

Bases: SyncFunctionTask, ForegroundTask

Foreground variant of SyncFunctionTask.

Source code in shutils/dag/task.py

class ForegroundSyncFunctionTask(SyncFunctionTask, ForegroundTask):
    """Foreground variant of SyncFunctionTask."""

    def _foreground_marker(self): pass

ForegroundSyncLoopTask

Bases: SyncLoopTask, ForegroundTask

Foreground variant of SyncLoopTask.

Source code in shutils/dag/task.py

class ForegroundSyncLoopTask(SyncLoopTask, ForegroundTask):
    """Foreground variant of SyncLoopTask."""

    def _foreground_marker(self): ...

ForegroundSyncStreamTask

Bases: SyncStreamTask, ForegroundTask

Foreground variant of SyncStreamTask.

Source code in shutils/dag/task.py

class ForegroundSyncStreamTask(SyncStreamTask, ForegroundTask):
    """Foreground variant of SyncStreamTask."""

    def _foreground_marker(self): pass

ForegroundTask

Bases: ABC

Marker class for tasks that must run on the main event loop thread.

Source code in shutils/dag/task.py

class ForegroundTask(ABC):
    """Marker class for tasks that must run on the main event loop thread."""

    @abstractmethod
    def _foreground_marker(self): ...

LongRunningTask

Bases: ABC

Marker class for long-running tasks that maintain their own thread or coroutine.

Source code in shutils/dag/task.py

class LongRunningTask(ABC):
    """Marker class for long-running tasks that maintain their own thread or coroutine."""

    @abstractmethod
    def _long_running_marker(self): ...

ProcessTask

Bases: AsyncTask

Async task that offloads execution to a process pool.

Source code in shutils/dag/task.py

class ProcessTask(AsyncTask):
    """Async task that offloads execution to a process pool."""

    def __init__(
        self,
        func: Callable[[SyncContext], list[SyncContext] | SyncContext],
        config: TaskConfig | None = None,
        name: str = "",
    ):
        """Initialize a process task.

        Args:
            func: The sync function to run in a process pool.
            config: Task execution configuration.
            name: Optional task name.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(func, config, name)
        self._func = func

    @staticmethod
    def _func_wrapper(func: Callable[[SyncContext], list[SyncContext] | SyncContext], data: dict):
        input_context = Context(None)
        input_context._data = data
        ret = func(input_context.sync_context)
        if isinstance(ret, SyncContext):
            return ret.context._data

        output_data_list = []
        for sync_ctx in ret:
            output_data_list.append(sync_ctx.context._data)
        return output_data_list

    async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
        """Execute the function in a process pool and return output contexts.

        Args:
            async_ctx: The async context for this execution.
            env: The runtime environment providing the process pool.

        Returns:
            List of output async contexts.

        Raises:
            RuntimeError: If no process pool is configured.
        """
        if env.process_pool is None:
            raise RuntimeError("process pool is not set")
        loop = asyncio.get_running_loop()
        async with async_ctx.wlock():
            ret = await loop.run_in_executor(env.process_pool, self._func_wrapper, self._func, async_ctx.context._data)

        if isinstance(ret, dict):
            async with async_ctx.wlock():
                async_ctx.context._data = ret
            return [async_ctx]

        output_context_list = []
        for data in ret:
            output_context = await async_ctx.create()
            output_context.context._data = data
            output_context_list.append(output_context)
        await async_ctx.destory()
        return output_context_list

__init__(func, config=None, name='')

Initialize a process task.

Parameters:

Name	Type	Description	Default
func	Callable[[SyncContext], list[SyncContext] | SyncContext]	The sync function to run in a process pool.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Source code in shutils/dag/task.py

def __init__(
    self,
    func: Callable[[SyncContext], list[SyncContext] | SyncContext],
    config: TaskConfig | None = None,
    name: str = "",
):
    """Initialize a process task.

    Args:
        func: The sync function to run in a process pool.
        config: Task execution configuration.
        name: Optional task name.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(func, config, name)
    self._func = func

call(async_ctx, env) async

Execute the function in a process pool and return output contexts.

Parameters:

Name	Type	Description	Default
async_ctx	AsyncContext	The async context for this execution.	required
env	Environment	The runtime environment providing the process pool.	required

Returns:

Type	Description
list[AsyncContext]	List of output async contexts.

Raises:

Type	Description
RuntimeError	If no process pool is configured.

Source code in shutils/dag/task.py

async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
    """Execute the function in a process pool and return output contexts.

    Args:
        async_ctx: The async context for this execution.
        env: The runtime environment providing the process pool.

    Returns:
        List of output async contexts.

    Raises:
        RuntimeError: If no process pool is configured.
    """
    if env.process_pool is None:
        raise RuntimeError("process pool is not set")
    loop = asyncio.get_running_loop()
    async with async_ctx.wlock():
        ret = await loop.run_in_executor(env.process_pool, self._func_wrapper, self._func, async_ctx.context._data)

    if isinstance(ret, dict):
        async with async_ctx.wlock():
            async_ctx.context._data = ret
        return [async_ctx]

    output_context_list = []
    for data in ret:
        output_context = await async_ctx.create()
        output_context.context._data = data
        output_context_list.append(output_context)
    await async_ctx.destory()
    return output_context_list

ShutdownCallableProtocol

Bases: Protocol

Protocol for sync callables that support a shutdown hook.

Source code in shutils/dag/task.py

class ShutdownCallableProtocol(Protocol):
    """Protocol for sync callables that support a shutdown hook."""

    def __call__(self, context: SyncContext) -> list[SyncContext] | SyncContext | None: ...
    def shutdown(self) -> None: ...

ShutdownTask

Bases: ABC

Abstract base for tasks that require a sync shutdown hook.

Source code in shutils/dag/task.py

class ShutdownTask(ABC):
    """Abstract base for tasks that require a sync shutdown hook."""

    @abstractmethod
    def shutdown(self):
        """Clean up resources when the executor stops."""

shutdown() abstractmethod

Clean up resources when the executor stops.

Source code in shutils/dag/task.py

@abstractmethod
def shutdown(self):
    """Clean up resources when the executor stops."""

SinkNode

Bases: AsyncTask

Sink node that converts contexts to OutputContext for final results.

Source code in shutils/dag/task.py

class SinkNode(AsyncTask):
    """Sink node that converts contexts to OutputContext for final results."""

    def __init__(self, name: str = "#SinkNode"):
        """Initialize the sink node.

        Args:
            name: Node name, defaults to "#SinkNode".
        """
        super().__init__(None, name=name)

    async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
        """Convert the context to an OutputContext if not already one."""
        if isinstance(async_ctx.context, OutputContext):
            return [async_ctx]
        else:
            output_context = OutputContext()
            await output_context.acopy(async_ctx.context)
            await async_ctx.destory(destory_parent=True)
            return [output_context.async_context]

    async def __call__(self, context: Context, env: Environment) -> list[Context]:
        asnc_ret = await self.call(context.async_context, env)
        return [ctx.context for ctx in asnc_ret]

__init__(name='#SinkNode')

Initialize the sink node.

Parameters:

Name	Type	Description	Default
name	str	Node name, defaults to "#SinkNode".	'#SinkNode'

Source code in shutils/dag/task.py

def __init__(self, name: str = "#SinkNode"):
    """Initialize the sink node.

    Args:
        name: Node name, defaults to "#SinkNode".
    """
    super().__init__(None, name=name)

call(async_ctx, env) async

Convert the context to an OutputContext if not already one.

Source code in shutils/dag/task.py

async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
    """Convert the context to an OutputContext if not already one."""
    if isinstance(async_ctx.context, OutputContext):
        return [async_ctx]
    else:
        output_context = OutputContext()
        await output_context.acopy(async_ctx.context)
        await async_ctx.destory(destory_parent=True)
        return [output_context.async_context]

SourceNode

Bases: AsyncTask

Source node that passes through input contexts unchanged.

Source code in shutils/dag/task.py

class SourceNode(AsyncTask):
    """Source node that passes through input contexts unchanged."""

    def __init__(self, name: str = "#SourceNode"):
        """Initialize the source node.

        Args:
            name: Node name, defaults to "#SourceNode".
        """
        super().__init__(None, name=name)

    async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
        """Pass through the context unchanged."""
        return [async_ctx]

    async def __call__(self, context: Context, env: Environment) -> list[Context]:
        async_ret = await self.call(context.async_context, env)
        return [ctx.context for ctx in async_ret]

__init__(name='#SourceNode')

Initialize the source node.

Parameters:

Name	Type	Description	Default
name	str	Node name, defaults to "#SourceNode".	'#SourceNode'

Source code in shutils/dag/task.py

def __init__(self, name: str = "#SourceNode"):
    """Initialize the source node.

    Args:
        name: Node name, defaults to "#SourceNode".
    """
    super().__init__(None, name=name)

call(async_ctx, env) async

Pass through the context unchanged.

Source code in shutils/dag/task.py

async def call(self, async_ctx: AsyncContext, env: Environment) -> list[AsyncContext]:
    """Pass through the context unchanged."""
    return [async_ctx]

SyncFunctionShutdownTask

Bases: SyncFunctionTask, ShutdownTask

Sync function task with a shutdown hook.

Source code in shutils/dag/task.py

class SyncFunctionShutdownTask(SyncFunctionTask, ShutdownTask):
    """Sync function task with a shutdown hook."""

    def __init__(self, shutdown_callable: ShutdownCallableProtocol, config: TaskConfig | None = None, name: str = ""):
        """Initialize with a shutdown-callable function.

        Args:
            shutdown_callable: A callable that also has a shutdown() method.
            config: Task execution configuration.
            name: Optional task name.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(shutdown_callable, config, name)

    def shutdown(self):
        """Delegate shutdown to the wrapped callable."""
        self._func.shutdown()

__init__(shutdown_callable, config=None, name='')

Initialize with a shutdown-callable function.

Parameters:

Name	Type	Description	Default
shutdown_callable	ShutdownCallableProtocol	A callable that also has a shutdown() method.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Source code in shutils/dag/task.py

def __init__(self, shutdown_callable: ShutdownCallableProtocol, config: TaskConfig | None = None, name: str = ""):
    """Initialize with a shutdown-callable function.

    Args:
        shutdown_callable: A callable that also has a shutdown() method.
        config: Task execution configuration.
        name: Optional task name.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(shutdown_callable, config, name)

shutdown()

Delegate shutdown to the wrapped callable.

Source code in shutils/dag/task.py

def shutdown(self):
    """Delegate shutdown to the wrapped callable."""
    self._func.shutdown()

SyncFunctionTask

Bases: SyncTask

Sync task that wraps a simple function.

Source code in shutils/dag/task.py

class SyncFunctionTask(SyncTask):
    """Sync task that wraps a simple function."""

    def __init__(
        self,
        func: Callable[[SyncContext], list[SyncContext] | SyncContext | None],
        config: TaskConfig | None = None,
        name: str = "",
    ):
        """Initialize a sync function task.

        Args:
            func: The sync function to execute.
            config: Task execution configuration.
            name: Optional task name.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(func, config, name)
        self._func = func

    def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
        """Execute the wrapped function and return output contexts."""
        ret = self._func(sync_ctx)
        if isinstance(ret, SyncContext):
            return [ret]
        elif ret is None:
            return []
        return ret
        if isinstance(ret, SyncContext):
            return [ret]
        elif ret is None:
            return []
        return ret

    def __repr__(self):
        return f"SyncFunctionTask(func={self._func}, {TaskBase.__repr__(self)})"

__init__(func, config=None, name='')

Initialize a sync function task.

Parameters:

Name	Type	Description	Default
func	Callable[[SyncContext], list[SyncContext] | SyncContext | None]	The sync function to execute.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Source code in shutils/dag/task.py

def __init__(
    self,
    func: Callable[[SyncContext], list[SyncContext] | SyncContext | None],
    config: TaskConfig | None = None,
    name: str = "",
):
    """Initialize a sync function task.

    Args:
        func: The sync function to execute.
        config: Task execution configuration.
        name: Optional task name.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(func, config, name)
    self._func = func

call(sync_ctx, env)

Execute the wrapped function and return output contexts.

Source code in shutils/dag/task.py

def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
    """Execute the wrapped function and return output contexts."""
    ret = self._func(sync_ctx)
    if isinstance(ret, SyncContext):
        return [ret]
    elif ret is None:
        return []
    return ret
    if isinstance(ret, SyncContext):
        return [ret]
    elif ret is None:
        return []
    return ret

SyncLoopTask

Bases: SyncStreamTask

Sync stream task that automatically re-enqueues until the generator exhausts.

Source code in shutils/dag/task.py

class SyncLoopTask(SyncStreamTask):
    """Sync stream task that automatically re-enqueues until the generator exhausts."""

    def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
        """Execute one iteration and append a LoopContext to continue the loop."""
        need_create_loop_context = not isinstance(sync_ctx.context, LoopContext)
        ret = super().call(sync_ctx, env)
        if ret:
            if need_create_loop_context:
                ret.append(LoopContext(sync_ctx.context._runtime, self).sync_context)
            else:
                ret.append(sync_ctx)
        elif not need_create_loop_context:
            sync_ctx.destory()

        return ret

call(sync_ctx, env)

Execute one iteration and append a LoopContext to continue the loop.

Source code in shutils/dag/task.py

def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
    """Execute one iteration and append a LoopContext to continue the loop."""
    need_create_loop_context = not isinstance(sync_ctx.context, LoopContext)
    ret = super().call(sync_ctx, env)
    if ret:
        if need_create_loop_context:
            ret.append(LoopContext(sync_ctx.context._runtime, self).sync_context)
        else:
            ret.append(sync_ctx)
    elif not need_create_loop_context:
        sync_ctx.destory()

    return ret

SyncStreamTask

Bases: SyncTask, ShutdownTask

Sync task backed by a generator that streams multiple outputs.

Source code in shutils/dag/task.py

class SyncStreamTask(SyncTask, ShutdownTask):
    """Sync task backed by a generator that streams multiple outputs."""

    def __init__(
        self,
        func: Callable[[], Generator[SyncContext | list[SyncContext] | None, SyncContext]],
        config: TaskConfig | None = None,
        name: str = "",
    ):
        """Initialize a sync stream task.

        Args:
            func: A generator function that yields output contexts.
            config: Task execution configuration.
            name: Optional task name.

        Raises:
            ValueError: If func does not return a generator.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(func, config, name)
        self._generator = func()
        if not isgenerator(self._generator):
            raise ValueError("func must be a generator function")
        next(self._generator)

    def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
        """Send the context to the generator and return yielded outputs."""
        try:
            ret = self._generator.send(sync_ctx)
            if ret is None:
                return []
            if isinstance(ret, SyncContext):
                return [ret]
            return ret
        except StopIteration:
            return []

    def __repr__(self):
        return f"{self.__class__.__name__}(func={self._generator}, {TaskBase.__repr__(self)})"

    def shutdown(self):
        """Close the underlying generator."""
        with contextlib.suppress(GeneratorExit):
            self._generator.close()

__init__(func, config=None, name='')

Initialize a sync stream task.

Parameters:

Name	Type	Description	Default
func	Callable[[], Generator[SyncContext | list[SyncContext] | None, SyncContext]]	A generator function that yields output contexts.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Raises:

Type	Description
ValueError	If func does not return a generator.

Source code in shutils/dag/task.py

def __init__(
    self,
    func: Callable[[], Generator[SyncContext | list[SyncContext] | None, SyncContext]],
    config: TaskConfig | None = None,
    name: str = "",
):
    """Initialize a sync stream task.

    Args:
        func: A generator function that yields output contexts.
        config: Task execution configuration.
        name: Optional task name.

    Raises:
        ValueError: If func does not return a generator.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(func, config, name)
    self._generator = func()
    if not isgenerator(self._generator):
        raise ValueError("func must be a generator function")
    next(self._generator)

call(sync_ctx, env)

Send the context to the generator and return yielded outputs.

Source code in shutils/dag/task.py

def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
    """Send the context to the generator and return yielded outputs."""
    try:
        ret = self._generator.send(sync_ctx)
        if ret is None:
            return []
        if isinstance(ret, SyncContext):
            return [ret]
        return ret
    except StopIteration:
        return []

shutdown()

Close the underlying generator.

Source code in shutils/dag/task.py

def shutdown(self):
    """Close the underlying generator."""
    with contextlib.suppress(GeneratorExit):
        self._generator.close()

SyncTask

Bases: TaskBase

Base class for synchronous tasks.

Source code in shutils/dag/task.py

class SyncTask(TaskBase):
    """Base class for synchronous tasks."""

    @abstractmethod
    def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
        """Execute the task synchronously.

        Args:
            sync_ctx: The sync context for this execution.
            env: The runtime environment.

        Returns:
            List of output sync contexts.
        """

    def __call__(self, context: Context, env: Environment) -> list[Context]:
        ret = self.call_before(context)
        if ret:
            return [ret]
        sync_ret = self.call(context.sync_context, env)
        ret = [ctx.context for ctx in sync_ret]
        self.call_after(ret)
        return ret

call(sync_ctx, env) abstractmethod

Execute the task synchronously.

Parameters:

Name	Type	Description	Default
sync_ctx	SyncContext	The sync context for this execution.	required
env	Environment	The runtime environment.	required

Returns:

Type	Description
list[SyncContext]	List of output sync contexts.

Source code in shutils/dag/task.py

@abstractmethod
def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
    """Execute the task synchronously.

    Args:
        sync_ctx: The sync context for this execution.
        env: The runtime environment.

    Returns:
        List of output sync contexts.
    """

SyncThreadTask

Bases: SyncTask, LongRunningTask, ShutdownTask

Sync task that runs in a dedicated background thread.

Source code in shutils/dag/task.py

class SyncThreadTask(SyncTask, LongRunningTask, ShutdownTask):
    """Sync task that runs in a dedicated background thread."""

    def __init__(
        self,
        func: Callable[[queue.Queue[tuple[Context, queue.Queue[list[Context] | Context | None]]]], None],
        config: TaskConfig | None = None,
        name: str = "",
    ):
        """Initialize a sync thread task.

        Args:
            func: A function that consumes from a queue of (context, response_queue) pairs.
            config: Task execution configuration.
            name: Optional task name.
        """
        if config is None:
            config = TaskConfig()
        super().__init__(func, config, name)
        self.__input_queue = queue.Queue()
        self._func = func
        self.__thread = None

    def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
        """Send context to the background thread and wait for the result."""
        if self.__thread is None:
            self.__thread = threading.Thread(target=self._func, args=(self.__input_queue,))
            self.__thread.start()
        if not self.__thread.is_alive():
            logger.error("[SyncThreadTask]: thread exit unexpectedly")
            raise RuntimeError("thread is not alive")
        future = queue.Queue()
        self.__input_queue.put((sync_ctx, future))
        result = future.get()
        if isinstance(result, SyncContext):
            return [sync_ctx]
        elif isinstance(result, list):
            return result
        return []

    def __repr__(self):
        return f"{self.__class__.__name__}(func={self._func}, {TaskBase.__repr__(self)})"

    def shutdown(self):
        """Send a stop signal to the background thread and join it."""
        if self.__thread and self.__thread.is_alive():
            self.__input_queue.put((StopContext().sync_context, queue.Queue()))
            self.__thread.join()

__init__(func, config=None, name='')

Initialize a sync thread task.

Parameters:

Name	Type	Description	Default
func	Callable[[Queue[tuple[Context, Queue[list[Context] | Context | None]]]], None]	A function that consumes from a queue of (context, response_queue) pairs.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name.	''

Source code in shutils/dag/task.py

def __init__(
    self,
    func: Callable[[queue.Queue[tuple[Context, queue.Queue[list[Context] | Context | None]]]], None],
    config: TaskConfig | None = None,
    name: str = "",
):
    """Initialize a sync thread task.

    Args:
        func: A function that consumes from a queue of (context, response_queue) pairs.
        config: Task execution configuration.
        name: Optional task name.
    """
    if config is None:
        config = TaskConfig()
    super().__init__(func, config, name)
    self.__input_queue = queue.Queue()
    self._func = func
    self.__thread = None

call(sync_ctx, env)

Send context to the background thread and wait for the result.

Source code in shutils/dag/task.py

def call(self, sync_ctx: SyncContext, env: Environment) -> list[SyncContext]:
    """Send context to the background thread and wait for the result."""
    if self.__thread is None:
        self.__thread = threading.Thread(target=self._func, args=(self.__input_queue,))
        self.__thread.start()
    if not self.__thread.is_alive():
        logger.error("[SyncThreadTask]: thread exit unexpectedly")
        raise RuntimeError("thread is not alive")
    future = queue.Queue()
    self.__input_queue.put((sync_ctx, future))
    result = future.get()
    if isinstance(result, SyncContext):
        return [sync_ctx]
    elif isinstance(result, list):
        return result
    return []

shutdown()

Send a stop signal to the background thread and join it.

Source code in shutils/dag/task.py

def shutdown(self):
    """Send a stop signal to the background thread and join it."""
    if self.__thread and self.__thread.is_alive():
        self.__input_queue.put((StopContext().sync_context, queue.Queue()))
        self.__thread.join()

TaskBase

Bases: ABC

Abstract base class for all DAG tasks.

Source code in shutils/dag/task.py

class TaskBase(ABC):  # noqa: B024
    """Abstract base class for all DAG tasks."""

    def __init__(self, func: Callable | None, config: TaskConfig | None = None, name: str = ""):
        """Initialize a task with optional function, config, and name.

        Args:
            func: The callable this task wraps.
            config: Task execution configuration.
            name: Optional task name; defaults to a UUID.
        """
        if config is None:
            config = TaskConfig()
        self.id = name if name else str(uuid.uuid4())
        self.upstream_tasks: set[TaskBase] = set()
        self.downstream_tasks: set[TaskBase] = set()
        self.config = config
        self.running_task_num = 0
        self.rate_limiter = self.config.limiter

    def add_upstream(self, task: "TaskBase"):
        """Add an upstream dependency to this task.

        Args:
            task: The upstream task to depend on.
        """
        self.upstream_tasks.add(task)
        task.downstream_tasks.add(self)

    def call_before(self, context: Context) -> Context | None:
        """Pre-execution hook for parallel and rate-limit control.

        Returns a context if the task should be skipped, or None to proceed.
        """
        # parallel control
        if self.config.parallel_num > 0 and self.running_task_num > self.config.parallel_num:
            return context
        self.running_task_num += 1

        # qps control
        if self.rate_limiter is not None and not self.rate_limiter.try_acquire().success:
            logger.debug(f"[Task {self.id}]: rate limit exceeded, throttling...")
            self.running_task_num -= 1
            return RateLimitContext(context)

        return None

    def call_after(self, context_list: list[Context]) -> None:
        """Post-execution hook to release rate-limit tokens and update parallel count."""
        self.running_task_num -= 1
        if self.rate_limiter is not None:
            self.rate_limiter.release()

    def __hash__(self):
        return hash(self.id)

    def __eq__(self, other: object):
        if not isinstance(other, TaskBase):
            return False
        return self.id == other.id

    def __repr__(self):
        return f"id={self.id}, config={self.config}"

__init__(func, config=None, name='')

Initialize a task with optional function, config, and name.

Parameters:

Name	Type	Description	Default
func	Callable | None	The callable this task wraps.	required
config	TaskConfig | None	Task execution configuration.	None
name	str	Optional task name; defaults to a UUID.	''

Source code in shutils/dag/task.py

def __init__(self, func: Callable | None, config: TaskConfig | None = None, name: str = ""):
    """Initialize a task with optional function, config, and name.

    Args:
        func: The callable this task wraps.
        config: Task execution configuration.
        name: Optional task name; defaults to a UUID.
    """
    if config is None:
        config = TaskConfig()
    self.id = name if name else str(uuid.uuid4())
    self.upstream_tasks: set[TaskBase] = set()
    self.downstream_tasks: set[TaskBase] = set()
    self.config = config
    self.running_task_num = 0
    self.rate_limiter = self.config.limiter

add_upstream(task)

Add an upstream dependency to this task.

Parameters:

Name	Type	Description	Default
task	TaskBase	The upstream task to depend on.	required

Source code in shutils/dag/task.py

def add_upstream(self, task: "TaskBase"):
    """Add an upstream dependency to this task.

    Args:
        task: The upstream task to depend on.
    """
    self.upstream_tasks.add(task)
    task.downstream_tasks.add(self)

call_after(context_list)

Post-execution hook to release rate-limit tokens and update parallel count.

Source code in shutils/dag/task.py

def call_after(self, context_list: list[Context]) -> None:
    """Post-execution hook to release rate-limit tokens and update parallel count."""
    self.running_task_num -= 1
    if self.rate_limiter is not None:
        self.rate_limiter.release()

call_before(context)

Pre-execution hook for parallel and rate-limit control.

Returns a context if the task should be skipped, or None to proceed.

Source code in shutils/dag/task.py

def call_before(self, context: Context) -> Context | None:
    """Pre-execution hook for parallel and rate-limit control.

    Returns a context if the task should be skipped, or None to proceed.
    """
    # parallel control
    if self.config.parallel_num > 0 and self.running_task_num > self.config.parallel_num:
        return context
    self.running_task_num += 1

    # qps control
    if self.rate_limiter is not None and not self.rate_limiter.try_acquire().success:
        logger.debug(f"[Task {self.id}]: rate limit exceeded, throttling...")
        self.running_task_num -= 1
        return RateLimitContext(context)

    return None

TaskConfig dataclass

Configuration for task execution behavior.

Attributes:

Name	Type	Description
retry_times	int	Maximum number of retry attempts on failure.
retry_interval	int | float | Callable[[Context], float | int]	Seconds between retries, or a callable returning the interval.
parallel_num	int	Maximum concurrent executions (0 means unlimited).
limiter	Limiter | None	Optional rate limiter for QPS/concurrency control.

Source code in shutils/dag/task.py

@dataclass
class TaskConfig:
    """Configuration for task execution behavior.

    Attributes:
        retry_times: Maximum number of retry attempts on failure.
        retry_interval: Seconds between retries, or a callable returning the interval.
        parallel_num: Maximum concurrent executions (0 means unlimited).
        limiter: Optional rate limiter for QPS/concurrency control.
    """
    retry_times: int = 0
    retry_interval: int | float | Callable[[Context], float | int] = 0
    parallel_num: int = 0
    limiter: Limiter | None = None