dagster-async-executor 0.0.1__py3-none-any.whl

dagster_async_executor/__init__.py

@@ -0,0 +1,9 @@
+from dagster._core.libraries import DagsterLibraryRegistry
+
+from dagster_async_executor.executor_definition import async_executor as async_executor
+
+__version__ = "0.0.1"
+
+DagsterLibraryRegistry.register(
+    "dagster-async-executor", __version__, is_dagster_package=False
+)

dagster_async_executor/executor.py

@@ -0,0 +1,316 @@
+import os
+import queue
+import sys
+from collections.abc import AsyncIterator, Iterator
+from contextlib import ExitStack
+from dataclasses import dataclass
+from typing import Any, Optional, ParamSpec, TypeVar
+
+import anyio
+import anyio.abc
+from anyio.from_thread import start_blocking_portal
+from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
+from dagster_shared.utils.timing import format_duration
+
+import dagster._check as check
+from dagster._core.events import DagsterEvent, EngineEventData
+from dagster._core.execution.api import ExecuteRunWithPlanIterable
+from dagster._core.execution.compute_logs import create_compute_log_file_key
+from dagster._core.execution.context.system import (
+    PlanExecutionContext,
+    PlanOrchestrationContext,
+)
+from dagster._core.execution.context_creation_job import PlanExecutionContextManager
+from dagster._core.execution.plan.active import ActiveExecution
+from dagster_async_executor.execution.plan.execute_plan import (
+    _trigger_hook,
+    dagster_event_sequence_for_step,
+)
+from dagster_async_executor.execution.plan.execute_step import _verify_if_complete
+from dagster._core.execution.plan.execute_plan import _handle_compute_log_setup_error
+from dagster._core.execution.plan.instance_concurrency_context import (
+    InstanceConcurrencyContext,
+)
+from dagster._core.execution.plan.objects import step_failure_event_from_exc_info
+from dagster._core.execution.plan.plan import ExecutionPlan
+from dagster._core.execution.plan.state import KnownExecutionState
+from dagster._core.execution.plan.step import ExecutionStep
+from dagster._core.execution.retries import RetryMode
+from dagster._core.execution.step_dependency_config import StepDependencyConfig
+from dagster._core.executor.base import Executor
+from dagster._utils.timing import time_execution_scope
+
+T = TypeVar("T")
+P = ParamSpec("P")
+
+_SENTINEL = object()
+_SYNC_ASYNC_BRIDGE_QUEUE_MAXSIZE = 0
+
+
+@dataclass
+class _ErrorWrapper:
+    exc: BaseException
+
+
+class AsyncExecutor(Executor):
+    def __init__(
+        self,
+        retries: RetryMode,
+        max_concurrent: Optional[int] = None,
+        tag_concurrency_limits: Optional[list[dict[str, Any]]] = None,
+        step_dependency_config: StepDependencyConfig = StepDependencyConfig.default(),
+    ):
+        self._retries = check.inst_param(retries, "retries", RetryMode)
+        self._step_dependency_config = check.inst_param(
+            step_dependency_config, "step_dependency_config", StepDependencyConfig
+        )
+        self._max_concurrent = check.opt_int_param(max_concurrent, "max_concurrent")
+        self._tag_concurrency_limits = check.opt_list_param(
+            tag_concurrency_limits, "tag_concurrency_limits"
+        )
+
+    @property
+    def retries(self) -> RetryMode:
+        return self._retries
+
+    @property
+    def step_dependency_config(self) -> StepDependencyConfig:
+        return self._step_dependency_config
+
+    def execute(
+        self,
+        plan_context: PlanOrchestrationContext,
+        execution_plan: ExecutionPlan,
+    ) -> Iterator[DagsterEvent]:
+        """Synchronous entrypoint.
+
+        Uses ExecuteRunWithPlanIterable to get a PlanExecutionContext, but the
+        actual scheduling logic is defined by this executor via
+        _async_execution_iterator.
+        """
+        check.inst_param(plan_context, "plan_context", PlanOrchestrationContext)
+        check.inst_param(execution_plan, "execution_plan", ExecutionPlan)
+
+        step_keys_to_execute = execution_plan.step_keys_to_execute
+
+        yield DagsterEvent.engine_event(
+            plan_context,
+            "Executing steps with AsyncExecutor",
+            event_specific_data=EngineEventData.in_process(
+                os.getpid(), step_keys_to_execute
+            ),
+        )
+
+        with time_execution_scope() as timer_result:
+            yield from ExecuteRunWithPlanIterable(
+                execution_plan=execution_plan,
+                iterator=self._execution_iterator_wrapper,
+                execution_context_manager=PlanExecutionContextManager(
+                    job=plan_context.job,
+                    retry_mode=plan_context.retry_mode,
+                    execution_plan=execution_plan,
+                    run_config=plan_context.run_config,
+                    dagster_run=plan_context.dagster_run,
+                    instance=plan_context.instance,
+                    raise_on_error=plan_context.raise_on_error,
+                    output_capture=plan_context.output_capture,
+                    step_dependency_config=self.step_dependency_config,
+                ),
+            )
+
+        yield DagsterEvent.engine_event(
+            plan_context,
+            f"Finished AsyncExecutor in {format_duration(timer_result.millis)}",
+            event_specific_data=EngineEventData.in_process(
+                os.getpid(), step_keys_to_execute
+            ),
+        )
+
+    def _execution_iterator_wrapper(
+        self,
+        job_context: PlanExecutionContext,
+        execution_plan: ExecutionPlan,
+    ) -> Iterator[DagsterEvent]:
+        """Synchronous wrapper around the async execution iterator.
+
+        Uses an anyio BlockingPortal and a queue to bridge between the async and sync worlds.
+        This allows the executor to and present a synchronous iterator interface while leveraging async execution under
+        the hood, avoid any issues with things like compute log capture that are not inherently async-aware or thread-safe.
+        """
+        event_queue: queue.Queue[DagsterEvent | _ErrorWrapper | object] = queue.Queue(
+            maxsize=_SYNC_ASYNC_BRIDGE_QUEUE_MAXSIZE
+        )
+
+        compute_log_manager = job_context.instance.compute_log_manager
+        step_keys = [s.key for s in execution_plan.get_steps_to_execute_in_topo_order()]
+        file_key = create_compute_log_file_key()
+        log_key = compute_log_manager.build_log_key_for_run(
+            job_context.run_id, file_key
+        )
+
+        with (
+            InstanceConcurrencyContext(
+                job_context.instance, job_context.dagster_run
+            ) as instance_concurrency_context,
+            execution_plan.start(
+                retry_mode=self.retries,
+                max_concurrent=self._max_concurrent,
+                tag_concurrency_limits=self._tag_concurrency_limits,
+                instance_concurrency_context=instance_concurrency_context,
+                step_dependency_config=self.step_dependency_config,
+            ) as active,
+            ExitStack() as capture_stack,
+        ):
+            # 1) Compute logs (still sync)
+            try:
+                log_context = capture_stack.enter_context(
+                    compute_log_manager.capture_logs(log_key)
+                )
+                yield DagsterEvent.capture_logs(
+                    job_context, step_keys, log_key, log_context
+                )
+            except Exception:
+                yield from _handle_compute_log_setup_error(job_context, sys.exc_info())
+
+            # 2) Define async to sync bridge
+            async def _async_iterator_to_queue() -> None:
+                try:
+                    async for i in self._async_execution_iterator(job_context, active):
+                        event_queue.put(i)
+                except BaseException as e:
+                    event_queue.put(_ErrorWrapper(e))
+                finally:
+                    event_queue.put(_SENTINEL)
+
+            # 3) Start the blocking portal and consume from the queue
+            with start_blocking_portal() as portal:
+                task = portal.start_task_soon(_async_iterator_to_queue)
+                try:
+                    while True:
+                        item = event_queue.get()
+                        if item is _SENTINEL:
+                            break
+                        if isinstance(item, _ErrorWrapper):
+                            try:
+                                task.cancel()
+                            except BaseException:
+                                pass
+                            raise item.exc
+                        assert isinstance(
+                            item, DagsterEvent
+                        )  # after narrowing, item must be DagsterEvent
+                        yield item
+                except BaseException:
+                    try:
+                        task.cancel()
+                    except Exception:
+                        pass
+                    raise
+
+    async def _async_execution_iterator(
+        self,
+        job_context: PlanExecutionContext,
+        active: ActiveExecution,
+    ) -> AsyncIterator[DagsterEvent]:
+        """Async orchestrator that runs steps in a TaskGroup and yields events.
+
+        - Uses a single anyio memory object stream as the central event bus.
+        - Each step worker clones the send_stream and writes its events there.
+        - We update ActiveExecution based on events and schedule new work as usual.
+        """
+        send_stream, recv_stream = anyio.create_memory_object_stream[DagsterEvent]()
+
+        async with recv_stream, anyio.create_task_group() as tg:
+            try:
+                while not active.is_complete:
+                    steps_to_execute = active.get_steps_to_execute(limit=None)
+
+                    for event in active.concurrency_event_iterator(job_context):
+                        yield event
+
+                    known_state = active.get_known_state()
+                    for step in steps_to_execute:
+                        tg.start_soon(
+                            self._run_step_worker,
+                            job_context,
+                            known_state,
+                            step,
+                            send_stream.clone(),
+                        )
+
+                    try:
+                        event = await recv_stream.receive()
+                    except anyio.EndOfStream:
+                        raise
+
+                    yield event
+                    active.handle_event(event)
+                    await _verify_if_complete(job_context, active, event)
+
+                    for plan_event in active.plan_events_iterator(job_context):
+                        yield plan_event
+            finally:
+                await send_stream.aclose()
+
+            async for event in self._drain_events(active, job_context, recv_stream):
+                yield event
+
+    async def _run_step_worker(
+        self,
+        job_context: PlanExecutionContext,
+        known_state: KnownExecutionState,
+        step: ExecutionStep,
+        send_stream: MemoryObjectSendStream[DagsterEvent],
+    ) -> None:
+        """Run a single step's async compute, emitting DagsterEvents via send_stream."""
+        step_context = job_context.for_step(step, known_state)
+        missing_resources = [
+            resource_key
+            for resource_key in step_context.required_resource_keys
+            if not hasattr(step_context.resources, resource_key)
+        ]
+        check.invariant(
+            len(missing_resources) == 0,
+            (
+                f"Expected step context for op {step_context.op.name} to have all required "
+                f"resources, but missing {missing_resources}."
+            ),
+        )
+        async with send_stream:
+            try:
+                step_events: list[DagsterEvent] = []
+                async for event in dagster_event_sequence_for_step(step_context):
+                    await send_stream.send(event)
+                    step_events.append(event)
+
+                async for hook_event in _trigger_hook(step_context, step_events):
+                    await send_stream.send(hook_event)
+            except BaseException:
+                failure_event = step_failure_event_from_exc_info(
+                    step_context,
+                    sys.exc_info(),
+                )
+                try:
+                    await send_stream.send(failure_event)
+                finally:
+                    raise
+
+    async def _drain_events(
+        self,
+        active: ActiveExecution,
+        job_context: PlanExecutionContext,
+        recv_stream: MemoryObjectReceiveStream[DagsterEvent],
+    ) -> AsyncIterator[DagsterEvent]:
+        """Drain any remaining events from recv_stream after all steps complete."""
+        while True:
+            try:
+                event = await recv_stream.receive()
+            except anyio.EndOfStream:
+                break
+            yield event
+
+            active.handle_event(event)
+            await _verify_if_complete(job_context, active, event)
+
+            for plan_event in active.plan_events_iterator(job_context):
+                yield plan_event

