diff --git a/src/orcapod/core/function_pod.py b/src/orcapod/core/function_pod.py index b835bb18b..2be8c7bb5 100644 --- a/src/orcapod/core/function_pod.py +++ b/src/orcapod/core/function_pod.py @@ -43,6 +43,7 @@ if TYPE_CHECKING: import polars as pl import pyarrow as pa + from orcapod.protocols.observability_protocols import ExecutionObserverProtocol else: pa = LazyModule("pyarrow") pl = LazyModule("polars") @@ -102,6 +103,11 @@ def executor(self, executor: DataFunctionExecutorProtocol | None) -> None: """Set or clear the executor on the underlying data function.""" self._data_function.executor = executor + @property + def pod_config(self) -> PodConfig: + """Per-pod executor configuration. Defaults to no concurrency limits.""" + return PodConfig() + def identity_structure(self) -> Any: return self.data_function.identity_structure() @@ -193,6 +199,69 @@ def handle_input_streams(self, *streams: StreamProtocol) -> StreamProtocol: return joined_stream return streams[0] + async def async_execute( + self, + inputs: Sequence[ReadableChannel[tuple[TagProtocol, DataProtocol]]], + output: WritableChannel[tuple[TagProtocol, DataProtocol]], + pipeline_config: PipelineConfig | None = None, + *, + observer: ExecutionObserverProtocol | None = None, + ) -> None: + """Streaming async execution with per-data concurrency control. + + Each input (tag, data) is dispatched as an independent async task. + A semaphore limits how many tasks are in-flight concurrently. + Observer hooks fire per item: ``on_data_start`` before processing, + ``on_data_end(cached=False)`` on success, ``on_data_crash`` on error. + + Args: + inputs: Single-element sequence containing the input channel. + output: Writable channel for output (tag, data) pairs. + pipeline_config: Optional pipeline-level concurrency config. + observer: Optional observer for per-item lifecycle hooks. + """ + from orcapod.pipeline.observer import NoOpObserver + + try: + pipeline_config = pipeline_config or PipelineConfig() + max_concurrency = resolve_concurrency(self.pod_config, pipeline_config) + obs = observer if observer is not None else NoOpObserver() + pod_label = self.label + + sem = ( + asyncio.Semaphore(max_concurrency) + if max_concurrency is not None + else None + ) + + async def process_one(tag: TagProtocol, data: DataProtocol) -> None: + obs.on_data_start(pod_label, tag, data) + pkt_logger = obs.create_data_logger(tag, data) + try: + out_tag, result_data = await self.async_process_data( + tag, data, logger=pkt_logger + ) + except Exception as exc: + logger.debug( + "Data processing failed, skipping: %s", exc, exc_info=True + ) + obs.on_data_crash(pod_label, tag, data, exc) + else: + obs.on_data_end(pod_label, tag, data, result_data, cached=False) + if result_data is not None: + await output.send((out_tag, result_data)) + finally: + if sem is not None: + sem.release() + + async with asyncio.TaskGroup() as tg: + async for tag, data in inputs[0]: + if sem is not None: + await sem.acquire() + tg.create_task(process_one(tag, data)) + finally: + await output.close() + @abstractmethod def process( self, *streams: StreamProtocol, label: str | None = None @@ -337,51 +406,6 @@ def from_config( return cls(data_function=data_function, pod_config=pod_config) - # ------------------------------------------------------------------ - # Async channel execution (streaming mode) - # ------------------------------------------------------------------ - - async def async_execute( - self, - inputs: Sequence[ReadableChannel[tuple[TagProtocol, DataProtocol]]], - output: WritableChannel[tuple[TagProtocol, DataProtocol]], - pipeline_config: PipelineConfig | None = None, - ) -> None: - """Streaming async execution with per-data concurrency control. - - Each input (tag, data) is processed independently. A semaphore - controls how many data are in-flight concurrently. - """ - try: - pipeline_config = pipeline_config or PipelineConfig() - max_concurrency = resolve_concurrency(self._pod_config, pipeline_config) - - sem = ( - asyncio.Semaphore(max_concurrency) - if max_concurrency is not None - else None - ) - - async def process_one(tag: TagProtocol, data: DataProtocol) -> None: - try: - tag, result_data = await self.async_process_data(tag, data) - if result_data is not None: - await output.send((tag, result_data)) - except Exception as e: - # Swallow data-level errors so remaining data continue. - logger.debug("Data processing failed, skipping: %s", e, exc_info=True) - finally: - if sem is not None: - sem.release() - - async with asyncio.TaskGroup() as tg: - async for tag, data in inputs[0]: - if sem is not None: - await sem.acquire() - tg.create_task(process_one(tag, data)) - finally: - await output.close() - class FunctionPodStream(StreamBase): """Recomputable stream wrapping a data function.""" @@ -769,6 +793,11 @@ def __init__( def computed_label(self) -> str | None: return self._function_pod.label + @property + def pod_config(self) -> PodConfig: + """Delegate to the inner pod's config so CachedFunctionPod respects limits.""" + return self._function_pod.pod_config + @property def uri(self) -> tuple[str, ...]: return self._function_pod.uri diff --git a/src/orcapod/core/nodes/function_node.py b/src/orcapod/core/nodes/function_node.py index 6acdaee80..56b2307d0 100644 --- a/src/orcapod/core/nodes/function_node.py +++ b/src/orcapod/core/nodes/function_node.py @@ -24,7 +24,7 @@ from typing import TYPE_CHECKING, Any, Literal, NamedTuple, cast from orcapod import contexts -from orcapod.channels import ReadableChannel, WritableChannel +from orcapod.channels import Channel, ReadableChannel, WritableChannel from orcapod.config import OrcapodConfig from orcapod.core.cached_function_pod import CachedFunctionPod from orcapod.core.streams.arrow_table_stream import ArrowTableStream @@ -50,10 +50,7 @@ ColumnConfig, ContentHash, NodeConfig, - PipelineConfig, - PodConfig, Schema, - resolve_concurrency, ) from orcapod.utils import arrow_utils, schema_utils from orcapod.utils.lazy_module import LazyModule @@ -82,6 +79,12 @@ # 0 for the first computation, N+1 for each miss-triggered recompute. _PIPELINE_RECOMPUTATION_INDEX_COL = "__pipeline_recomputation_index" +# Private meta-column name stamped into tags routed to the compute channel. +# Carries the correlation key (bytes) that links an execution result back to +# its original (tag, input_data) pair. Stripped from output tags in +# record_and_forward() before downstream emission. +_TAG_NODE_INPUT_REF = "_tag_node_input_ref" + def _executor_supports_concurrent( data_function: DataFunctionProtocol, @@ -2086,22 +2089,27 @@ async def async_execute( ctx_obs = obs.contextualize(*self.node_identity_path) try: - # Resolve concurrency limit from pod config (pipeline config is not - # threaded through the orchestrator, so we fall back to defaults). - pod_config = getattr(self._function_pod, "pod_config", PodConfig()) - max_concurrency = resolve_concurrency(pod_config, PipelineConfig()) - - sem = ( - asyncio.Semaphore(max_concurrency) - if max_concurrency is not None - else None - ) - tag_schema = self._input_stream.output_schema(columns={"system_tags": True})[0] ctx_obs.on_node_start(node_label, node_hash, tag_schema=tag_schema) if self._cached_function_pod is not None: - # Phase 1: build cache lookup from pipeline DB + # ---------------------------------------------------------- + # DB path — 3-stage concurrent pipeline + # ---------------------------------------------------------- + is_ephemeral = bool(self._node_config.is_result_ephemeral) + if is_ephemeral: + if self._ephemeral_cached_pod is None: + raise RuntimeError( + f"FunctionJobNode '{self.label}' has is_result_ephemeral=True " + "but no ephemeral store has been assigned. Call " + "set_ephemeral_store() with an ArrowDatabaseProtocol before " + "executing this node." + ) + execution_pod = self._ephemeral_cached_pod + else: + execution_pod = self._cached_function_pod + + # Load pipeline DB → in-memory cache keyed by base_entry_id. loaded = self._load_cached_entries() self._cached_output_datas.update(loaded) if loaded: @@ -2109,93 +2117,169 @@ async def async_execute( self._cached_content_hash_column = None cached_by_base_entry_id: dict[bytes, tuple[TagProtocol, DataProtocol]] = dict(loaded) - # Phase 2: drive output from input channel — cached or compute - async def _process_one_db( - tag: TagProtocol, data: DataProtocol - ) -> None: - base_entry_id = self.compute_base_entry_id(tag, data) - if base_entry_id in cached_by_base_entry_id: - tag_out, result_data = cached_by_base_entry_id[base_entry_id] - ctx_obs.on_data_start(node_label, tag, data) - ctx_obs.on_data_end( - node_label, tag, data, result_data, cached=True + # Intermediate channels (bounded for backpressure). + compute_channel: Channel[tuple[TagProtocol, DataProtocol]] = Channel(buffer_size=16) + result_channel: Channel[tuple[TagProtocol, DataProtocol]] = Channel(buffer_size=16) + + # Local dict: correlation_key → (original_tag, original_input_data) + input_store: dict[bytes, tuple[TagProtocol, DataProtocol]] = {} + + async def route_inputs() -> None: + """Stage 1: send cache hits to output; stamp misses for computation.""" + try: + async for tag, data in input_channel: + base_entry_id = self.compute_base_entry_id(tag, data) + if base_entry_id in cached_by_base_entry_id: + cached_tag, cached_data = cached_by_base_entry_id[base_entry_id] + ctx_obs.on_data_start(node_label, tag, data) + ctx_obs.on_data_end( + node_label, tag, data, cached_data, cached=True + ) + await output.send((cached_tag, cached_data)) + else: + correlation_key = uuid.uuid4().bytes + input_store[correlation_key] = (tag, data) + stamped_tag = tag.with_meta_columns( + **{_TAG_NODE_INPUT_REF: correlation_key} + ) + await compute_channel.writer.send((stamped_tag, data)) + finally: + await compute_channel.writer.close() + + async def record_and_forward() -> None: + """Stage 3: record pipeline entry, strip key, emit to output.""" + async for output_tag, output_data in result_channel.reader: + correlation_key = output_tag.get_meta_value(_TAG_NODE_INPUT_REF) + original_tag, input_data = input_store.pop(correlation_key) + result_computed = bool( + output_data.get_meta_value( + execution_pod.RESULT_COMPUTED_FLAG, False + ) + ) + self.add_pipeline_record( + original_tag, + input_data, + data_record_id=output_data.datagram_uuid, + computed=result_computed, + is_ephemeral=is_ephemeral, + ) + # Update in-memory cache so iter_data() sees the result. + base_entry_id = self.compute_base_entry_id(original_tag, input_data) + clean_tag = output_tag.drop_meta_columns( + _TAG_NODE_INPUT_REF, ignore_missing=True + ) + self._cached_output_datas[base_entry_id] = (clean_tag, output_data) + self._cached_output_table = None + self._cached_content_hash_column = None + await output.send((clean_tag, output_data)) + + # Capture outer self so _NodeLabelObserver methods can + # update FunctionJobNode state without shadowing with 'self'. + fn_node = self + + # Wrap ctx_obs so that data-level events emitted by the pod + # carry the node label, not the pod's own label — preserving + # the observable contract of the old DB path. + class _NodeLabelObserver: + """Relay all observer calls, replacing any label with node_label.""" + + def contextualize(self, *path: str) -> "_NodeLabelObserver": + return self + + def on_run_start(self, run_id: str, pipeline_uri: str = "") -> None: + ctx_obs.on_run_start(run_id, pipeline_uri) + + def on_run_end(self, run_id: str) -> None: + ctx_obs.on_run_end(run_id) + + def on_node_start(self, lbl: str, h: str, **kw: Any) -> None: + ctx_obs.on_node_start(node_label, h, **kw) + + def on_node_end(self, lbl: str, h: str) -> None: + ctx_obs.on_node_end(node_label, h) + + def on_data_start(self, lbl: str, tag: Any, data: Any) -> None: + clean = tag.drop_meta_columns( + _TAG_NODE_INPUT_REF, ignore_missing=True + ) + ctx_obs.on_data_start(node_label, clean, data) + + def on_data_end( + self, lbl: str, tag: Any, inp: Any, out: Any, *, cached: bool + ) -> None: + clean = tag.drop_meta_columns( + _TAG_NODE_INPUT_REF, ignore_missing=True + ) + ctx_obs.on_data_end(node_label, clean, inp, out, cached=cached) + if out is None: + # The function filtered this item (returned None). + # The pod does not emit to result_channel, so + # record_and_forward() never sees it. Pop the + # input_store entry and update the in-memory cache + # so iter_data() (the sync path) sees a cached None + # result — restoring the old _async_process_data_internal + # behaviour of always writing (tag, None) to + # _cached_output_datas. + # Note: route_inputs() checks cached_by_base_entry_id + # (pipeline DB) and filtered items are never written + # to the pipeline DB, so subsequent async_execute() + # calls still re-send filtered inputs to the pod. + correlation_key = tag.get_meta_value( + _TAG_NODE_INPUT_REF, None + ) + if correlation_key is not None: + original = input_store.pop(correlation_key, None) + if original is not None: + orig_tag, orig_data = original + base_entry_id = fn_node.compute_base_entry_id( + orig_tag, orig_data + ) + fn_node._cached_output_datas[base_entry_id] = ( + clean, + None, + ) + fn_node._cached_output_table = None + fn_node._cached_content_hash_column = None + + def on_data_crash( + self, lbl: str, tag: Any, data: Any, error: Any + ) -> None: + clean = tag.drop_meta_columns( + _TAG_NODE_INPUT_REF, ignore_missing=True ) - await output.send((tag_out, result_data)) - else: - await self._async_execute_one_data( - tag, data, output, - observer=ctx_obs, - node_label=node_label, - node_hash=node_hash, + ctx_obs.on_data_crash(node_label, clean, data, error) + # Clean up the dangling input_store entry so it doesn't + # accumulate memory during long-running calls with errors. + correlation_key = tag.get_meta_value(_TAG_NODE_INPUT_REF, None) + if correlation_key is not None: + input_store.pop(correlation_key, None) + + def create_data_logger(self, tag: Any, data: Any) -> Any: + # Strip the correlation key before creating the logger + # so that internal columns never leak into log records. + clean = tag.drop_meta_columns( + _TAG_NODE_INPUT_REF, ignore_missing=True ) + return ctx_obs.create_data_logger(clean, data) async with asyncio.TaskGroup() as tg: - async for tag, data in input_channel: - async def _guarded_db( - t: TagProtocol = tag, p: DataProtocol = data - ) -> None: - try: - await _process_one_db(t, p) - finally: - if sem is not None: - sem.release() - - if sem is not None: - await sem.acquire() - tg.create_task(_guarded_db()) + tg.create_task(route_inputs()) + tg.create_task( + execution_pod.async_execute( + [compute_channel.reader], + result_channel.writer, + observer=_NodeLabelObserver(), + ) + ) + tg.create_task(record_and_forward()) else: - # Simple async execution without DB - async with asyncio.TaskGroup() as tg: - async for tag, data in input_channel: - async def _guarded_simple( - t: TagProtocol = tag, p: DataProtocol = data - ) -> None: - try: - await self._async_execute_one_data( - t, p, output, - observer=ctx_obs, - node_label=node_label, - node_hash=node_hash, - ) - finally: - if sem is not None: - sem.release() - - if sem is not None: - await sem.acquire() - tg.create_task(_guarded_simple()) + # No-DB path: delegate dispatch entirely to the pod. + # The pod's async_execute() closes `output` in its own finally. + await self._function_pod.async_execute( + [input_channel], output, observer=ctx_obs + ) ctx_obs.on_node_end(node_label, node_hash) finally: await output.close() - async def _async_execute_one_data( - self, - tag: TagProtocol, - data: DataProtocol, - output: "WritableChannel[tuple[TagProtocol, DataProtocol]]", - *, - observer: ExecutionObserverProtocol, - node_label: str, - node_hash: str, - ) -> None: - """Process one non-cached data in the async execute path.""" - observer.on_data_start(node_label, tag, data) - pkt_logger = observer.create_data_logger(tag, data) - - try: - tag_out, result_data = await self._async_process_data_internal( - tag, data, logger=pkt_logger - ) - except Exception as exc: - logger.warning( - "Data execution failed in %s: %s", node_label, exc, - exc_info=True, - ) - observer.on_data_crash(node_label, tag, data, exc) - else: - observer.on_data_end( - node_label, tag, data, result_data, cached=False - ) - if result_data is not None: - await output.send((tag_out, result_data)) diff --git a/src/orcapod/protocols/core_protocols/function_pod.py b/src/orcapod/protocols/core_protocols/function_pod.py index 9c8bab91f..3f0de3ff6 100644 --- a/src/orcapod/protocols/core_protocols/function_pod.py +++ b/src/orcapod/protocols/core_protocols/function_pod.py @@ -4,6 +4,7 @@ from orcapod.protocols.core_protocols.data_function import DataFunctionProtocol from orcapod.protocols.core_protocols.pod import PodProtocol from orcapod.protocols.hashing_protocols import PipelineElementProtocol +from orcapod.types import PodConfig @runtime_checkable @@ -19,6 +20,11 @@ def data_function(self) -> DataFunctionProtocol: """ ... + @property + def pod_config(self) -> PodConfig: + """Per-pod executor configuration.""" + ... + def process_data( self, tag: TagProtocol, data: DataProtocol ) -> tuple[TagProtocol, DataProtocol | None]: ... diff --git a/superpowers/plans/2026-07-10-itl-516-unify-async-execute.md b/superpowers/plans/2026-07-10-itl-516-unify-async-execute.md new file mode 100644 index 000000000..f1e3cc046 --- /dev/null +++ b/superpowers/plans/2026-07-10-itl-516-unify-async-execute.md @@ -0,0 +1,1357 @@ +# ITL-516: Unify FunctionPod/FunctionJobNode async_execute() Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use sensei:subagent-driven-development (recommended) or sensei:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Move `async_execute()` to `_FunctionPodBase`, redesign `FunctionJobNode.async_execute()` as a 3-stage concurrent pipeline (route → compute → record), eliminating duplicate dispatch loops. + +**Architecture:** `_FunctionPodBase.async_execute()` becomes the single authoritative async dispatch loop (with observer support), used by all pod types via polymorphism. `FunctionJobNode.async_execute()` delegates computation to `execution_pod.async_execute()` via bounded `compute_channel`/`result_channel`, retaining only routing (cache hits vs misses), correlation-key stamping, `add_pipeline_record` calls, and key stripping. Three coroutines run concurrently in a `TaskGroup`. + +**Tech Stack:** Python asyncio (`asyncio.TaskGroup`, `asyncio.Semaphore`), `orcapod.channels.Channel`, `pytest-asyncio`. + +--- + +## File Structure + +| File | Action | Responsibility | +|------|--------|----------------| +| `src/orcapod/core/function_pod.py` | Modify | Add `pod_config` to `_FunctionPodBase` and `WrappedFunctionPod`; move `async_execute()` from `FunctionPod` to `_FunctionPodBase` with observer hooks | +| `src/orcapod/core/nodes/function_node.py` | Modify | Add `_TAG_NODE_INPUT_REF` constant; replace `FunctionJobNode.async_execute()` with 3-stage pipeline; delete `_async_execute_one_data`; remove dead imports | +| `tests/test_core/test_regression_fixes.py` | Modify | Remove `TestAsyncExecuteBackpressure` (orphan test for orphan method) | +| `tests/test_core/nodes/test_function_job_node_async_execute.py` | Create | All async_execute coverage: no-DB, DB (persistent + ephemeral), cache hits/misses, correlation key absence, concurrency, observer hooks | + +--- + +## Task 1: Move `async_execute()` to `_FunctionPodBase` with observer support + +`FunctionPod.async_execute()` currently lives on `FunctionPod` (lines 344–383 of `function_pod.py`). +`CachedFunctionPod` inherits from `WrappedFunctionPod → _FunctionPodBase`, not from `FunctionPod`, so +it currently has no `async_execute()`. Moving the method to `_FunctionPodBase` gives all pod types +the dispatch loop via polymorphism. + +Two supporting changes are needed first: a `pod_config` property on `_FunctionPodBase` (default +`PodConfig()`) and a delegating override on `WrappedFunctionPod`. Without these, `CachedFunctionPod` +would silently ignore concurrency limits set on the inner `FunctionPod`. + +**Files:** +- Modify: `src/orcapod/core/function_pod.py` + +- [ ] **Step 1: Write a failing test for observer hooks on `_FunctionPodBase.async_execute()`** + +Create a temporary test in an existing test file or a scratch location. We'll move this to the +dedicated file in Task 4. For now, paste it into the bottom of +`tests/test_core/test_regression_fixes.py` (we'll relocate in Task 4). + +```python +# Add these imports at the top of test_regression_fixes.py if not present: +# from orcapod.types import PodConfig + +class TestFunctionPodBaseAsyncExecuteObserver: + """Observer hooks fire correctly in _FunctionPodBase.async_execute().""" + + @pytest.mark.asyncio + async def test_observer_on_data_start_and_end_fire_per_item(self): + """on_data_start fires before and on_data_end fires after each item.""" + starts = [] + ends = [] + + class _Spy: + def contextualize(self, *path): + return self + def on_node_start(self, *a, **kw): pass + def on_node_end(self, *a, **kw): pass + def on_data_start(self, label, tag, data): + starts.append((label, tag, data)) + def on_data_end(self, label, tag, inp, out, *, cached): + ends.append((label, tag, inp, out, cached)) + def on_data_crash(self, label, tag, data, exc): + pass + def create_data_logger(self, tag, data): + return None + + def double(x: int) -> int: + return x * 2 + + pf = PythonDataFunction(double, output_keys="result") + pod = FunctionPod(pf) + + stream = make_stream(3) + input_ch = Channel(buffer_size=16) + output_ch = Channel(buffer_size=16) + await feed_stream_to_channel(stream, input_ch) + + spy = _Spy() + await pod.async_execute([input_ch.reader], output_ch.writer, observer=spy) + + results = await output_ch.reader.collect() + assert len(results) == 3 + assert len(starts) == 3 + assert len(ends) == 3 + # All items were computed, so cached=False + assert all(cached is False for _, _, _, _, cached in ends) + + @pytest.mark.asyncio + async def test_observer_on_data_crash_fires_on_failure(self): + """on_data_crash fires (not on_data_end) when async_process_data raises.""" + crashes = [] + ends = [] + + class _Spy: + def contextualize(self, *path): return self + def on_node_start(self, *a, **kw): pass + def on_node_end(self, *a, **kw): pass + def on_data_start(self, *a): pass + def on_data_end(self, label, tag, inp, out, *, cached): + ends.append((label, tag, inp, out, cached)) + def on_data_crash(self, label, tag, data, exc): + crashes.append(exc) + def create_data_logger(self, tag, data): + return None + + def boom(x: int) -> int: + raise ValueError("explode") + + pf = PythonDataFunction(boom, output_keys="result") + pod = FunctionPod(pf) + + stream = make_stream(2) + input_ch = Channel(buffer_size=16) + output_ch = Channel(buffer_size=16) + await feed_stream_to_channel(stream, input_ch) + + spy = _Spy() + await pod.async_execute([input_ch.reader], output_ch.writer, observer=spy) + + results = await output_ch.reader.collect() + assert len(results) == 0 # all failed → nothing forwarded + assert len(crashes) == 2 # crash fires per failed item + assert len(ends) == 0 # on_data_end never fires on crash +``` + +- [ ] **Step 2: Run the tests to confirm they fail** + +```bash +uv run pytest tests/test_core/test_regression_fixes.py::TestFunctionPodBaseAsyncExecuteObserver -v +``` + +Expected: `FAILED` — `pod.async_execute()` currently accepts no `observer` parameter. + +- [ ] **Step 3: Add `pod_config` default to `_FunctionPodBase` and delegating override to `WrappedFunctionPod`** + +In `src/orcapod/core/function_pod.py`, locate `_FunctionPodBase` (line 60) and add the property +**after** the `executor` setter (around line 103): + +```python +# --- add to _FunctionPodBase, after the executor setter --- +@property +def pod_config(self) -> PodConfig: + """Per-pod executor configuration. Defaults to no concurrency limits.""" + return PodConfig() +``` + +Then locate `WrappedFunctionPod` (line 750) and add a delegating override **after** `computed_label`: + +```python +# --- add to WrappedFunctionPod, after computed_label --- +@property +def pod_config(self) -> PodConfig: + """Delegate to the inner pod's config so CachedFunctionPod respects limits.""" + return getattr(self._function_pod, "pod_config", PodConfig()) +``` + +- [ ] **Step 4: Add `ExecutionObserverProtocol` import to `function_pod.py`** + +`function_pod.py` currently doesn't import `ExecutionObserverProtocol`. Add it to the `TYPE_CHECKING` +block (it's only used in a type hint): + +```python +# In the TYPE_CHECKING block at the top of function_pod.py, add: +if TYPE_CHECKING: + ... + from orcapod.protocols.observability_protocols import ExecutionObserverProtocol +``` + +- [ ] **Step 5: Move `async_execute()` from `FunctionPod` to `_FunctionPodBase` with observer support** + +Remove the existing `async_execute()` method from `FunctionPod` (lines 340–383, starting from the +`# Async channel execution (streaming mode)` section comment). + +Then add the following method to `_FunctionPodBase` **just before the abstract `process()` method** +(around line 196 in the original file — adjust for line shifts from Step 3): + +```python +async def async_execute( + self, + inputs: "Sequence[ReadableChannel[tuple[TagProtocol, DataProtocol]]]", + output: "WritableChannel[tuple[TagProtocol, DataProtocol]]", + pipeline_config: PipelineConfig | None = None, + *, + observer: "ExecutionObserverProtocol | None" = None, +) -> None: + """Streaming async execution with per-data concurrency control. + + Each input (tag, data) is dispatched as an independent async task. + A semaphore limits how many tasks are in-flight concurrently. + Observer hooks fire per item: ``on_data_start`` before processing, + ``on_data_end(cached=False)`` on success, ``on_data_crash`` on error. + + Args: + inputs: Single-element sequence containing the input channel. + output: Writable channel for output (tag, data) pairs. + pipeline_config: Optional pipeline-level concurrency config. + observer: Optional observer for per-item lifecycle hooks. + """ + from orcapod.pipeline.observer import NoOpObserver + + try: + pipeline_config = pipeline_config or PipelineConfig() + max_concurrency = resolve_concurrency(self.pod_config, pipeline_config) + obs = observer if observer is not None else NoOpObserver() + pod_label = self.label + + sem = ( + asyncio.Semaphore(max_concurrency) + if max_concurrency is not None + else None + ) + + async def process_one(tag: TagProtocol, data: DataProtocol) -> None: + obs.on_data_start(pod_label, tag, data) + try: + out_tag, result_data = await self.async_process_data(tag, data) + except Exception as exc: + logger.debug( + "Data processing failed, skipping: %s", exc, exc_info=True + ) + obs.on_data_crash(pod_label, tag, data, exc) + else: + obs.on_data_end(pod_label, tag, data, result_data, cached=False) + if result_data is not None: + await output.send((out_tag, result_data)) + finally: + if sem is not None: + sem.release() + + async with asyncio.TaskGroup() as tg: + async for tag, data in inputs[0]: + if sem is not None: + await sem.acquire() + tg.create_task(process_one(tag, data)) + finally: + await output.close() +``` + +- [ ] **Step 6: Run ALL tests to verify nothing broke** + +```bash +uv run pytest tests/ -v --timeout=30 +``` + +Expected: all previously-passing tests still pass, the new observer tests pass. + +Key tests to verify: +- `tests/test_core/test_regression_fixes.py::TestAsyncExecuteChannelCloseOnError` — still passes (FunctionPod inherits) +- `tests/test_core/test_regression_fixes.py::TestAsyncExecuteBackpressure` — still passes (FunctionPod still works) +- `tests/test_core/test_regression_fixes.py::TestFunctionPodBaseAsyncExecuteObserver` — now passes + +- [ ] **Step 7: Commit** + +```bash +git add src/orcapod/core/function_pod.py tests/test_core/test_regression_fixes.py +git commit -m "feat(function_pod): move async_execute to _FunctionPodBase with observer support + +Promotes async_execute() from FunctionPod to _FunctionPodBase so all pod +types (FunctionPod, CachedFunctionPod) inherit the dispatch loop. Adds +observer hooks (on_data_start, on_data_end, on_data_crash) per item. +Adds pod_config property to _FunctionPodBase (default PodConfig()) and +WrappedFunctionPod (delegates to inner pod) so concurrency limits are +honoured through CachedFunctionPod. + +Closes ITL-516 partial." +``` + +--- + +## Task 2: `FunctionJobNode.async_execute()` — no-DB delegation path + +Currently the no-DB path in `FunctionJobNode.async_execute()` (lines 2147–2166) reimplements +the same semaphore-guarded TaskGroup loop that's already in `_FunctionPodBase.async_execute()`. +Replace it with a single delegation call. This mirrors the OperatorNode pattern. + +**Files:** +- Modify: `src/orcapod/core/nodes/function_node.py` + +- [ ] **Step 1: Write tests for the no-DB async delegation** + +Create `tests/test_core/nodes/test_function_job_node_async_execute.py` with just the no-DB tests: + +```python +"""Tests for FunctionJobNode.async_execute() — unified async execution path.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import pyarrow as pa +import pytest + +from orcapod.channels import Channel +from orcapod.core.data_function import PythonDataFunction +from orcapod.core.function_pod import FunctionPod +from orcapod.core.nodes.function_node import FunctionJobNode +from orcapod.core.sources import ArrowTableSource +from orcapod.databases import InMemoryArrowDatabase +from orcapod.types import PodConfig + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _make_source(n: int = 3) -> ArrowTableSource: + """ArrowTableSource with tag=key (large_string), data=value (int64).""" + table = pa.table( + { + "key": pa.array([f"k{i}" for i in range(n)], type=pa.large_string()), + "value": pa.array(list(range(n)), type=pa.int64()), + } + ) + return ArrowTableSource(table, tag_columns=["key"], infer_nullable=True) + + +def _make_no_db_node(n: int = 3, max_concurrency: int | None = None) -> FunctionJobNode: + """FunctionJobNode with no DB attached.""" + src = _make_source(n) + pod = FunctionPod( + PythonDataFunction(lambda value: value * 2, output_keys="result"), + pod_config=PodConfig(max_concurrency=max_concurrency), + ) + return FunctionJobNode(pod, src) + + +async def _run_node(node: FunctionJobNode) -> list[tuple]: + """Feed the node's own input stream through async_execute and collect results.""" + input_ch: Channel = Channel(buffer_size=32) + output_ch: Channel = Channel(buffer_size=32) + + async def feed(): + for tag, data in node._input_stream.iter_data(): + await input_ch.writer.send((tag, data)) + await input_ch.writer.close() + + await asyncio.gather( + feed(), + node.async_execute(input_ch.reader, output_ch.writer), + ) + return await output_ch.reader.collect() + + +# --------------------------------------------------------------------------- +# No-DB path +# --------------------------------------------------------------------------- + + +class TestFunctionJobNodeAsyncExecuteNoDB: + @pytest.mark.asyncio + async def test_all_items_processed_and_forwarded(self): + """All input items reach the output channel.""" + node = _make_no_db_node(n=3) + results = await _run_node(node) + assert len(results) == 3 + values = sorted(data.as_dict()["result"] for _, data in results) + assert values == [0, 2, 4] + + @pytest.mark.asyncio + async def test_output_channel_closed_after_execution(self): + """Output channel is closed even if execution completes without error.""" + node = _make_no_db_node(n=2) + input_ch: Channel = Channel(buffer_size=8) + output_ch: Channel = Channel(buffer_size=8) + + async def feed(): + for tag, data in node._input_stream.iter_data(): + await input_ch.writer.send((tag, data)) + await input_ch.writer.close() + + await asyncio.gather(feed(), node.async_execute(input_ch.reader, output_ch.writer)) + # collect() should return immediately (channel closed), not hang. + results = await output_ch.reader.collect() + assert isinstance(results, list) + + @pytest.mark.asyncio + async def test_max_concurrency_limits_tasks(self): + """With max_concurrency=1, at most one task runs at a time.""" + concurrent_count = 0 + max_observed = 0 + + async def slow_double(value: int) -> int: + nonlocal concurrent_count, max_observed + concurrent_count += 1 + max_observed = max(max_observed, concurrent_count) + await asyncio.sleep(0.01) + concurrent_count -= 1 + return value * 2 + + pf = PythonDataFunction(slow_double, output_keys="result") + # Patch async_call to be the async function + import functools + original_async_call = pf.async_call + + async def patched_async_call(data, **kwargs): + nonlocal concurrent_count, max_observed + concurrent_count += 1 + max_observed = max(max_observed, concurrent_count) + await asyncio.sleep(0.01) + concurrent_count -= 1 + return await original_async_call(data, **kwargs) + + pf.async_call = patched_async_call # type: ignore[method-assign] + + src = _make_source(5) + pod = FunctionPod(pf, pod_config=PodConfig(max_concurrency=1)) + node = FunctionJobNode(pod, src) + results = await _run_node(node) + assert len(results) == 5 + assert max_observed <= 1 +``` + +- [ ] **Step 2: Run the tests to confirm they pass (they should with existing code)** + +```bash +uv run pytest tests/test_core/nodes/test_function_job_node_async_execute.py::TestFunctionJobNodeAsyncExecuteNoDB -v +``` + +Expected: all 3 tests PASS (existing code already handles no-DB). If any fail, investigate before +proceeding. + +- [ ] **Step 3: Replace the no-DB path in `FunctionJobNode.async_execute()` with delegation** + +In `src/orcapod/core/nodes/function_node.py`, find the else branch at lines 2146–2166: + +```python +else: + # Simple async execution without DB + async with asyncio.TaskGroup() as tg: + async for tag, data in input_channel: + async def _guarded_simple( + t: TagProtocol = tag, p: DataProtocol = data + ) -> None: + try: + await self._async_execute_one_data( + t, p, output, + observer=ctx_obs, + node_label=node_label, + node_hash=node_hash, + ) + finally: + if sem is not None: + sem.release() + + if sem is not None: + await sem.acquire() + tg.create_task(_guarded_simple()) +``` + +Replace the entire else block with: + +```python +else: + # No-DB path: delegate dispatch entirely to the pod. + await self._function_pod.async_execute( + [input_channel], output, observer=ctx_obs + ) + return # output is already closed by the pod's finally; skip outer finally +``` + +Wait — there is a problem here. The outer `finally: await output.close()` will try to close the +output channel a second time (it was already closed by `_function_pod.async_execute()`). Looking +at `_ChannelWriter.close()`: + +```python +async def close(self) -> None: + if not self._channel._closed.is_set(): + self._channel._closed.set() + await self._channel._queue.put(_CLOSED) +``` + +It's idempotent (guarded by `_closed.is_set()`). Double-close is safe. So `return` is not needed; +just replace the else block: + +```python +else: + # No-DB path: delegate dispatch entirely to the pod. + # The pod's async_execute() closes `output` in its own finally. + await self._function_pod.async_execute( + [input_channel], output, observer=ctx_obs + ) +``` + +The outer `finally: await output.close()` will be a no-op since the channel is already closed. + +Also remove the semaphore setup that now only applies to the DB path (lines 2089–2097): + +```python +# REMOVE these lines from async_execute(): +pod_config = getattr(self._function_pod, "pod_config", PodConfig()) +max_concurrency = resolve_concurrency(pod_config, PipelineConfig()) + +sem = ( + asyncio.Semaphore(max_concurrency) + if max_concurrency is not None + else None +) +``` + +- [ ] **Step 4: Delete `_async_execute_one_data`** + +Find and delete the entire `_async_execute_one_data` method (lines 2172–2201): + +```python +async def _async_execute_one_data( + self, + tag: TagProtocol, + data: DataProtocol, + output: "WritableChannel[tuple[TagProtocol, DataProtocol]]", + *, + observer: ExecutionObserverProtocol, + node_label: str, + node_hash: str, +) -> None: + ... +``` + +Delete this method entirely. + +- [ ] **Step 5: Remove dead imports from `function_node.py`** + +After removing the semaphore setup from `async_execute()`, `PipelineConfig`, `PodConfig`, and +`resolve_concurrency` are no longer referenced in `function_node.py`. Remove them from the import: + +Before: +```python +from orcapod.types import ( + ColumnConfig, + ContentHash, + NodeConfig, + PipelineConfig, + PodConfig, + Schema, + resolve_concurrency, +) +``` + +After: +```python +from orcapod.types import ( + ColumnConfig, + ContentHash, + NodeConfig, + Schema, +) +``` + +Verify by running a grep to confirm neither `PipelineConfig`, `PodConfig`, nor `resolve_concurrency` +appear elsewhere in the file: + +```bash +grep -n "PipelineConfig\|PodConfig\|resolve_concurrency" \ + src/orcapod/core/nodes/function_node.py +``` + +Expected: no output. + +- [ ] **Step 6: Run all tests** + +```bash +uv run pytest tests/ -v --timeout=30 +``` + +Expected: all tests still pass, including `TestFunctionJobNodeAsyncExecuteNoDB`. + +- [ ] **Step 7: Commit** + +```bash +git add src/orcapod/core/nodes/function_node.py \ + tests/test_core/nodes/test_function_job_node_async_execute.py +git commit -m "refactor(function_node): delegate no-DB async_execute to pod + +FunctionJobNode.async_execute() no longer reimplements the semaphore ++TaskGroup dispatch loop for the no-DB path. It delegates directly to +self._function_pod.async_execute(), matching the OperatorNode pattern. +Deletes _async_execute_one_data (dead after delegation). Removes the +now-dead PipelineConfig, PodConfig, and resolve_concurrency imports." +``` + +--- + +## Task 3: `FunctionJobNode.async_execute()` — DB 3-stage concurrent pipeline + +Redesign the DB path into three concurrent coroutines: + +1. **`route_inputs`** — routes cache hits directly to output, stamps cache misses with a + correlation key and forwards to `compute_channel`. +2. **`execution_pod.async_execute()`** — reads `compute_channel`, writes `result_channel` + (this is now `_FunctionPodBase.async_execute()`). +3. **`record_and_forward`** — reads `result_channel`, calls `add_pipeline_record`, strips the + correlation key from the output tag, writes to `output`. + +Covers both the persistent (`_cached_function_pod`, `is_ephemeral=False`) and ephemeral +(`_ephemeral_cached_pod`, `is_ephemeral=True`) sub-paths — only the pod and the flag differ. + +**Files:** +- Modify: `src/orcapod/core/nodes/function_node.py` +- Modify: `tests/test_core/nodes/test_function_job_node_async_execute.py` + +- [ ] **Step 1: Write failing tests for the DB path** + +Add the following test classes to `tests/test_core/nodes/test_function_job_node_async_execute.py`. + +First, add the `_TAG_NODE_INPUT_REF` constant import at the top of the test file so we can assert +the correlation key is absent: + +```python +# Add to imports at the top: +from orcapod.core.nodes.function_node import _TAG_NODE_INPUT_REF # private constant +from orcapod.system_constants import constants as _constants +``` + +Then add these test classes at the bottom of the file: + +```python +# --------------------------------------------------------------------------- +# Helpers — observer spy +# --------------------------------------------------------------------------- + + +class _SpyObserver: + """Captures all observer hook invocations.""" + + def __init__(self): + self.node_starts: list[str] = [] + self.node_ends: list[str] = [] + self.data_starts: list[tuple] = [] + self.data_ends: list[tuple] = [] # (label, tag, input_data, output_data, cached) + self.data_crashes: list[tuple] = [] + + def contextualize(self, *path): + return self + + def on_node_start(self, node_label, node_hash, *, tag_schema=None): + self.node_starts.append(node_label) + + def on_node_end(self, node_label, node_hash): + self.node_ends.append(node_label) + + def on_data_start(self, node_label, tag, data): + self.data_starts.append((node_label, tag, data)) + + def on_data_end(self, node_label, tag, input_data, output_data, *, cached): + self.data_ends.append((node_label, tag, input_data, output_data, cached)) + + def on_data_crash(self, node_label, tag, data, error): + self.data_crashes.append((node_label, tag, data, error)) + + def create_data_logger(self, tag, data): + return None + + +# --------------------------------------------------------------------------- +# Helpers — DB node fixtures +# --------------------------------------------------------------------------- + + +def _make_db_node(n: int = 3) -> tuple[FunctionJobNode, InMemoryArrowDatabase, InMemoryArrowDatabase]: + """FunctionJobNode with persistent pipeline + result databases.""" + src = _make_source(n) + pod = FunctionPod( + PythonDataFunction(lambda value: value * 2, output_keys="result"), + ) + pipeline_db = InMemoryArrowDatabase() + result_db = InMemoryArrowDatabase() + node = FunctionJobNode( + pod, src, pipeline_database=pipeline_db, result_database=result_db + ) + return node, pipeline_db, result_db + + +# --------------------------------------------------------------------------- +# DB path — basic flow +# --------------------------------------------------------------------------- + + +class TestFunctionJobNodeAsyncExecuteDB: + + @pytest.mark.asyncio + async def test_all_items_processed_and_forwarded(self): + """All input items are computed and forwarded via the DB path.""" + node, pipeline_db, result_db = _make_db_node(3) + results = await _run_node(node) + assert len(results) == 3 + values = sorted(data.as_dict()["result"] for _, data in results) + assert values == [0, 2, 4] + + @pytest.mark.asyncio + async def test_pipeline_records_written_for_each_item(self): + """add_pipeline_record is called once per cache miss → one row per item.""" + node, pipeline_db, _ = _make_db_node(3) + await _run_node(node) + records = pipeline_db.get_all_records(node.node_identity_path) + assert records is not None + assert records.num_rows == 3 + + @pytest.mark.asyncio + async def test_correlation_key_absent_from_output_tags(self): + """Output tags must not contain the internal _tag_node_input_ref column.""" + node, _, _ = _make_db_node(3) + results = await _run_node(node) + full_key = f"{_constants.META_PREFIX}{_TAG_NODE_INPUT_REF}" + for tag, _ in results: + meta = tag.get_meta_info() + assert full_key not in meta, ( + f"Output tag leaked internal correlation key {full_key!r}" + ) + + @pytest.mark.asyncio + async def test_cache_hits_emitted_without_recomputation(self): + """Items already in the pipeline DB are served from cache (no recomputation).""" + node, _, _ = _make_db_node(3) + + # First run — populates both pipeline DB and result DB. + first_results = await _run_node(node) + assert len(first_results) == 3 + + # Reset in-memory state so async_execute re-runs from scratch. + node.clear_cache() + + call_count = 0 + original_process = node._function_pod.data_function.call + + def counted_call(data, **kw): + nonlocal call_count + call_count += 1 + return original_process(data, **kw) + + node._function_pod.data_function.call = counted_call # type: ignore[method-assign] + + second_results = await _run_node(node) + assert len(second_results) == 3 + # All 3 were hits — data function should NOT have been called again. + assert call_count == 0 + + @pytest.mark.asyncio + async def test_observer_cached_false_for_misses_cached_true_for_hits(self): + """Cache hits emit on_data_end(cached=True); misses emit on_data_end(cached=False).""" + node, _, _ = _make_db_node(2) + + spy = _SpyObserver() + + # First run — both are misses. + await _run_node(node) # populates DB + node.clear_cache() + + spy = _SpyObserver() + input_ch: Channel = Channel(buffer_size=8) + output_ch: Channel = Channel(buffer_size=8) + + async def feed(): + for tag, data in node._input_stream.iter_data(): + await input_ch.writer.send((tag, data)) + await input_ch.writer.close() + + await asyncio.gather( + feed(), + node.async_execute(input_ch.reader, output_ch.writer, observer=spy), + ) + await output_ch.reader.collect() + + # Second run was all hits. + assert len(spy.data_ends) == 2 + assert all(cached is True for _, _, _, _, cached in spy.data_ends) + + @pytest.mark.asyncio + async def test_semaphore_not_leaked_on_computation_crash(self): + """When a data function raises, processing continues for remaining items.""" + call_count = 0 + + def sometimes_fail(value: int) -> int: + nonlocal call_count + call_count += 1 + if value == 1: + raise ValueError("deliberate failure") + return value * 2 + + src = _make_source(3) + pod = FunctionPod( + PythonDataFunction(sometimes_fail, output_keys="result"), + pod_config=PodConfig(max_concurrency=1), + ) + pipeline_db = InMemoryArrowDatabase() + result_db = InMemoryArrowDatabase() + node = FunctionJobNode(pod, src, pipeline_database=pipeline_db, result_database=result_db) + + results = await _run_node(node) + # value=0 → 0, value=1 → crash (skipped), value=2 → 4 + assert len(results) == 2 + values = sorted(data.as_dict()["result"] for _, data in results) + assert values == [0, 4] + + +# --------------------------------------------------------------------------- +# DB path — ephemeral +# --------------------------------------------------------------------------- + + +class TestFunctionJobNodeAsyncExecuteEphemeral: + + @pytest.mark.asyncio + async def test_ephemeral_pipeline_records_have_is_ephemeral_true(self): + """Ephemeral execution writes pipeline records with is_ephemeral=True.""" + from orcapod.system_constants import constants + from orcapod.types import NodeConfig + + src = _make_source(2) + pod = FunctionPod( + PythonDataFunction(lambda value: value * 2, output_keys="result"), + ) + pipeline_db = InMemoryArrowDatabase() + result_db = InMemoryArrowDatabase() + ephemeral_db = InMemoryArrowDatabase() + node = FunctionJobNode( + pod, src, pipeline_database=pipeline_db, result_database=result_db + ) + node.node_config = NodeConfig(is_result_ephemeral=True) + node.set_ephemeral_store(ephemeral_db) + + results = await _run_node(node) + assert len(results) == 2 + + records = pipeline_db.get_all_records(node.node_identity_path) + assert records is not None + assert records.num_rows == 2 + # All rows should have is_ephemeral=True + is_eph_col = records.column(constants.IS_EPHEMERAL_COL).to_pylist() + assert all(v is True for v in is_eph_col) + + @pytest.mark.asyncio + async def test_ephemeral_missing_store_raises_runtime_error(self): + """Raising RuntimeError when is_result_ephemeral=True but no store set.""" + from orcapod.types import NodeConfig + + src = _make_source(1) + pod = FunctionPod( + PythonDataFunction(lambda value: value * 2, output_keys="result"), + ) + pipeline_db = InMemoryArrowDatabase() + result_db = InMemoryArrowDatabase() + node = FunctionJobNode( + pod, src, pipeline_database=pipeline_db, result_database=result_db + ) + node.node_config = NodeConfig(is_result_ephemeral=True) + # Deliberately do NOT call set_ephemeral_store() + + with pytest.raises(RuntimeError, match="ephemeral"): + await _run_node(node) +``` + +- [ ] **Step 2: Run the tests to confirm they fail** + +```bash +uv run pytest tests/test_core/nodes/test_function_job_node_async_execute.py::TestFunctionJobNodeAsyncExecuteDB \ + tests/test_core/nodes/test_function_job_node_async_execute.py::TestFunctionJobNodeAsyncExecuteEphemeral \ + -v +``` + +Expected: most or all tests FAIL because the 3-stage pipeline isn't implemented yet. +The test importing `_TAG_NODE_INPUT_REF` will cause an `ImportError` — that's expected. + +- [ ] **Step 3: Add the `_TAG_NODE_INPUT_REF` module constant to `function_node.py`** + +Near the top of `function_node.py` where the other private constants are defined (after +`_PIPELINE_RECOMPUTATION_INDEX_COL`), add: + +```python +# Private meta-column name stamped into tags routed to the compute channel. +# Carries the correlation key (bytes) that links an execution result back to +# its original (tag, input_data) pair. Stripped from output tags in +# record_and_forward() before downstream emission. +_TAG_NODE_INPUT_REF = "_tag_node_input_ref" +``` + +- [ ] **Step 4: Implement the 3-stage DB pipeline in `FunctionJobNode.async_execute()`** + +Find the DB path block in `async_execute()` (the `if self._cached_function_pod is not None:` branch, +starting around line 2103 after earlier edits). Replace the entire block with the 3-stage pipeline. + +The complete new `async_execute()` method body (replacing from after the UNAVAILABLE check to the +final `finally`) is: + +```python + from orcapod.pipeline.observer import NoOpObserver + + node_label = self.label + node_hash = self.content_hash().to_string() + + obs = observer if observer is not None else NoOpObserver() + ctx_obs = obs.contextualize(*self.node_identity_path) + + try: + tag_schema = self._input_stream.output_schema(columns={"system_tags": True})[0] + ctx_obs.on_node_start(node_label, node_hash, tag_schema=tag_schema) + + if self._cached_function_pod is not None: + # ---------------------------------------------------------- + # DB path — 3-stage concurrent pipeline + # ---------------------------------------------------------- + is_ephemeral = bool(self._node_config.is_result_ephemeral) + if is_ephemeral: + if self._ephemeral_cached_pod is None: + raise RuntimeError( + f"FunctionJobNode '{self.label}' has is_result_ephemeral=True " + "but no ephemeral store has been assigned. Call " + "set_ephemeral_store() with an ArrowDatabaseProtocol before " + "executing this node." + ) + execution_pod = self._ephemeral_cached_pod + else: + execution_pod = self._cached_function_pod + + # Phase 1: load pipeline DB → in-memory cache keyed by base_entry_id. + loaded = self._load_cached_entries() + self._cached_output_datas.update(loaded) + if loaded: + self._cached_output_table = None + self._cached_content_hash_column = None + cached_by_base_entry_id: dict[bytes, tuple[TagProtocol, DataProtocol]] = dict(loaded) + + # Intermediate channels (bounded for backpressure). + compute_channel: Channel[tuple[TagProtocol, DataProtocol]] = Channel(buffer_size=16) + result_channel: Channel[tuple[TagProtocol, DataProtocol]] = Channel(buffer_size=16) + + # Local dict: correlation_key → (original_tag, original_input_data) + input_store: dict[bytes, tuple[TagProtocol, DataProtocol]] = {} + + async def route_inputs() -> None: + """Stage 1: send cache hits to output; stamp misses for computation.""" + try: + async for tag, data in input_channel: + base_entry_id = self.compute_base_entry_id(tag, data) + if base_entry_id in cached_by_base_entry_id: + cached_tag, cached_data = cached_by_base_entry_id[base_entry_id] + ctx_obs.on_data_start(node_label, tag, data) + ctx_obs.on_data_end( + node_label, tag, data, cached_data, cached=True + ) + await output.send((cached_tag, cached_data)) + else: + correlation_key = uuid.uuid4().bytes + input_store[correlation_key] = (tag, data) + stamped_tag = tag.with_meta_columns( + **{_TAG_NODE_INPUT_REF: correlation_key} + ) + await compute_channel.writer.send((stamped_tag, data)) + finally: + await compute_channel.writer.close() + + async def record_and_forward() -> None: + """Stage 3: record pipeline entry, strip key, emit to output.""" + async for output_tag, output_data in result_channel.reader: + correlation_key = output_tag.get_meta_value(_TAG_NODE_INPUT_REF) + original_tag, input_data = input_store.pop(correlation_key) + result_computed = bool( + output_data.get_meta_value( + execution_pod.RESULT_COMPUTED_FLAG, False + ) + ) + self.add_pipeline_record( + original_tag, + input_data, + data_record_id=output_data.datagram_uuid, + computed=result_computed, + is_ephemeral=is_ephemeral, + ) + # Update in-memory cache so iter_data() sees the result. + base_entry_id = self.compute_base_entry_id(original_tag, input_data) + clean_tag = output_tag.drop_meta_columns(_TAG_NODE_INPUT_REF) + self._cached_output_datas[base_entry_id] = (clean_tag, output_data) + self._cached_output_table = None + self._cached_content_hash_column = None + await output.send((clean_tag, output_data)) + + async with asyncio.TaskGroup() as tg: + tg.create_task(route_inputs()) + tg.create_task( + execution_pod.async_execute( + [compute_channel.reader], + result_channel.writer, + observer=ctx_obs, + ) + ) + tg.create_task(record_and_forward()) + + else: + # No-DB path: delegate dispatch entirely to the pod. + await self._function_pod.async_execute( + [input_channel], output, observer=ctx_obs + ) + + ctx_obs.on_node_end(node_label, node_hash) + finally: + await output.close() +``` + +**Important:** The `Channel` type needs to be imported in `function_node.py`. Verify the import +already exists (it should since it's in `orcapod.channels`): + +```python +from orcapod.channels import ReadableChannel, WritableChannel +``` + +You also need `Channel` itself for the intermediate channels. Add it: + +```python +from orcapod.channels import Channel, ReadableChannel, WritableChannel +``` + +- [ ] **Step 5: Run all tests** + +```bash +uv run pytest tests/ -v --timeout=30 +``` + +Expected: all tests pass, including all new DB and ephemeral tests. + +If `test_cache_hits_emitted_without_recomputation` fails, check that `clear_cache()` properly +resets `_cached_output_datas` (it should — it's an existing method). + +If `test_correlation_key_absent_from_output_tags` fails, check that `drop_meta_columns` is +called with the right key in `record_and_forward`. + +- [ ] **Step 6: Commit** + +```bash +git add src/orcapod/core/nodes/function_node.py \ + tests/test_core/nodes/test_function_job_node_async_execute.py +git commit -m "feat(function_node): redesign async_execute as 3-stage concurrent pipeline + +FunctionJobNode.async_execute() now delegates computation to +execution_pod.async_execute() (CachedFunctionPod inheriting the loop +from _FunctionPodBase) via bounded compute_channel/result_channel. +Three asyncio.TaskGroup tasks run concurrently: + 1. route_inputs: cache hits → output, misses → compute_channel (stamped) + 2. execution_pod.async_execute: compute_channel → result_channel + 3. record_and_forward: add_pipeline_record, strip key, emit to output +Both persistent and ephemeral DB sub-paths use the same pipeline. +Adds _TAG_NODE_INPUT_REF module constant (private, never in public API). +Imports Channel for intermediate channels." +``` + +--- + +## Task 4: Cleanup — remove orphan test, relocate new observer tests + +The `TestAsyncExecuteBackpressure` in `test_regression_fixes.py` tested +`FunctionPod.async_execute()` in isolation — a method that was an orphan (never called by +the orchestrator). Its equivalent coverage now lives in `TestFunctionJobNodeAsyncExecuteNoDB.test_max_concurrency_limits_tasks`. Remove it. + +Also relocate the `TestFunctionPodBaseAsyncExecuteObserver` test added in Task 1 from +`test_regression_fixes.py` to its natural home (or keep in place — it tests pod-level observer +behavior, and `test_regression_fixes.py` is an acceptable home). Move it to the dedicated file +for cleanliness. + +**Files:** +- Modify: `tests/test_core/test_regression_fixes.py` +- Modify: `tests/test_core/nodes/test_function_job_node_async_execute.py` + +- [ ] **Step 1: Remove `TestAsyncExecuteBackpressure` from `test_regression_fixes.py`** + +Find and delete the entire `TestAsyncExecuteBackpressure` class (lines 293–343 in the original +file, may have shifted). It begins with: + +```python +class TestAsyncExecuteBackpressure: + """With max_concurrency set, pending tasks should be bounded.""" + ... +``` + +Delete from the class definition through the end of the last test in the class. + +Also remove the section header comment above it: +```python +# =========================================================================== +# 4. FunctionPod.async_execute backpressure bounds pending tasks +# =========================================================================== +``` + +Update the module-level docstring to remove item 4: + +```python +""" +Regression tests for bugs fixed in the data-function-executor-system branch. + +Covers: +1. async_execute output channel closed on exception (try/finally) +2. DataFunctionWrapper.direct_call/direct_async_call bypass executor routing +3. Concurrent iteration falls back to sequential inside a running event loop +4. _materialize_to_stream preserves source_info provenance tokens +5. RayExecutor._ensure_ray_initialized uses ray_address +6. DataFunctionExecutorProtocol uses DataFunctionProtocol (not Any) +""" +``` + +(Renumber items 5-7 → 4-6.) + +- [ ] **Step 2: Move `TestFunctionPodBaseAsyncExecuteObserver` to the dedicated test file** + +Remove `TestFunctionPodBaseAsyncExecuteObserver` from the bottom of `test_regression_fixes.py` +(added in Task 1). + +Add it to `tests/test_core/nodes/test_function_job_node_async_execute.py` — or optionally to +a separate `tests/test_core/function_pod/test_function_pod_base_async_execute.py`. Since these +tests test pod behavior (not node behavior), place them in a new file: + +Create `tests/test_core/function_pod/test_function_pod_async_execute.py`: + +```python +"""Tests for _FunctionPodBase.async_execute() — observer hooks and backpressure.""" + +from __future__ import annotations + +import asyncio + +import pyarrow as pa +import pytest + +from orcapod.channels import Channel +from orcapod.core.data_function import PythonDataFunction +from orcapod.core.function_pod import FunctionPod +from orcapod.core.streams.arrow_table_stream import ArrowTableStream +from orcapod.types import PodConfig + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_stream(n: int = 3) -> ArrowTableStream: + schema = pa.schema( + [pa.field("id", pa.int64(), nullable=False), pa.field("x", pa.int64(), nullable=False)] + ) + table = pa.table( + { + "id": pa.array(list(range(n)), type=pa.int64()), + "x": pa.array(list(range(n)), type=pa.int64()), + }, + schema=schema, + ) + return ArrowTableStream(table, tag_columns=["id"]) + + +async def _feed(stream: ArrowTableStream, ch: Channel) -> None: + for tag, data in stream.iter_data(): + await ch.writer.send((tag, data)) + await ch.writer.close() + + +class _SpyObserver: + def __init__(self): + self.starts = [] + self.ends = [] # [(label, tag, input, output, cached)] + self.crashes = [] + + def contextualize(self, *path): + return self + + def on_node_start(self, *a, **kw): pass + def on_node_end(self, *a, **kw): pass + + def on_data_start(self, label, tag, data): + self.starts.append((label, tag, data)) + + def on_data_end(self, label, tag, inp, out, *, cached): + self.ends.append((label, tag, inp, out, cached)) + + def on_data_crash(self, label, tag, data, exc): + self.crashes.append(exc) + + def create_data_logger(self, tag, data): + return None + + +# --------------------------------------------------------------------------- +# Observer hooks +# --------------------------------------------------------------------------- + + +class TestFunctionPodBaseAsyncExecuteObserver: + @pytest.mark.asyncio + async def test_on_data_start_and_end_fire_per_item(self): + """on_data_start fires before and on_data_end(cached=False) after each item.""" + pf = PythonDataFunction(lambda x: x * 2, output_keys="result") + pod = FunctionPod(pf) + stream = _make_stream(3) + input_ch: Channel = Channel(buffer_size=16) + output_ch: Channel = Channel(buffer_size=16) + await _feed(stream, input_ch) + + spy = _SpyObserver() + await pod.async_execute([input_ch.reader], output_ch.writer, observer=spy) + results = await output_ch.reader.collect() + + assert len(results) == 3 + assert len(spy.starts) == 3 + assert len(spy.ends) == 3 + assert all(cached is False for _, _, _, _, cached in spy.ends) + + @pytest.mark.asyncio + async def test_on_data_crash_fires_not_on_data_end(self): + """on_data_crash fires (and on_data_end does not) when processing raises.""" + def boom(x: int) -> int: + raise ValueError("bang") + + pf = PythonDataFunction(boom, output_keys="result") + pod = FunctionPod(pf) + stream = _make_stream(2) + input_ch: Channel = Channel(buffer_size=16) + output_ch: Channel = Channel(buffer_size=16) + await _feed(stream, input_ch) + + spy = _SpyObserver() + await pod.async_execute([input_ch.reader], output_ch.writer, observer=spy) + results = await output_ch.reader.collect() + + assert len(results) == 0 + assert len(spy.crashes) == 2 + assert len(spy.ends) == 0 + + +# --------------------------------------------------------------------------- +# Backpressure +# --------------------------------------------------------------------------- + + +class TestFunctionPodBaseAsyncExecuteBackpressure: + @pytest.mark.asyncio + async def test_max_concurrency_limits_concurrent_tasks(self): + """With max_concurrency=1, at most one task runs concurrently.""" + concurrent_count = 0 + max_observed = 0 + + async def track(x: int) -> int: + nonlocal concurrent_count, max_observed + concurrent_count += 1 + max_observed = max(max_observed, concurrent_count) + await asyncio.sleep(0.01) + concurrent_count -= 1 + return x * 2 + + pf = PythonDataFunction(lambda x: x * 2, output_keys="result") + original_async_call = pf.async_call + + async def patched(data, **kw): + nonlocal concurrent_count, max_observed + concurrent_count += 1 + max_observed = max(max_observed, concurrent_count) + await asyncio.sleep(0.01) + concurrent_count -= 1 + return await original_async_call(data, **kw) + + pf.async_call = patched # type: ignore[method-assign] + + pod = FunctionPod(pf, pod_config=PodConfig(max_concurrency=1)) + stream = _make_stream(5) + input_ch: Channel = Channel(buffer_size=32) + output_ch: Channel = Channel(buffer_size=32) + await _feed(stream, input_ch) + await pod.async_execute([input_ch.reader], output_ch.writer) + results = await output_ch.reader.collect() + assert len(results) == 5 + assert max_observed <= 1 +``` + +- [ ] **Step 3: Run the full test suite** + +```bash +uv run pytest tests/ -v --timeout=30 +``` + +Expected: all tests pass. Verify: +- `TestAsyncExecuteBackpressure` is gone +- `TestFunctionPodBaseAsyncExecuteBackpressure` in the new file covers equivalent behavior +- All `TestFunctionJobNodeAsyncExecuteDB` and `TestFunctionJobNodeAsyncExecuteEphemeral` pass + +- [ ] **Step 4: Final import cleanup — verify `function_pod.py`** + +The observer types are used only in the `TYPE_CHECKING` block. Confirm the import is present +and correct: + +```bash +grep -n "ExecutionObserverProtocol" src/orcapod/core/function_pod.py +``` + +Expected: one line in the `TYPE_CHECKING` block. + +Also verify `PipelineConfig`, `PodConfig`, `resolve_concurrency` are still present in +`function_pod.py` (they ARE used there — `async_execute` on `_FunctionPodBase` still uses them): + +```bash +grep -n "PipelineConfig\|resolve_concurrency" src/orcapod/core/function_pod.py +``` + +Expected: lines referencing both in the import and in `async_execute`. + +- [ ] **Step 5: Commit** + +```bash +git add tests/test_core/test_regression_fixes.py \ + tests/test_core/nodes/test_function_job_node_async_execute.py \ + tests/test_core/function_pod/test_function_pod_async_execute.py +git commit -m "test(async_execute): remove orphan backpressure test, add dedicated coverage + +Removes TestAsyncExecuteBackpressure from test_regression_fixes.py +(tested an orphan FunctionPod method, never called by orchestrator). +Equivalent coverage lives in TestFunctionJobNodeAsyncExecuteNoDB and +TestFunctionPodBaseAsyncExecuteBackpressure. Adds: +- tests/test_core/function_pod/test_function_pod_async_execute.py: + observer hooks and backpressure via _FunctionPodBase.async_execute() +- tests/test_core/nodes/test_function_job_node_async_execute.py: + full FunctionJobNode.async_execute() coverage (no-DB, DB, ephemeral)" +``` + +--- + +## Self-Review + +### Spec coverage + +| Spec requirement | Task | +|---|---| +| `async_execute()` moves to `_FunctionPodBase` | Task 1 | +| `observer` parameter on `_FunctionPodBase.async_execute()` | Task 1 | +| `on_data_start`, `on_data_end`, `on_data_crash` per item | Task 1 | +| No-DB path delegates to `self._function_pod.async_execute()` | Task 2 | +| Delete orphan `_async_execute_one_data` | Task 2 | +| Dead import cleanup (PipelineConfig, PodConfig, resolve_concurrency) | Task 2 | +| `_TAG_NODE_INPUT_REF` private constant in `function_node.py` | Task 3 | +| 3-stage concurrent pipeline (route → compute → record) | Task 3 | +| route_inputs: cache hits → output, misses → compute_channel | Task 3 | +| record_and_forward: add_pipeline_record, strip key, emit | Task 3 | +| Bounded `compute_channel` + `result_channel` | Task 3 | +| Ephemeral path uses `_ephemeral_cached_pod`, `is_ephemeral=True` | Task 3 | +| RuntimeError guard when ephemeral=True but no store set | Task 3 | +| Output tags carry no `_tag_node_input_ref` | Task 3 (tested) | +| `add_pipeline_record` called once per cache miss | Task 3 (tested) | +| Remove `TestAsyncExecuteBackpressure` | Task 4 | +| New test file `test_function_job_node_async_execute.py` | Tasks 2–4 | +| CACHE_ONLY / UNAVAILABLE modes unchanged | ✓ preserved (no changes to those branches) | + +### Key invariants verification + +1. **Output tags carry no `_tag_node_input_ref`** — verified by `test_correlation_key_absent_from_output_tags` +2. **`add_pipeline_record` called once per cache miss** — verified by `test_pipeline_records_written_for_each_item` +3. **Cache hits bypass computation** — verified by `test_cache_hits_emitted_without_recomputation` +4. **`_FunctionPodBase.async_execute()` is unaware of pipeline DB** — architectural constraint, no DB imports in function_pod.py +5. **No-DB path uses `_function_pod`** — enforced by code structure (else branch) +6. **Ephemeral path uses `_ephemeral_cached_pod` + `is_ephemeral=True`** — verified by `test_ephemeral_pipeline_records_have_is_ephemeral_true` +7. **Persistent path uses `_cached_function_pod` + `is_ephemeral=False`** — verified by `test_pipeline_records_written_for_each_item` diff --git a/superpowers/specs/2026-07-09-itl-516-unify-async-execute-design.md b/superpowers/specs/2026-07-09-itl-516-unify-async-execute-design.md new file mode 100644 index 000000000..58bc75fda --- /dev/null +++ b/superpowers/specs/2026-07-09-itl-516-unify-async-execute-design.md @@ -0,0 +1,255 @@ +# ITL-516: Unify FunctionPod.async_execute() and FunctionJobNode.async_execute() Concurrency Paths + +**Date:** 2026-07-09 +**Issue:** [ITL-516](https://linear.app/enigma-metamorphic/issue/ITL-516) + +--- + +## Overview + +`FunctionPod.async_execute()` and `FunctionJobNode.async_execute()` are currently two independent +implementations of the same async dispatch loop. Investigation revealed that `FunctionPod.async_execute()` +is an orphan method: it is not part of `FunctionPodProtocol`, never called by the orchestrator, and +has no sync counterpart on `FunctionPod` (which uses `process()` → `FunctionPodStream` for its own +execution model). The method was written alongside `FunctionJobNode.async_execute()` as the concurrency +pattern was being developed, leaving two divergent copies. + +The correct model is the one OperatorNode already uses: `OperatorNode.async_execute()` delegates the +actual computation to `self._operator.async_execute()` and wraps only the DB-persistence concerns. +`FunctionJobNode` should do the same — delegate concurrent dispatch to `FunctionPod.async_execute()` +and own only the routing, caching, and pipeline-record concerns. + +--- + +## Goals & Success Criteria + +- `FunctionPod.async_execute()` is the single authoritative async dispatch loop, with observer support +- `FunctionJobNode.async_execute()` owns routing (cache hits vs misses), pipeline-record writing, and + observer lifecycle — but never re-implements the concurrent dispatch loop +- The three stages (routing, computation, recording) run concurrently via `asyncio.TaskGroup`, so + FunctionPod begins processing cache misses immediately as the router produces them +- Downstream nodes never receive tags carrying internal correlation metadata +- Existing behaviour is preserved: two-phase DB-backed execution for FunctionJobNode, standalone + streaming for FunctionPod +- Tests cover the unified path + +--- + +## Scope & Boundaries + +**In scope:** + +- Add `observer` parameter to `FunctionPod.async_execute()` +- Redesign `FunctionJobNode.async_execute()` as a 3-stage concurrent pipeline using + `compute_channel`, `result_channel`, and a correlation-key tag-stamp mechanism +- Delete the orphan `FunctionPod.async_execute()` backpressure regression test; add equivalent + coverage via `FunctionJobNode.async_execute()` +- Clean up any dead imports in `function_node.py` after removing the manual semaphore setup + (e.g. `PipelineConfig` if it is no longer referenced outside `async_execute`) + +**Out of scope:** + +- `FunctionPod.process_data()` / `async_process_data()` — unchanged +- `FunctionJobNode.execute()` sync path — unchanged +- Operator node execution — unchanged +- `NodeConfig` / `PodConfig` semantics (ITL-512) +- Changes to synchronous execution paths (`process_data`, `_process_data_internal`) + +--- + +## Design + +### Why FunctionPod.async_execute() was an orphan + +`FunctionPod` has no `execute()` method. Its sync model is `process(*streams)` → lazy +`FunctionPodStream`. `async_execute()` was written without a sync counterpart, without being +added to `FunctionPodProtocol`, and without being called by the orchestrator. The orchestrator +dispatches only to `SourceNodeProtocol`, `FunctionNodeProtocol`, and `OperatorNodeProtocol` — +FunctionPod is none of these. The method's sole caller was a single regression test. + +### The OperatorNode analogy + +`OperatorNode.async_execute()` delegates to `self._operator.async_execute()` in the no-DB path +and wraps it in a forwarding channel in the LOG path. It never reimplements the operator's logic. +This is the correct pattern: the pod owns computation, the node owns DB concerns. + +For function execution, the same separation applies: `FunctionPod.async_execute()` owns concurrent +dispatch; `FunctionJobNode.async_execute()` owns routing, caching, and pipeline recording. + +### async_execute() moves to _FunctionPodBase + +Currently `async_execute()` is defined only on `FunctionPod`. `CachedFunctionPod` (used for +both persistent and ephemeral execution) inherits from `WrappedFunctionPod → _FunctionPodBase` +— not from `FunctionPod` — so it has no `async_execute()` today. + +`_FunctionPodBase` already defines `async_process_data()`, which each subclass overrides +correctly (`FunctionPod` calls the data function, `CachedFunctionPod` checks the cache then +stores). Moving `async_execute()` to `_FunctionPodBase` gives all pod types the dispatch loop +for free via polymorphism: the loop calls `self.async_process_data()`, which resolves to the +right implementation at runtime. + +`FunctionPod` retains no override — it inherits the loop from `_FunctionPodBase`. + +Add an optional `observer` parameter (same pattern as all node `async_execute` signatures): + +```python +# on _FunctionPodBase +async def async_execute( + self, + inputs: Sequence[ReadableChannel[tuple[TagProtocol, DataProtocol]]], + output: WritableChannel[tuple[TagProtocol, DataProtocol]], + pipeline_config: PipelineConfig | None = None, + *, + observer: ExecutionObserverProtocol | None = None, +) -> None: +``` + +When `observer` is provided, fire `on_data_start`, `on_data_end`, and `on_data_crash` per item. +When `None`, behaviour is unchanged (no-op). All pod types — `FunctionPod`, `CachedFunctionPod` +(persistent), `CachedFunctionPod` (ephemeral) — are now valid delegation targets from +`FunctionJobNode`. + +### FunctionJobNode.async_execute() — 3-stage concurrent pipeline + +#### No-DB path (simple delegation) + +When `_cached_function_pod is None`, delegate directly — no routing, no recording: + +```python +await self._function_pod.async_execute( + [input_channel], output, observer=ctx_obs +) +``` + +Mirrors OperatorNode's no-DB delegation. + +#### DB path (3-stage concurrent pipeline) + +Before entering the TaskGroup, select the execution pod and `is_ephemeral` flag based on +`_node_config.is_result_ephemeral`: + +```python +if self._node_config.is_result_ephemeral: + execution_pod = self._ephemeral_cached_pod # CachedFunctionPod → ephemeral DB + is_ephemeral = True +else: + execution_pod = self._cached_function_pod # CachedFunctionPod → persistent DB + is_ephemeral = False +``` + +Both are `CachedFunctionPod` instances that now inherit `async_execute()` from `_FunctionPodBase`. +Their `async_process_data()` overrides handle the correct store (ephemeral vs persistent) and +set `RESULT_COMPUTED_FLAG` on the output data. The 3-stage pipeline is identical for both; +only the pod and the flag differ. + +Three coroutines run inside a single `asyncio.TaskGroup`: + +``` +input_channel + │ + ▼ route_inputs() — task 1 + ├── cache hit → emit to output (observer: on_data_start + on_data_end cached=True) + └── cache miss → stamp tag with correlation key → compute_channel + │ + ▼ execution_pod.async_execute() — task 2 + (reads compute_channel, writes result_channel) + result_channel + │ + ▼ record_and_forward() — task 3 + read correlation key from output_tag + → look up (original_tag, input_data) in input_store + → add_pipeline_record(original_tag, input_data, + is_ephemeral=is_ephemeral, ...) + → strip correlation key from output_tag + → emit (clean_tag, output_data) to output +``` + +**Backpressure:** `compute_channel` and `result_channel` are bounded channels. If the execution +pod is slower than the router, `route_inputs` blocks on `compute_channel.writer.send()`; if the +recorder is slower than the pod, the pod blocks on `result_channel.writer.send()`. + +**TaskGroup lifecycle:** The TaskGroup exits only when all three tasks complete. `route_inputs` +closes `compute_channel.writer` when the input is exhausted, signalling the pod to finish. The +pod closes `result_channel.writer` in its `finally`, signalling the recorder to finish. +`FunctionJobNode.async_execute()`'s own `finally` closes the final `output`. + +### Correlation key mechanism + +To give the recorder access to `(original_tag, input_data)` — needed by `add_pipeline_record` — +without changing `result_channel`'s type or FunctionPod's interface: + +1. **Router** generates a correlation key (`bytes`, e.g. `uuid.uuid4().bytes`) per cache miss, + stamps it into the tag as a private meta column (`_tag_node_input_ref`), stores + `{key: (original_tag, input_data)}` in a local dict (`input_store`), then sends the + stamped tag + data to `compute_channel`. + +2. **FunctionPod** passes the tag through completely unchanged — the stamp travels to + `result_channel` without modification. + +3. **Recorder** reads `_tag_node_input_ref` from the output tag, retrieves + `(original_tag, input_data)` from `input_store`, calls `add_pipeline_record`, then + **strips the `_tag_node_input_ref` column** from the output tag before emitting to `output`. + +Stripping the correlation key before downstream emission is a hard correctness requirement: +downstream nodes must never receive tags with internal bookkeeping columns. + +The meta column name `_tag_node_input_ref` is a private constant defined in `function_node.py`. +It must not appear in any public schema or protocol. + +### CACHE_ONLY and UNAVAILABLE modes + +These are handled before the 3-stage pipeline, unchanged from the current implementation: + +```python +status = self.load_status +if status == LoadStatus.CACHE_ONLY: + await self._async_execute_cache_only(output, observer=observer) + return +if status == LoadStatus.UNAVAILABLE: + await output.close() + raise RuntimeError(...) +``` + +### Observer responsibilities by stage + +`on_node_start` is called once before the TaskGroup is entered. `on_node_end` is called once +after the TaskGroup exits. Per-item hooks are split across stages: + +| Stage | Observer calls | +|---|---| +| **before TaskGroup** | `on_node_start` (once) | +| **route_inputs** | `on_data_start` + `on_data_end(cached=True)` for cache hits | +| **FunctionPod.async_execute** | `on_data_start` + `on_data_end(cached=False)` or `on_data_crash` for misses | +| **record_and_forward** | none | +| **after TaskGroup** | `on_node_end` (once) | + +### Dead import cleanup in function_node.py + +After the refactor, `FunctionJobNode.async_execute()` no longer calls `resolve_concurrency` or +creates semaphores directly — that logic moves into `FunctionPod.async_execute()`. If `PipelineConfig` +is imported in `function_node.py` solely for the `resolve_concurrency(pod_config, PipelineConfig())` +call in `async_execute`, it becomes dead and should be removed. Verify during implementation. + +### Test changes + +- **Remove** `TestAsyncExecuteBackpressure` from `test_regression_fixes.py` (tests the method + in isolation, not through any real execution path) +- **Add** `tests/test_core/nodes/test_function_job_node_async_execute.py`: + - All items processed via `FunctionJobNode.async_execute()` (no-DB and DB paths) + - Concurrency limit respected: `max_concurrency=1` limits concurrent tasks to 1 + - Cache hits emitted directly, cache misses computed + - Correlation key is absent from output tags + - Semaphore released on computation failure + - Observer callbacks fire correctly for hits and misses + +--- + +## Key invariants + +1. Output tags carry no `_tag_node_input_ref` column — verified by test +2. `add_pipeline_record` is called exactly once per cache miss +3. Cache hits bypass computation entirely +4. `_FunctionPodBase.async_execute()` is unaware of pipeline DB, base_entry_id, or correlation keys +5. The no-DB path delegates to `_function_pod`, not `_cached_function_pod` or `_ephemeral_cached_pod` +6. The ephemeral path uses `_ephemeral_cached_pod` as task 2 and passes `is_ephemeral=True` to `add_pipeline_record` +7. The persistent path uses `_cached_function_pod` as task 2 and passes `is_ephemeral=False` to `add_pipeline_record` diff --git a/tests/test_channels/test_node_async_execute.py b/tests/test_channels/test_node_async_execute.py index 60e1fb497..0b46d6c82 100644 --- a/tests/test_channels/test_node_async_execute.py +++ b/tests/test_channels/test_node_async_execute.py @@ -695,21 +695,22 @@ def patched(tag, data, *, logger=None): assert len(call_log) == 3 @pytest.mark.asyncio - async def test_function_node_async_uses_async_process_data_internal(self): - """Verify FunctionJobNode.async_execute routes through _async_process_data_internal.""" + async def test_function_node_async_delegates_to_pod(self): + """Verify FunctionJobNode.async_execute (no-DB path) delegates to pod.async_execute.""" call_log = [] _, pod = make_double_pod() stream = make_stream(3) node = FunctionJobNode(pod, stream) - original = node._async_process_data_internal + # No-DB path delegates to pod.async_execute(); verify pod.async_process_data is called. + original = pod.async_process_data async def patched(tag, data, **kwargs): - call_log.append("_async_process_data_internal") + call_log.append("async_process_data") return await original(tag, data, **kwargs) - node._async_process_data_internal = patched + pod.async_process_data = patched # type: ignore[method-assign] input_ch = Channel(buffer_size=16) output_ch = Channel(buffer_size=16) diff --git a/tests/test_core/function_pod/test_function_pod_async_execute.py b/tests/test_core/function_pod/test_function_pod_async_execute.py new file mode 100644 index 000000000..74a1a9fcd --- /dev/null +++ b/tests/test_core/function_pod/test_function_pod_async_execute.py @@ -0,0 +1,162 @@ +"""Tests for _FunctionPodBase.async_execute() — observer hooks and backpressure.""" + +from __future__ import annotations + +import asyncio + +import pyarrow as pa +import pytest + +from orcapod.channels import Channel +from orcapod.core.data_function import PythonDataFunction +from orcapod.core.function_pod import FunctionPod +from orcapod.core.streams.arrow_table_stream import ArrowTableStream +from orcapod.types import PodConfig + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_stream(n: int = 3) -> ArrowTableStream: + schema = pa.schema( + [ + pa.field("id", pa.int64(), nullable=False), + pa.field("x", pa.int64(), nullable=False), + ] + ) + table = pa.table( + { + "id": pa.array(list(range(n)), type=pa.int64()), + "x": pa.array(list(range(n)), type=pa.int64()), + }, + schema=schema, + ) + return ArrowTableStream(table, tag_columns=["id"]) + + +async def _feed(stream: ArrowTableStream, ch: Channel) -> None: + for tag, data in stream.iter_data(): + await ch.writer.send((tag, data)) + await ch.writer.close() + + +class _SpyObserver: + def __init__(self): + self.starts = [] + self.ends = [] # [(label, tag, input, output, cached)] + self.crashes = [] + + def contextualize(self, *path): + return self + + def on_run_start(self, *a, **kw): pass + def on_run_end(self, *a, **kw): pass + def on_node_start(self, *a, **kw): pass + def on_node_end(self, *a, **kw): pass + + def on_data_start(self, label, tag, data): + self.starts.append((label, tag, data)) + + def on_data_end(self, label, tag, inp, out, *, cached): + self.ends.append((label, tag, inp, out, cached)) + + def on_data_crash(self, label, tag, data, exc): + self.crashes.append(exc) + + def create_data_logger(self, tag, data): + return None + + +# --------------------------------------------------------------------------- +# Observer hooks +# --------------------------------------------------------------------------- + + +class TestFunctionPodBaseAsyncExecuteObserver: + @pytest.mark.asyncio + async def test_on_data_start_and_end_fire_per_item(self): + """on_data_start fires before and on_data_end(cached=False) fires after each item.""" + + def double(x: int) -> int: + return x * 2 + + pf = PythonDataFunction(double, output_keys="result") + pod = FunctionPod(pf) + stream = _make_stream(3) + input_ch: Channel = Channel(buffer_size=16) + output_ch: Channel = Channel(buffer_size=16) + await _feed(stream, input_ch) + + spy = _SpyObserver() + await pod.async_execute([input_ch.reader], output_ch.writer, observer=spy) + results = await output_ch.reader.collect() + + assert len(results) == 3 + assert len(spy.starts) == 3 + assert len(spy.ends) == 3 + assert all(cached is False for _, _, _, _, cached in spy.ends) + + @pytest.mark.asyncio + async def test_on_data_crash_fires_not_on_data_end(self): + """on_data_crash fires (and on_data_end does not) when processing raises.""" + + def boom(x: int) -> int: + raise ValueError("bang") + + pf = PythonDataFunction(boom, output_keys="result") + pod = FunctionPod(pf) + stream = _make_stream(2) + input_ch: Channel = Channel(buffer_size=16) + output_ch: Channel = Channel(buffer_size=16) + await _feed(stream, input_ch) + + spy = _SpyObserver() + await pod.async_execute([input_ch.reader], output_ch.writer, observer=spy) + results = await output_ch.reader.collect() + + assert len(results) == 0 + assert len(spy.crashes) == 2 + assert len(spy.ends) == 0 + + +# --------------------------------------------------------------------------- +# Backpressure +# --------------------------------------------------------------------------- + + +class TestFunctionPodBaseAsyncExecuteBackpressure: + @pytest.mark.asyncio + async def test_max_concurrency_limits_concurrent_tasks(self): + """With max_concurrency=1, at most one task runs concurrently.""" + concurrent_count = 0 + max_observed = 0 + + def double(x: int) -> int: + return x * 2 + + pf = PythonDataFunction(double, output_keys="result") + original_async_call = pf.async_call + + async def patched(data, **kw): + nonlocal concurrent_count, max_observed + concurrent_count += 1 + max_observed = max(max_observed, concurrent_count) + try: + await asyncio.sleep(0.01) + return await original_async_call(data, **kw) + finally: + concurrent_count -= 1 + + pf.async_call = patched # type: ignore[method-assign] + + pod = FunctionPod(pf, pod_config=PodConfig(max_concurrency=1)) + stream = _make_stream(5) + input_ch: Channel = Channel(buffer_size=32) + output_ch: Channel = Channel(buffer_size=32) + await _feed(stream, input_ch) + await pod.async_execute([input_ch.reader], output_ch.writer) + results = await output_ch.reader.collect() + assert len(results) == 5 + assert max_observed <= 1 diff --git a/tests/test_core/nodes/test_function_job_node_async_execute.py b/tests/test_core/nodes/test_function_job_node_async_execute.py new file mode 100644 index 000000000..ef2508b13 --- /dev/null +++ b/tests/test_core/nodes/test_function_job_node_async_execute.py @@ -0,0 +1,470 @@ +"""Tests for FunctionJobNode.async_execute() — unified async execution path.""" + +from __future__ import annotations + +import asyncio + +import pyarrow as pa +import pytest + +from orcapod.channels import Channel +from orcapod.core.data_function import PythonDataFunction +from orcapod.core.function_pod import FunctionPod +from orcapod.core.nodes.function_node import FunctionJobNode +from orcapod.core.sources import ArrowTableSource +from orcapod.databases import InMemoryArrowDatabase +from orcapod.types import PodConfig + + +# --------------------------------------------------------------------------- +# Shared helpers +# --------------------------------------------------------------------------- + + +def _make_source(n: int = 3) -> ArrowTableSource: + """ArrowTableSource with tag=key (large_string), data=value (int64).""" + table = pa.table( + { + "key": pa.array([f"k{i}" for i in range(n)], type=pa.large_string()), + "value": pa.array(list(range(n)), type=pa.int64()), + } + ) + return ArrowTableSource(table, tag_columns=["key"], infer_nullable=True) + + +def double_no_db(value: int) -> int: + return value * 2 + + +def _make_no_db_node( + n: int = 3, max_concurrency: int | None = None +) -> FunctionJobNode: + """FunctionJobNode with no DB attached.""" + src = _make_source(n) + pod = FunctionPod( + PythonDataFunction(double_no_db, output_keys="result"), + pod_config=PodConfig(max_concurrency=max_concurrency), + ) + return FunctionJobNode(pod, src) + + +async def _run_node(node: FunctionJobNode) -> list[tuple]: + """Feed the node's own input stream through async_execute and collect results.""" + input_ch: Channel = Channel(buffer_size=32) + output_ch: Channel = Channel(buffer_size=32) + + async def feed() -> None: + for tag, data in node._input_stream.iter_data(): + await input_ch.writer.send((tag, data)) + await input_ch.writer.close() + + await asyncio.gather( + feed(), + node.async_execute(input_ch.reader, output_ch.writer), + ) + return await output_ch.reader.collect() + + +# --------------------------------------------------------------------------- +# No-DB path +# --------------------------------------------------------------------------- + + +class TestFunctionJobNodeAsyncExecuteNoDB: + @pytest.mark.asyncio + async def test_all_items_processed_and_forwarded(self): + """All input items reach the output channel.""" + node = _make_no_db_node(n=3) + results = await _run_node(node) + assert len(results) == 3 + values = sorted(data.as_dict()["result"] for _, data in results) + assert values == [0, 2, 4] + + @pytest.mark.asyncio + async def test_output_channel_closed_after_execution(self): + """Output channel is closed even if execution completes without error.""" + node = _make_no_db_node(n=2) + input_ch: Channel = Channel(buffer_size=8) + output_ch: Channel = Channel(buffer_size=8) + + async def feed() -> None: + for tag, data in node._input_stream.iter_data(): + await input_ch.writer.send((tag, data)) + await input_ch.writer.close() + + await asyncio.gather( + feed(), node.async_execute(input_ch.reader, output_ch.writer) + ) + # collect() should return immediately (channel closed), not hang. + results = await output_ch.reader.collect() + assert isinstance(results, list) + + @pytest.mark.asyncio + async def test_max_concurrency_limits_tasks(self): + """With max_concurrency=1, at most one task runs at a time.""" + concurrent_count = 0 + max_observed = 0 + + def double_base(value: int) -> int: + return value * 2 + + pf = PythonDataFunction(double_base, output_keys="result") + original_async_call = pf.async_call + + async def patched_async_call(data, **kwargs): + nonlocal concurrent_count, max_observed + concurrent_count += 1 + max_observed = max(max_observed, concurrent_count) + try: + await asyncio.sleep(0.01) + return await original_async_call(data, **kwargs) + finally: + concurrent_count -= 1 + + pf.async_call = patched_async_call # type: ignore[method-assign] + + src = _make_source(5) + pod = FunctionPod(pf, pod_config=PodConfig(max_concurrency=1)) + node = FunctionJobNode(pod, src) + results = await _run_node(node) + assert len(results) == 5 + assert max_observed <= 1 + + +# --------------------------------------------------------------------------- +# Helpers — observer spy +# --------------------------------------------------------------------------- + + +class _SpyObserver: + """Captures all observer hook invocations.""" + + def __init__(self): + self.node_starts: list[str] = [] + self.node_ends: list[str] = [] + self.data_starts: list[tuple] = [] + self.data_ends: list[tuple] = [] # (label, tag, input_data, output_data, cached) + self.data_crashes: list[tuple] = [] + + def contextualize(self, *path): + return self + + def on_node_start(self, node_label, node_hash, *, tag_schema=None): + self.node_starts.append(node_label) + + def on_node_end(self, node_label, node_hash): + self.node_ends.append(node_label) + + def on_data_start(self, node_label, tag, data): + self.data_starts.append((node_label, tag, data)) + + def on_data_end(self, node_label, tag, input_data, output_data, *, cached): + self.data_ends.append((node_label, tag, input_data, output_data, cached)) + + def on_data_crash(self, node_label, tag, data, error): + self.data_crashes.append((node_label, tag, data, error)) + + def create_data_logger(self, tag, data): + return None + + +# --------------------------------------------------------------------------- +# Helpers — DB node fixtures +# --------------------------------------------------------------------------- + + +def _make_db_node( + n: int = 3, +) -> tuple[FunctionJobNode, InMemoryArrowDatabase, InMemoryArrowDatabase]: + """FunctionJobNode with persistent pipeline + result databases.""" + + def double(value: int) -> int: + return value * 2 + + src = _make_source(n) + pod = FunctionPod(PythonDataFunction(double, output_keys="result")) + pipeline_db = InMemoryArrowDatabase() + result_db = InMemoryArrowDatabase() + node = FunctionJobNode( + pod, src, pipeline_database=pipeline_db, result_database=result_db + ) + return node, pipeline_db, result_db + + +# --------------------------------------------------------------------------- +# DB path — basic flow +# --------------------------------------------------------------------------- + + +class TestFunctionJobNodeAsyncExecuteDB: + @pytest.mark.asyncio + async def test_all_items_processed_and_forwarded(self): + """All input items are computed and forwarded via the DB path.""" + node, _, _ = _make_db_node(3) + results = await _run_node(node) + assert len(results) == 3 + values = sorted(data.as_dict()["result"] for _, data in results) + assert values == [0, 2, 4] + + @pytest.mark.asyncio + async def test_pipeline_records_written_for_each_item(self): + """add_pipeline_record is called once per cache miss → one row per item.""" + node, pipeline_db, _ = _make_db_node(3) + await _run_node(node) + records = pipeline_db.get_all_records(node.node_identity_path) + assert records is not None + assert records.num_rows == 3 + + @pytest.mark.asyncio + async def test_correlation_key_absent_from_output_tags(self): + """Output tags must not contain the internal _tag_node_input_ref column.""" + from orcapod.core.nodes.function_node import _TAG_NODE_INPUT_REF + from orcapod.system_constants import constants as _constants + + node, _, _ = _make_db_node(3) + results = await _run_node(node) + full_key = f"{_constants.META_PREFIX}{_TAG_NODE_INPUT_REF}" + for tag, _ in results: + meta = tag.get_meta_info() + assert full_key not in meta, ( + f"Output tag leaked internal correlation key {full_key!r}" + ) + + @pytest.mark.asyncio + async def test_cache_hits_emitted_without_recomputation(self): + """Items already in the pipeline DB are served from cache (no recomputation).""" + node, _, _ = _make_db_node(3) + + # First run — populates both pipeline DB and result DB. + first_results = await _run_node(node) + assert len(first_results) == 3 + + # Reset in-memory state so async_execute re-runs from scratch. + node.clear_cache() + + call_count = 0 + # async_execute() calls data_function.async_call() (not .call()), so patch + # async_call to reliably detect any recomputation via the async path. + original_async_call = node._function_pod.data_function.async_call + + async def counted_async_call(data, **kw): + nonlocal call_count + call_count += 1 + return await original_async_call(data, **kw) + + node._function_pod.data_function.async_call = counted_async_call # type: ignore[method-assign] + + second_results = await _run_node(node) + assert len(second_results) == 3 + # All 3 were hits — data function should NOT have been called again. + assert call_count == 0 + + @pytest.mark.asyncio + async def test_observer_cached_false_for_misses_cached_true_for_hits(self): + """Cache hits emit on_data_end(cached=True); misses emit on_data_end(cached=False).""" + node, _, _ = _make_db_node(2) + + # First run — both are misses. Just to populate the DB. + await _run_node(node) + node.clear_cache() + + spy = _SpyObserver() + input_ch: Channel = Channel(buffer_size=8) + output_ch: Channel = Channel(buffer_size=8) + + async def feed() -> None: + for tag, data in node._input_stream.iter_data(): + await input_ch.writer.send((tag, data)) + await input_ch.writer.close() + + await asyncio.gather( + feed(), + node.async_execute(input_ch.reader, output_ch.writer, observer=spy), + ) + await output_ch.reader.collect() + + # Second run was all hits. + assert len(spy.data_ends) == 2 + assert all(cached is True for _, _, _, _, cached in spy.data_ends) + + @pytest.mark.asyncio + async def test_correlation_key_absent_from_observer_callbacks(self): + """Observer on_data_start/on_data_end must not see _tag_node_input_ref (first run = cache misses).""" + from orcapod.core.nodes.function_node import _TAG_NODE_INPUT_REF + from orcapod.system_constants import constants as _constants + + node, _, _ = _make_db_node(3) + spy = _SpyObserver() + + # First run — all are cache misses, pod path is used with stamped tags. + input_ch: Channel = Channel(buffer_size=8) + output_ch: Channel = Channel(buffer_size=8) + + async def feed() -> None: + for tag, data in node._input_stream.iter_data(): + await input_ch.writer.send((tag, data)) + await input_ch.writer.close() + + await asyncio.gather( + feed(), + node.async_execute(input_ch.reader, output_ch.writer, observer=spy), + ) + await output_ch.reader.collect() + + full_key = f"{_constants.META_PREFIX}{_TAG_NODE_INPUT_REF}" + for _, tag, _ in spy.data_starts: + assert full_key not in tag.get_meta_info(), ( + f"on_data_start leaked correlation key {full_key!r}" + ) + for _, tag, _, _, _ in spy.data_ends: + assert full_key not in tag.get_meta_info(), ( + f"on_data_end leaked correlation key {full_key!r}" + ) + + @pytest.mark.asyncio + async def test_semaphore_not_leaked_on_computation_crash(self): + """When a data function raises, processing continues for remaining items.""" + call_count = 0 + + def sometimes_fail(value: int) -> int: + nonlocal call_count + call_count += 1 + if value == 1: + raise ValueError("deliberate failure") + return value * 2 + + src = _make_source(3) + pod = FunctionPod( + PythonDataFunction(sometimes_fail, output_keys="result"), + pod_config=PodConfig(max_concurrency=1), + ) + pipeline_db = InMemoryArrowDatabase() + result_db = InMemoryArrowDatabase() + node = FunctionJobNode( + pod, src, pipeline_database=pipeline_db, result_database=result_db + ) + + results = await _run_node(node) + # value=0 → 0, value=1 → crash (skipped), value=2 → 4 + assert len(results) == 2 + values = sorted(data.as_dict()["result"] for _, data in results) + assert values == [0, 4] + + @pytest.mark.asyncio + async def test_filtered_items_update_in_memory_cache(self): + """Items filtered by the pod (output_data=None) must be stored in + _cached_output_datas so the sync iter_data() path sees them. + + Regression: before the fix, _NodeLabelObserver.on_data_end with + out=None never popped the input_store entry or updated + _cached_output_datas, so filtered inputs were always recomputed and + the input_store accumulated entries for the lifetime of the call. + """ + node, _, _ = _make_db_node(3) + + # Deactivate the data function so every input is filtered (call → None). + node._function_pod.data_function.set_active(False) + + results = await _run_node(node) + # All items filtered — nothing emitted. + assert len(results) == 0 + + # _cached_output_datas must have one entry per input, with None data, + # confirming that on_data_end cleaned up input_store and updated the cache. + assert len(node._cached_output_datas) == 3 + for _tag, data in node._cached_output_datas.values(): + assert data is None + + @pytest.mark.asyncio + async def test_crashed_items_cleaned_from_input_store(self): + """Crashed items must not leave dangling entries in input_store. + + Regression: before the fix, _NodeLabelObserver.on_data_crash never + popped the input_store entry, so the dict accumulated one entry per + crash for the entire lifetime of the async_execute() call. + + Verification is indirect: if input_store cleanup works, no KeyError or + deadlock occurs when the same inputs are re-run after a crash run. + """ + + def sometimes_fail(value: int) -> int: + if value == 1: + raise ValueError("deliberate failure") + return value * 2 + + src = _make_source(3) + pod = FunctionPod(PythonDataFunction(sometimes_fail, output_keys="result")) + pipeline_db = InMemoryArrowDatabase() + result_db = InMemoryArrowDatabase() + node = FunctionJobNode( + pod, src, pipeline_database=pipeline_db, result_database=result_db + ) + + # First run — value=1 crashes. + first = await _run_node(node) + assert len(first) == 2 + + # Second run — value=1 still crashes, others are cache hits. + node.clear_cache() + second = await _run_node(node) + assert len(second) == 2 + values = sorted(data.as_dict()["result"] for _, data in second) + assert values == [0, 4] + + +# --------------------------------------------------------------------------- +# DB path — ephemeral +# --------------------------------------------------------------------------- + + +class TestFunctionJobNodeAsyncExecuteEphemeral: + @pytest.mark.asyncio + async def test_ephemeral_pipeline_records_have_is_ephemeral_true(self): + """Ephemeral execution writes pipeline records with is_ephemeral=True.""" + from orcapod.system_constants import constants + from orcapod.types import NodeConfig + + def double(value: int) -> int: + return value * 2 + + src = _make_source(2) + pod = FunctionPod(PythonDataFunction(double, output_keys="result")) + pipeline_db = InMemoryArrowDatabase() + result_db = InMemoryArrowDatabase() + ephemeral_db = InMemoryArrowDatabase() + node = FunctionJobNode( + pod, src, pipeline_database=pipeline_db, result_database=result_db + ) + node.node_config = NodeConfig(is_result_ephemeral=True) + node.set_ephemeral_store(ephemeral_db) + + results = await _run_node(node) + assert len(results) == 2 + + records = pipeline_db.get_all_records(node.node_identity_path) + assert records is not None + assert records.num_rows == 2 + # All rows should have is_ephemeral=True + is_eph_col = records.column(constants.IS_EPHEMERAL_COL).to_pylist() + assert all(v is True for v in is_eph_col) + + @pytest.mark.asyncio + async def test_ephemeral_missing_store_raises_runtime_error(self): + """Raises RuntimeError when is_result_ephemeral=True but no store set.""" + from orcapod.types import NodeConfig + + def double(value: int) -> int: + return value * 2 + + src = _make_source(1) + pod = FunctionPod(PythonDataFunction(double, output_keys="result")) + pipeline_db = InMemoryArrowDatabase() + result_db = InMemoryArrowDatabase() + node = FunctionJobNode( + pod, src, pipeline_database=pipeline_db, result_database=result_db + ) + node.node_config = NodeConfig(is_result_ephemeral=True) + # Deliberately do NOT call set_ephemeral_store() + + with pytest.raises(RuntimeError, match="ephemeral"): + await _run_node(node) diff --git a/tests/test_core/test_regression_fixes.py b/tests/test_core/test_regression_fixes.py index ad018581a..5bba54c32 100644 --- a/tests/test_core/test_regression_fixes.py +++ b/tests/test_core/test_regression_fixes.py @@ -5,10 +5,9 @@ 1. async_execute output channel closed on exception (try/finally) 2. DataFunctionWrapper.direct_call/direct_async_call bypass executor routing 3. Concurrent iteration falls back to sequential inside a running event loop -4. FunctionPod.async_execute backpressure bounds pending tasks -5. _materialize_to_stream preserves source_info provenance tokens -6. RayExecutor._ensure_ray_initialized uses ray_address -7. DataFunctionExecutorProtocol uses DataFunctionProtocol (not Any) +4. _materialize_to_stream preserves source_info provenance tokens +5. RayExecutor._ensure_ray_initialized uses ray_address +6. DataFunctionExecutorProtocol uses DataFunctionProtocol (not Any) """ from __future__ import annotations @@ -286,65 +285,7 @@ def test_uses_asyncio_run_when_no_loop(self): # =========================================================================== -# 4. FunctionPod.async_execute backpressure bounds pending tasks -# =========================================================================== - - -class TestAsyncExecuteBackpressure: - """With max_concurrency set, pending tasks should be bounded.""" - - @pytest.mark.asyncio - async def test_semaphore_limits_concurrent_tasks(self): - """With max_concurrency=1, at most one task should be running.""" - concurrent_count = 0 - max_concurrent = 0 - - async def track_concurrency(x: int) -> int: - nonlocal concurrent_count, max_concurrent - concurrent_count += 1 - max_concurrent = max(max_concurrent, concurrent_count) - await asyncio.sleep(0.01) # simulate work - concurrent_count -= 1 - return x * 2 - - def double(x: int) -> int: - return x * 2 - - # Build a PythonDataFunction that uses our async-aware tracker. - # We override async_call to directly call our async function. - pf = PythonDataFunction(double, output_keys="result") - - # Patch async_call to use our concurrency-tracking function - original_async_call = pf.async_call - - async def tracked_async_call(data: DataProtocol, **kwargs): - nonlocal concurrent_count, max_concurrent - concurrent_count += 1 - max_concurrent = max(max_concurrent, concurrent_count) - await asyncio.sleep(0.01) - concurrent_count -= 1 - return await original_async_call(data, **kwargs) - - pf.async_call = tracked_async_call # type: ignore - - pod = FunctionPod(pf, pod_config=PodConfig(max_concurrency=1)) - - stream = make_stream(5) - input_ch = Channel(buffer_size=32) - output_ch = Channel(buffer_size=32) - - await feed_stream_to_channel(stream, input_ch) - await pod.async_execute([input_ch.reader], output_ch.writer) - - results = await output_ch.reader.collect() - assert len(results) == 5 - # With semaphore acquired before task creation and max_concurrency=1, - # at most 1 should be in-flight at a time. - assert max_concurrent <= 1 - - -# =========================================================================== -# 5. _materialize_to_stream preserves source_info provenance +# 4. _materialize_to_stream preserves source_info provenance # =========================================================================== @@ -404,7 +345,7 @@ def test_materialize_source_columns_in_table(self): # =========================================================================== -# 6. RayExecutor._ensure_ray_initialized uses ray_address +# 5. RayExecutor._ensure_ray_initialized uses ray_address # =========================================================================== @@ -706,7 +647,7 @@ def decorator(wrapper): # =========================================================================== -# 7. DataFunctionExecutorProtocol type safety +# 6. DataFunctionExecutorProtocol type safety # ===========================================================================