From cc4dbba6642bdf4586237933cb7f277c23ac5889 Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Thu, 9 Jul 2026 19:24:15 +0000 Subject: [PATCH 01/15] docs(specs): add ITL-516 async_execute unification design spec Closes ITL-516 Co-Authored-By: Claude Sonnet 4.5 --- ...7-09-itl-516-unify-async-execute-design.md | 221 ++++++++++++++++++ 1 file changed, 221 insertions(+) create mode 100644 superpowers/specs/2026-07-09-itl-516-unify-async-execute-design.md 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 00000000..bb80643c --- /dev/null +++ b/superpowers/specs/2026-07-09-itl-516-unify-async-execute-design.md @@ -0,0 +1,221 @@ +# 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. + +### FunctionPod.async_execute() — add observer support + +Add an optional `observer` parameter (same pattern as all node `async_execute` signatures): + +```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: +``` + +When `observer` is provided, fire `on_data_start`, `on_data_end`, and `on_data_crash` per item. +When `None`, behaviour is unchanged (no-op). This makes `FunctionPod.async_execute()` a +first-class execution method suitable for delegation 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) + +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 + │ + ▼ _cached_function_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, ...) + → strip correlation key from output_tag + → emit (clean_tag, output_data) to output +``` + +**Backpressure:** `compute_channel` and `result_channel` are bounded channels. If FunctionPod is +slower than the router, `route_inputs` blocks on `compute_channel.writer.send()`; if the recorder +is slower than FunctionPod, FunctionPod 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 FunctionPod to finish. +FunctionPod 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. FunctionPod.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` From e7f9e9b11d527d8fbf25f1bc8b3f6985166b6d3b Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 02:38:47 +0000 Subject: [PATCH 02/15] docs(specs): update ITL-516 spec with ephemeral path and _FunctionPodBase promotion Co-Authored-By: Claude Sonnet 4.5 --- ...7-09-itl-516-unify-async-execute-design.md | 58 +++++++++++++++---- 1 file changed, 46 insertions(+), 12 deletions(-) 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 index bb80643c..58bc75fd 100644 --- 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 @@ -76,11 +76,24 @@ 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. -### FunctionPod.async_execute() — add observer support +### 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]]], @@ -92,8 +105,9 @@ async def async_execute( ``` When `observer` is provided, fire `on_data_start`, `on_data_end`, and `on_data_crash` per item. -When `None`, behaviour is unchanged (no-op). This makes `FunctionPod.async_execute()` a -first-class execution method suitable for delegation from `FunctionJobNode`. +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 @@ -111,6 +125,23 @@ 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`: ``` @@ -120,25 +151,26 @@ input_channel ├── cache hit → emit to output (observer: on_data_start + on_data_end cached=True) └── cache miss → stamp tag with correlation key → compute_channel │ - ▼ _cached_function_pod.async_execute() — task 2 + ▼ 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, ...) + → 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 FunctionPod is -slower than the router, `route_inputs` blocks on `compute_channel.writer.send()`; if the recorder -is slower than FunctionPod, FunctionPod blocks on `result_channel.writer.send()`. +**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 FunctionPod to finish. -FunctionPod closes `result_channel.writer` in its `finally`, signalling the recorder to finish. +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 @@ -217,5 +249,7 @@ call in `async_execute`, it becomes dead and should be removed. Verify during im 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. FunctionPod.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` +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` From 89ba83732453696ff50897f97c0c84dd1c2238ba Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 03:53:21 +0000 Subject: [PATCH 03/15] 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. Co-Authored-By: Claude Sonnet 4.6 --- src/orcapod/core/function_pod.py | 116 ++++++++++++++--------- tests/test_core/test_regression_fixes.py | 80 ++++++++++++++++ 2 files changed, 151 insertions(+), 45 deletions(-) diff --git a/src/orcapod/core/function_pod.py b/src/orcapod/core/function_pod.py index b835bb18..ba4fcd91 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,66 @@ 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) + 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() + @abstractmethod def process( self, *streams: StreamProtocol, label: str | None = None @@ -337,51 +403,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 +790,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 getattr(self._function_pod, "pod_config", PodConfig()) + @property def uri(self) -> tuple[str, ...]: return self._function_pod.uri diff --git a/tests/test_core/test_regression_fixes.py b/tests/test_core/test_regression_fixes.py index ad018581..c62c4b60 100644 --- a/tests/test_core/test_regression_fixes.py +++ b/tests/test_core/test_regression_fixes.py @@ -745,3 +745,83 @@ def test_protocol_async_execute_callable_annotation_is_typed(self): f"async_execute_callable() fn should reference Callable, " f"got {fn_annotation!r}" ) + + +# =========================================================================== +# 8. _FunctionPodBase.async_execute observer hooks +# =========================================================================== + + +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(cached=False) 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 + 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 + assert len(crashes) == 2 + assert len(ends) == 0 From f31d66824d0e2d0da1f72ae00c49dc704b45e54f Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 04:54:56 +0000 Subject: [PATCH 04/15] fix(function_pod): add pod_config to FunctionPodProtocol, clean up annotation quotes Adds pod_config property to FunctionPodProtocol so WrappedFunctionPod can use self._function_pod.pod_config directly without getattr fallback. Removes redundant explicit string quotes from async_execute signature (file has __future__ annotations). --- src/orcapod/core/function_pod.py | 6 +++--- src/orcapod/protocols/core_protocols/function_pod.py | 6 ++++++ 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/orcapod/core/function_pod.py b/src/orcapod/core/function_pod.py index ba4fcd91..67471a1b 100644 --- a/src/orcapod/core/function_pod.py +++ b/src/orcapod/core/function_pod.py @@ -201,11 +201,11 @@ def handle_input_streams(self, *streams: StreamProtocol) -> StreamProtocol: async def async_execute( self, - inputs: "Sequence[ReadableChannel[tuple[TagProtocol, DataProtocol]]]", - output: "WritableChannel[tuple[TagProtocol, DataProtocol]]", + inputs: Sequence[ReadableChannel[tuple[TagProtocol, DataProtocol]]], + output: WritableChannel[tuple[TagProtocol, DataProtocol]], pipeline_config: PipelineConfig | None = None, *, - observer: "ExecutionObserverProtocol | None" = None, + observer: ExecutionObserverProtocol | None = None, ) -> None: """Streaming async execution with per-data concurrency control. diff --git a/src/orcapod/protocols/core_protocols/function_pod.py b/src/orcapod/protocols/core_protocols/function_pod.py index 9c8bab91..3f0de3ff 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]: ... From f1d048c0a891c3b3a6e50817dababf411415cb1a Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 05:46:02 +0000 Subject: [PATCH 05/15] refactor(function_node): delegate no-DB async_execute to pod MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit FunctionJobNode.async_execute() no longer reimplements the semaphore +TaskGroup dispatch loop for the no-DB path. Delegates directly to self._function_pod.async_execute(), matching the OperatorNode pattern. Deletes _async_execute_one_data (dead after delegation); its logic is inlined into the DB path's _process_one_db() closure so the DB path retains concurrency-controlled dispatch until Task 3 redesigns it. Semaphore setup moved from the shared preamble into the DB branch only. PipelineConfig, PodConfig, and resolve_concurrency imports are retained because they are still used in the DB branch. Updates test_function_node_async_uses_async_process_data_internal → test_function_node_async_delegates_to_pod to reflect that the no-DB path now routes through pod.async_process_data, not node._async_process_data_internal. Adds tests/test_core/nodes/test_function_job_node_async_execute.py with TestFunctionJobNodeAsyncExecuteNoDB (3 tests). Co-Authored-By: Claude Sonnet 4.6 --- src/orcapod/core/nodes/function_node.py | 100 +++++--------- .../test_channels/test_node_async_execute.py | 11 +- .../test_function_job_node_async_execute.py | 129 ++++++++++++++++++ 3 files changed, 168 insertions(+), 72 deletions(-) create mode 100644 tests/test_core/nodes/test_function_job_node_async_execute.py diff --git a/src/orcapod/core/nodes/function_node.py b/src/orcapod/core/nodes/function_node.py index 6acdaee8..233c397c 100644 --- a/src/orcapod/core/nodes/function_node.py +++ b/src/orcapod/core/nodes/function_node.py @@ -2086,21 +2086,20 @@ 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: + # Resolve concurrency limit for DB path (pipeline config not + # threaded through the orchestrator, so 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 + ) + # Phase 1: build cache lookup from pipeline DB loaded = self._load_cached_entries() self._cached_output_datas.update(loaded) @@ -2122,12 +2121,24 @@ async def _process_one_db( ) 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_start(node_label, tag, data) + pkt_logger = ctx_obs.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, + ) + ctx_obs.on_data_crash(node_label, tag, data, exc) + else: + ctx_obs.on_data_end( + node_label, tag, data, result_data, cached=False + ) + if result_data is not None: + await output.send((tag_out, result_data)) async with asyncio.TaskGroup() as tg: async for tag, data in input_channel: @@ -2144,58 +2155,13 @@ async def _guarded_db( await sem.acquire() tg.create_task(_guarded_db()) 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/tests/test_channels/test_node_async_execute.py b/tests/test_channels/test_node_async_execute.py index 60e1fb49..0b46d6c8 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/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 00000000..e9f1c313 --- /dev/null +++ b/tests/test_core/nodes/test_function_job_node_async_execute.py @@ -0,0 +1,129 @@ +"""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) + 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 From c300ae10e56bddb24d72b12f3f6c612fb319d9c8 Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 06:29:52 +0000 Subject: [PATCH 06/15] fix(test_function_job_node): remove unused InMemoryArrowDatabase import Co-Authored-By: Claude Sonnet 4.6 --- tests/test_core/nodes/test_function_job_node_async_execute.py | 1 - 1 file changed, 1 deletion(-) 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 index e9f1c313..2351f17d 100644 --- a/tests/test_core/nodes/test_function_job_node_async_execute.py +++ b/tests/test_core/nodes/test_function_job_node_async_execute.py @@ -12,7 +12,6 @@ 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 From 42411325aa891cb56a8a828b74f36ed2fb08e154 Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 07:06:06 +0000 Subject: [PATCH 07/15] feat(function_node): redesign DB 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). Removes dead imports: PipelineConfig, PodConfig, resolve_concurrency. Also fixes _FunctionPodBase.async_execute() to call create_data_logger() and pass the returned logger to async_process_data(), restoring logging support for the DB execution path. Co-Authored-By: Claude Sonnet 4.6 --- src/orcapod/core/function_pod.py | 5 +- src/orcapod/core/nodes/function_node.py | 175 +++++++++---- .../test_function_job_node_async_execute.py | 243 ++++++++++++++++++ 3 files changed, 366 insertions(+), 57 deletions(-) diff --git a/src/orcapod/core/function_pod.py b/src/orcapod/core/function_pod.py index 67471a1b..1d03dcff 100644 --- a/src/orcapod/core/function_pod.py +++ b/src/orcapod/core/function_pod.py @@ -236,8 +236,11 @@ async def async_execute( 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) + 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 diff --git a/src/orcapod/core/nodes/function_node.py b/src/orcapod/core/nodes/function_node.py index 233c397c..da3d15f8 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, @@ -2090,17 +2093,23 @@ async def async_execute( ctx_obs.on_node_start(node_label, node_hash, tag_schema=tag_schema) if self._cached_function_pod is not None: - # Resolve concurrency limit for DB path (pipeline config not - # threaded through the orchestrator, so 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 - ) + # ---------------------------------------------------------- + # 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: build cache lookup from pipeline DB + # Load pipeline DB → in-memory cache keyed by base_entry_id. loaded = self._load_cached_entries() self._cached_output_datas.update(loaded) if loaded: @@ -2108,52 +2117,106 @@ 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, ) - await output.send((tag_out, result_data)) - else: + # 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)) + + # 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_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: ctx_obs.on_data_start(node_label, tag, data) - pkt_logger = ctx_obs.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, - ) - ctx_obs.on_data_crash(node_label, tag, data, exc) - else: - ctx_obs.on_data_end( - node_label, tag, data, result_data, cached=False - ) - if result_data is not None: - await output.send((tag_out, result_data)) + + def on_data_end( + self, lbl: str, tag: Any, inp: Any, out: Any, *, cached: bool + ) -> None: + ctx_obs.on_data_end(node_label, tag, inp, out, cached=cached) + + def on_data_crash( + self, lbl: str, tag: Any, data: Any, error: Any + ) -> None: + ctx_obs.on_data_crash(node_label, tag, data, error) + + 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: # No-DB path: delegate dispatch entirely to the pod. # The pod's async_execute() closes `output` in its own finally. 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 index 2351f17d..4785704b 100644 --- a/tests/test_core/nodes/test_function_job_node_async_execute.py +++ b/tests/test_core/nodes/test_function_job_node_async_execute.py @@ -12,6 +12,7 @@ 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 @@ -126,3 +127,245 @@ async def patched_async_call(data, **kwargs): 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 + original_call = node._function_pod.data_function.call + + def counted_call(data, **kw): + nonlocal call_count + call_count += 1 + return original_call(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) + + # 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_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 + + 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) From fe3a68ef6289706983d353f868872aa7afebd79a Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 14:01:49 +0000 Subject: [PATCH 08/15] fix(function_node): use ignore_missing=True in record_and_forward drop_meta_columns The invariant that _NodeLabelObserver only receives stamped tags holds today, but drop_meta_columns in record_and_forward should be defensive against future call-site changes that could violate that invariant. Passing ignore_missing=True matches the create_data_logger pattern in the same class and avoids a hard crash if the key is ever absent. Co-Authored-By: Claude Sonnet 4.6 --- src/orcapod/core/nodes/function_node.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/orcapod/core/nodes/function_node.py b/src/orcapod/core/nodes/function_node.py index da3d15f8..f202955b 100644 --- a/src/orcapod/core/nodes/function_node.py +++ b/src/orcapod/core/nodes/function_node.py @@ -2165,7 +2165,9 @@ async def record_and_forward() -> None: ) # 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) + 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 From aacf8169a889877e4dad572e42e4d3ab4032dc21 Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 14:52:59 +0000 Subject: [PATCH 09/15] docs(plans): add ITL-516 async_execute unification implementation plan Co-Authored-By: Claude Sonnet 4.6 --- .../2026-07-10-itl-516-unify-async-execute.md | 1357 +++++++++++++++++ 1 file changed, 1357 insertions(+) create mode 100644 superpowers/plans/2026-07-10-itl-516-unify-async-execute.md 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 00000000..f1e3cc04 --- /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` From c52445073ebb1fbaad8bd8059a21380dc9c5c650 Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 15:00:28 +0000 Subject: [PATCH 10/15] fix(function_node): strip correlation key from _NodeLabelObserver tag forwarding on_data_start/on_data_end/on_data_crash were forwarding the stamped tag (containing _tag_node_input_ref) to external observers for cache-miss items. Apply drop_meta_columns(ignore_missing=True) before forwarding so observers never see internal correlation key metadata. Adds missing protocol method delegation. Adds regression test for observer callback key absence. --- src/orcapod/core/nodes/function_node.py | 21 ++++++++++-- .../test_function_job_node_async_execute.py | 34 +++++++++++++++++++ 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/src/orcapod/core/nodes/function_node.py b/src/orcapod/core/nodes/function_node.py index f202955b..4ec58356 100644 --- a/src/orcapod/core/nodes/function_node.py +++ b/src/orcapod/core/nodes/function_node.py @@ -2182,6 +2182,12 @@ class _NodeLabelObserver: 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) @@ -2189,17 +2195,26 @@ 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: - ctx_obs.on_data_start(node_label, tag, data) + 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: - ctx_obs.on_data_end(node_label, tag, inp, out, cached=cached) + clean = tag.drop_meta_columns( + _TAG_NODE_INPUT_REF, ignore_missing=True + ) + ctx_obs.on_data_end(node_label, clean, inp, out, cached=cached) def on_data_crash( self, lbl: str, tag: Any, data: Any, error: Any ) -> None: - ctx_obs.on_data_crash(node_label, tag, data, error) + clean = tag.drop_meta_columns( + _TAG_NODE_INPUT_REF, ignore_missing=True + ) + ctx_obs.on_data_crash(node_label, clean, data, error) def create_data_logger(self, tag: Any, data: Any) -> Any: # Strip the correlation key before creating the logger 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 index 4785704b..7b5f4ca0 100644 --- a/tests/test_core/nodes/test_function_job_node_async_execute.py +++ b/tests/test_core/nodes/test_function_job_node_async_execute.py @@ -283,6 +283,40 @@ async def feed() -> None: 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.""" From 8868df67a021eff573b1fe51da7c324e6f2c0372 Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 15:09:13 +0000 Subject: [PATCH 11/15] test(async_execute): remove orphan backpressure test, add dedicated pod-level coverage Removes TestAsyncExecuteBackpressure from test_regression_fixes.py (tested an orphan FunctionPod method never called by the orchestrator). Equivalent coverage lives in TestFunctionJobNodeAsyncExecuteNoDB and TestFunctionPodBaseAsyncExecuteBackpressure. Moves TestFunctionPodBase- AsyncExecuteObserver to dedicated file. Creates: tests/test_core/function_pod/test_function_pod_async_execute.py: observer hooks and backpressure for _FunctionPodBase.async_execute() Co-Authored-By: Claude Sonnet 4.6 --- .../test_function_pod_async_execute.py | 160 ++++++++++++++++++ tests/test_core/test_regression_fixes.py | 151 +---------------- 2 files changed, 166 insertions(+), 145 deletions(-) create mode 100644 tests/test_core/function_pod/test_function_pod_async_execute.py 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 00000000..b89512fe --- /dev/null +++ b/tests/test_core/function_pod/test_function_pod_async_execute.py @@ -0,0 +1,160 @@ +"""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) + 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 diff --git a/tests/test_core/test_regression_fixes.py b/tests/test_core/test_regression_fixes.py index c62c4b60..5bba54c3 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 # =========================================================================== @@ -745,83 +686,3 @@ def test_protocol_async_execute_callable_annotation_is_typed(self): f"async_execute_callable() fn should reference Callable, " f"got {fn_annotation!r}" ) - - -# =========================================================================== -# 8. _FunctionPodBase.async_execute observer hooks -# =========================================================================== - - -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(cached=False) 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 - 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 - assert len(crashes) == 2 - assert len(ends) == 0 From 20a2c99d9e293ad645456f1c6703f2908fe61347 Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 22:47:42 +0000 Subject: [PATCH 12/15] fix(function_node): clean up input_store and cache filtered/crashed items in DB pipeline Two issues in _NodeLabelObserver (PR #224 Copilot review): 1. Filtered items (output_data=None): when the pod filters an item, it calls on_data_end(out=None) but does not emit to result_channel, so record_and_forward() never sees the item and input_store[correlation_key] was never removed. Also, _cached_output_datas was not updated, breaking the old _async_process_data_internal behaviour of always storing (tag, None) for filtered inputs so iter_data() returns the cached None on the next call. Fix: on_data_end now pops the input_store entry and updates _cached_output_datas when out is None. 2. Crashed items: on_data_crash() never removed the input_store entry, leaving stale entries for the lifetime of the async_execute() call. Fix: on_data_crash() now pops the correlation key from input_store. Adds fn_node = self alias so _NodeLabelObserver methods can reference FunctionJobNode state without shadowing the inner class's self. Adds two regression tests: - test_filtered_items_update_in_memory_cache - test_crashed_items_cleaned_from_input_store Co-Authored-By: Claude Sonnet 4.6 --- src/orcapod/core/nodes/function_node.py | 33 ++++++++++ .../test_function_job_node_async_execute.py | 61 +++++++++++++++++++ 2 files changed, 94 insertions(+) diff --git a/src/orcapod/core/nodes/function_node.py b/src/orcapod/core/nodes/function_node.py index 4ec58356..dd144c7d 100644 --- a/src/orcapod/core/nodes/function_node.py +++ b/src/orcapod/core/nodes/function_node.py @@ -2173,6 +2173,10 @@ async def record_and_forward() -> 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. @@ -2207,6 +2211,30 @@ def on_data_end( _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 this input is not recomputed on the next call + # (restoring the old _async_process_data_internal + # behaviour of always storing even None results). + 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 @@ -2215,6 +2243,11 @@ def on_data_crash( _TAG_NODE_INPUT_REF, ignore_missing=True ) 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 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 index 7b5f4ca0..134c42eb 100644 --- a/tests/test_core/nodes/test_function_job_node_async_execute.py +++ b/tests/test_core/nodes/test_function_job_node_async_execute.py @@ -346,6 +346,67 @@ def sometimes_fail(value: int) -> int: 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 From 29000d3411321ca375eb6d8490680e0cb2357f4b Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Fri, 10 Jul 2026 23:25:54 +0000 Subject: [PATCH 13/15] refactor(function_pod): use direct .pod_config access in WrappedFunctionPod The getattr fallback was added when FunctionPodProtocol didn't yet include pod_config. Now that the property is part of the protocol, self._function_pod is guaranteed to have it and the defensive getattr is unnecessary. Co-Authored-By: Claude Sonnet 4.6 --- src/orcapod/core/function_pod.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/orcapod/core/function_pod.py b/src/orcapod/core/function_pod.py index 1d03dcff..2be8c7bb 100644 --- a/src/orcapod/core/function_pod.py +++ b/src/orcapod/core/function_pod.py @@ -796,7 +796,7 @@ def computed_label(self) -> str | None: @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()) + return self._function_pod.pod_config @property def uri(self) -> tuple[str, ...]: From e82a062d235c47809b85e85c95d0616396dc1678 Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Sat, 11 Jul 2026 01:09:26 +0000 Subject: [PATCH 14/15] fix(function_node): correct comment and test for filtered-item async behavior Two follow-up fixes from second Copilot review: 1. Patch async_call instead of call in test_cache_hits_emitted_without_recomputation. async_execute() dispatches via data_function.async_call(), so patching .call() would miss any recomputation that went through the async path. 2. Correct the comment in _NodeLabelObserver.on_data_end for the out-is-None branch. The in-memory cache update only benefits the sync iter_data() path. The async route_inputs() checks cached_by_base_entry_id (pipeline DB), which never contains filtered items, so subsequent async_execute() calls still re-send them to the pod. Co-Authored-By: Claude Sonnet 4.6 --- src/orcapod/core/nodes/function_node.py | 11 ++++++++--- .../nodes/test_function_job_node_async_execute.py | 10 ++++++---- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/orcapod/core/nodes/function_node.py b/src/orcapod/core/nodes/function_node.py index dd144c7d..56b2307d 100644 --- a/src/orcapod/core/nodes/function_node.py +++ b/src/orcapod/core/nodes/function_node.py @@ -2216,9 +2216,14 @@ def on_data_end( # 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 this input is not recomputed on the next call - # (restoring the old _async_process_data_internal - # behaviour of always storing even None results). + # 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 ) 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 index 134c42eb..b3aaa7f1 100644 --- a/tests/test_core/nodes/test_function_job_node_async_execute.py +++ b/tests/test_core/nodes/test_function_job_node_async_execute.py @@ -241,14 +241,16 @@ async def test_cache_hits_emitted_without_recomputation(self): node.clear_cache() call_count = 0 - original_call = node._function_pod.data_function.call + # 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 - def counted_call(data, **kw): + async def counted_async_call(data, **kw): nonlocal call_count call_count += 1 - return original_call(data, **kw) + return await original_async_call(data, **kw) - node._function_pod.data_function.call = counted_call # type: ignore[method-assign] + 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 From 603443f0db2ab3095df1374019fed99f436d10da Mon Sep 17 00:00:00 2001 From: "agent-kurodo[bot]" <268466204+agent-kurodo[bot]@users.noreply.github.com> Date: Sat, 11 Jul 2026 01:16:47 +0000 Subject: [PATCH 15/15] test(async_execute): move concurrent_count decrement into finally block Both concurrency tests decremented concurrent_count before awaiting original_async_call, so any overlap that occurred during the underlying call was not counted. Moving the decrement to a finally block ensures the counter spans the full duration of the patched coroutine and the assertion is reliable even when the underlying call itself yields. Co-Authored-By: Claude Sonnet 4.6 --- .../function_pod/test_function_pod_async_execute.py | 8 +++++--- .../nodes/test_function_job_node_async_execute.py | 8 +++++--- 2 files changed, 10 insertions(+), 6 deletions(-) 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 index b89512fe..74a1a9fc 100644 --- a/tests/test_core/function_pod/test_function_pod_async_execute.py +++ b/tests/test_core/function_pod/test_function_pod_async_execute.py @@ -143,9 +143,11 @@ 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) + 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] 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 index b3aaa7f1..ef2508b1 100644 --- a/tests/test_core/nodes/test_function_job_node_async_execute.py +++ b/tests/test_core/nodes/test_function_job_node_async_execute.py @@ -115,9 +115,11 @@ 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) + 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]