dagster_async_executor/executor_definition.py

@@ -0,0 +1,53 @@
+from typing import TYPE_CHECKING
+
+from dagster._core.definitions.executor_definition import ExecutorConfig, executor
+from dagster_shared import check
+from dagster._builtins import Int
+from dagster._config import Field, Noneable
+from dagster._core.execution.retries import RetryMode, get_retries_config
+from dagster._core.execution.step_dependency_config import (
+    StepDependencyConfig,
+    get_step_dependency_config_field,
+)
+from dagster._core.execution.tags import get_tag_concurrency_limits_config
+
+if TYPE_CHECKING:
+    from dagster_async_executor.executor import AsyncExecutor
+
+ASYNC_CONFIG = Field(
+    {
+        "max_concurrent": Field(
+            Noneable(Int),
+            default_value=None,
+            description=(
+                "The number of asynchronous tasks that may run concurrently. "
+                "By default, this is set to None, which allows for an unlimited number of "
+                "concurrent tasks."
+            ),
+        ),
+        "tag_concurrency_limits": get_tag_concurrency_limits_config(),
+        "retries": get_retries_config(),
+        "step_dependency_config": get_step_dependency_config_field(),
+    },
+)
+
+
+@executor(
+    name="async_executor",
+    config_schema=ASYNC_CONFIG,
+)
+def async_executor(init_context):
+    return _core_async_executor_creation(init_context.executor_config)
+
+
+def _core_async_executor_creation(config: ExecutorConfig) -> "AsyncExecutor":
+    from dagster_async_executor.executor import AsyncExecutor
+
+    return AsyncExecutor(
+        retries=RetryMode.from_config(check.dict_elem(config, "retries")),  # type: ignore  # (possible none)
+        max_concurrent=check.opt_int_elem(config, "max_concurrent"),
+        tag_concurrency_limits=check.opt_list_elem(config, "tag_concurrency_limits"),
+        step_dependency_config=StepDependencyConfig.from_config(
+            check.opt_nullable_dict_elem(config, "step_dependency_config")
+        ),
+    )

dagster_async_executor-0.0.1.dist-info/METADATA

@@ -0,0 +1,180 @@
+Metadata-Version: 2.4
+Name: dagster-async-executor
+Version: 0.0.1
+Summary: Async Support for Dagster
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: anyio>=4.12.0
+Requires-Dist: dagster>=1.12.2
+
+# dagster-async-executor
+
+An executor for [Dagster](https://github.com/dagster-io/dagster) that adds **native `asyncio` support** for ops and assets.
+
+`dagster-async-executor` lets you:
+
+- Run `async def` ops without manually managing event loops or thread pools.
+- Mix sync and async ops in the same job.
+- Use dynamic / fan-out graphs with async upstream and downstream dependencies.
+- Keep the familiar Dagster executor interface, while enabling more scalable, concurrency‑friendly workloads (e.g. I/O‑heavy tasks, service calls, and streaming patterns).
+
+This integration is a community‑maintained port of the original core PR: [dagster-io/dagster#32833](https://github.com/dagster-io/dagster/pull/32833).
+
+---
+
+## Installation
+
+```bash
+pip install dagster-async-executor
+```
+
+---
+
+## Quickstart
+
+Use the `async_executor` when defining your job and write your ops as `async def`:
+
+```python
+import anyio
+import dagster as dg
+from dagster_async_executor import async_executor
+
+NUM_FANOUTS = 300
+SLEEP_SECONDS = 3
+
+
+@dg.op(out=dg.DynamicOut())
+async def create_dynamic_outputs():
+    """Creates a dynamic number of outputs."""
+    for i in range(NUM_FANOUTS):
+        yield dg.DynamicOutput(value=f"item_{i}", mapping_key=f"key_{i}")
+
+
+@dg.op
+async def process_item(context: dg.OpExecutionContext, item: str):
+    """Process each item from the fan-out."""
+    context.log.info(f"[{context.op_handle}] sleeping...")
+    await anyio.sleep(SLEEP_SECONDS)
+    context.log.info(f"[{context.op_handle}] completed")
+    return item
+
+
+@dg.op
+async def collect_results(context: dg.OpExecutionContext, results: list):
+    """Collect all results from the fan-out."""
+    context.log.info(f"[{context.op_handle}] collected {len(results)} results")
+    return results
+
+
+@dg.job(executor_def=async_executor)
+def simple_fanout_job():
+    # no need to use await
+    dynamic_items = create_dynamic_outputs()
+    processed = dynamic_items.map(process_item)
+    collect_results(processed.collect())
+```
+
+Run the job as usual (e.g. via Dagit, `dagster job execute`, or your orchestration environment). From the outside, this executor behaves like a standard Dagster executor – but internally it uses async orchestration.
+
+---
+
+## How it works
+
+At a high level, `dagster-async-executor` introduces an `AsyncExecutor` that:
+
+- Reuses Dagster’s **existing execution plan** machinery.
+- Runs steps in an **async orchestration loop** backed by an `anyio.TaskGroup`.
+- Bridges between **async step execution** and Dagster’s **synchronous executor interface** via a queue‑based sync–async bridge.
+
+### Execution model
+
+Conceptually, execution looks like this:
+
+1. **Plan creation (sync)**  
+   The run, plan, and context are created synchronously, just like with `in_process` and other executors.
+
+2. **Async orchestration loop**  
+   Once the plan is ready, an async orchestrator:
+   - Schedules each ready step as an async task.
+   - Uses `dagster_event_sequence_for_step` to obtain an **async sequence of `DagsterEvent`s** for each step.
+   - Sends those events through async streams.
+
+3. **Sync–async event bridge**  
+   A synchronous wrapper:
+   - Starts the async orchestrator inside an `anyio` `BlockingPortal`.
+   - Streams `DagsterEvent`s into a `queue.Queue`.
+   - Exposes a standard **`Iterator[DagsterEvent]`** to Dagster’s core execution machinery.
+
+From the rest of the system’s perspective, this executor still “looks like” a normal synchronous executor, which keeps:
+
+- Resource initialization behavior consistent.
+- Logging and event semantics unchanged.
+- Compatibility with existing Dagster entrypoints and tooling.
+
+### Sync + async ops
+
+Per-step behavior:
+
+- Async orchestration drives all steps.
+- Each step:
+  - Builds a `step_context`.
+  - Iterates over `dagster_event_sequence_for_step(step_context)` in an async `for` loop.
+  - Sends each `DagsterEvent` back through the async stream → queue → iterator bridge.
+
+Because the core execution semantics are reused, you can mix sync and async ops in the same graph:
+
+- Async upstream → sync downstream
+- Sync upstream → async downstream
+- Dynamic outputs and mapped steps that interleave async work
+
+---
+
+## Performance
+
+The executor is designed for I/O‑bound and highly concurrent workloads. An initial performance test (`test_async_executor_performance.py::test_async_performance_basic`) shows improved scaling with increased parallelism.
+
+Example results (fan‑out of async ops sleeping for 3 seconds):
+
+| Number of ops | Sleep (seconds) | Job duration (seconds) |
+|--------------:|----------------:|------------------------:|
+|             1 |               3 | 3.25                    |
+|             5 |               3 | 3.51                    |
+|            20 |               3 | 4.23                    |
+|           100 |               3 | 6.86                    |
+|           300 |               3 | 14.17                   |
+
+These numbers are illustrative; real‑world performance depends on your environment, I/O characteristics, and concurrency limits.
+
+---
+
+## Testing & behavior guarantees
+
+The test suite focuses on validating behavior across a representative set of job shapes:
+
+- **Basic async jobs**
+  - Single async op producing a simple output.
+  - Multiple async ops with dependencies and parallelism where possible.
+
+- **Mixed sync/async graphs**
+  - Async upstream feeding into sync downstream.
+  - Sync upstream feeding into async downstream.
+  - Ensuring consistent materializations, events, and success/failure semantics across both kinds of ops.
+
+- **Dynamic / fan‑out graphs**
+  - Async producers yielding dynamic outputs.
+  - Downstream mapping over dynamic keys.
+  - Interleaving async mapped steps and verifying all mapped outputs are awaited and collected correctly.
+
+- **Error handling**
+  - Exceptions raised from async ops (including inside dynamic maps).
+  - Failures reported on the correct steps.
+  - Downstream steps cancelled or skipped according to normal Dagster rules.
+
+---
+
+## Limitations & notes
+
+- This executor has not been tested on Python 3.14 free-threaded mode.
+- The executor targets **I/O‑bound** concurrency; CPU‑bound workloads should still be offloaded to processes or threads.
+- Cancellation, backpressure, and resource lifetime semantics follow Dagster’s existing execution model, but async nuances may still evolve.
+- This is a **community‑maintained** integration; behavior may change more rapidly than core Dagster executors as we iterate.

dagster_async_executor-0.0.1.dist-info/RECORD

@@ -0,0 +1,7 @@
+dagster_async_executor/__init__.py,sha256=vC3BtUjz0Mu8jQNsr7O6bulzLAUA7T2x6I14QoI1ebA,275
+dagster_async_executor/executor.py,sha256=f898L65a65hgyMzpEQo0joaEc3jCvPE6a9PgX0DU0L8,12641
+dagster_async_executor/executor_definition.py,sha256=AEElqwpAlyDBa2GYOG28Vrwbc-v9-CJl4okHRZVE07w,1926
+dagster_async_executor-0.0.1.dist-info/METADATA,sha256=UnfNGRE8L_IzVr40CqmCZHmSAyFT843TtO_cB4j_5x8,6727
+dagster_async_executor-0.0.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+dagster_async_executor-0.0.1.dist-info/top_level.txt,sha256=zc6K7GEfOGGmdYjRlAfpGPjFpVSmUV7_IIguUpslAbc,23
+dagster_async_executor-0.0.1.dist-info/RECORD,,

dagster_async_executor-0.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

dagster_async_executor-0.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+dagster_async_executor