diff --git a/src/radical/asyncflow/workflow_manager.py b/src/radical/asyncflow/workflow_manager.py index 5c4acaf..4721c93 100644 --- a/src/radical/asyncflow/workflow_manager.py +++ b/src/radical/asyncflow/workflow_manager.py @@ -752,9 +752,9 @@ def _register_component( "name", "" ) - # Detect dependencies + # Detect dependencies (futures may be passed positionally or as kwargs) comp_deps, input_files_deps, output_files_deps = self._detect_dependencies( - comp_desc["args"] + list(comp_desc["args"]) + list(comp_desc["kwargs"].values()) ) comp_desc["metadata"] = { @@ -785,6 +785,10 @@ def _register_component( self._update_dependency_tracking(comp_desc["uid"]) self._component_change_event.set() + # Wake the run loop when this component completes, so dependents are + # noticed immediately instead of on the next poll interval or timeout. + comp_fut.add_done_callback(lambda _: self._component_change_event.set()) + # Track block membership: if this component is registered from within a block's # execution context, record it so it gets cancelled when the block is cancelled. # Read ContextVar directly — not comp_desc["workflow_id"] which may be overridden @@ -919,6 +923,7 @@ def _detect_dependencies(self, possible_dependencies): dependencies = [] input_files = [] output_files = [] + seen_dep_uids = set() for possible_dep in possible_dependencies: # Flow component dependency @@ -927,6 +932,19 @@ def _detect_dependencies(self, possible_dependencies): possible_dep = possible_dep.task elif hasattr(possible_dep, BLOCK): possible_dep = possible_dep.block + else: + # not created by an asyncflow task/block — dependency + # tracking cannot resolve it + raise TypeError( + "Only futures returned by asyncflow tasks or blocks " + "can be used as dependencies, got a plain " + f"{type(possible_dep).__name__}" + ) + # Deduplicate: the same future passed more than once must only + # count once, or the dependency count can never reach zero. + if possible_dep["uid"] in seen_dep_uids: + continue + seen_dep_uids.add(possible_dep["uid"]) dependencies.append(possible_dep) # Input file dependency elif isinstance(possible_dep, InputFile): @@ -1093,7 +1111,8 @@ async def run(self): Note: - Runs indefinitely until cancelled or shutdown is signaled - - Uses sleep intervals to prevent busy-waiting + - Event-driven: waits on component change events (registration + and completion) instead of a fixed poll interval - Handles both implicit and explicit data dependencies - Trigger internal shutdown on loop failure """ @@ -1309,7 +1328,12 @@ async def run(self): task.cancel() raise else: - await asyncio.sleep(0.01) + # There was activity this pass; yield to let backend + # coroutines and callbacks run, then re-check immediately. + # Completion wake-ups arrive via _component_change_event + # (set by each component future's done callback), so no + # fixed poll interval is needed. + await asyncio.sleep(0) except asyncio.CancelledError: logger.debug("Run component stopped") diff --git a/tests/unit/test_dependencies_detection.py b/tests/unit/test_dependencies_detection.py index 7d346a7..d31a6dc 100644 --- a/tests/unit/test_dependencies_detection.py +++ b/tests/unit/test_dependencies_detection.py @@ -1,6 +1,12 @@ +import asyncio + import pytest -from radical.asyncflow import NoopExecutionBackend, WorkflowEngine +from radical.asyncflow import ( + LocalExecutionBackend, + NoopExecutionBackend, + WorkflowEngine, +) from radical.asyncflow.data import InputFile, OutputFile @@ -37,3 +43,65 @@ async def task2(): assert len(task_deps) == 1 assert task1 in task_deps[0]["args"] + + +@pytest.mark.asyncio +async def test_kwarg_future_is_tracked_as_dependency(): + """A future passed as kwarg must delay submission until it is resolved. + + Regression test: kwarg futures were resolved at submission but never + registered as dependencies, so a consumer could be submitted while a + slow kwarg producer was still running (InvalidStateError). + """ + backend = await LocalExecutionBackend() + flow = await WorkflowEngine.create(backend=backend) + try: + + @flow.function_task + async def slow_producer(): + await asyncio.sleep(0.05) + return "value" + + @flow.function_task + async def consumer(kw=None): + return kw + + result = await asyncio.wait_for(consumer(kw=slow_producer()), timeout=10) + assert result == "value" + finally: + await flow.shutdown() + + +@pytest.mark.asyncio +async def test_duplicate_future_as_arg_and_kwarg(): + """The same future passed twice must count as one dependency. + + Without deduplication the dependency count is incremented twice but only decremented + once on completion, deadlocking the consumer. + """ + backend = await LocalExecutionBackend() + flow = await WorkflowEngine.create(backend=backend) + try: + + @flow.function_task + async def producer(): + return 7 + + @flow.function_task + async def consumer(pos, kw=None): + return pos + kw + + fut = producer() + result = await asyncio.wait_for(consumer(fut, kw=fut), timeout=10) + assert result == 14 + finally: + await flow.shutdown() + + +@pytest.mark.asyncio +async def test_plain_future_dependency_rejected(): + """A plain asyncio.Future (not produced by a task or block) fails with a clear + TypeError instead of an unhelpful error downstream.""" + engine = await WorkflowEngine.create(backend=NoopExecutionBackend()) + with pytest.raises(TypeError, match="asyncflow tasks or blocks"): + engine._detect_dependencies([asyncio.Future()]) diff --git a/tests/unit/test_scheduler_responsiveness.py b/tests/unit/test_scheduler_responsiveness.py new file mode 100644 index 0000000..c85be54 --- /dev/null +++ b/tests/unit/test_scheduler_responsiveness.py @@ -0,0 +1,67 @@ +"""Scheduler responsiveness: the run loop must react to component completion +via events, not a fixed poll interval. + +Regression tests for two latency defects of the former poll-based loop: + +- a hardcoded 10ms sleep after every active pass put a ~10ms floor on each + task round-trip, and +- a dependent task behind a dependency running longer than that poll + interval was only submitted when the 1s event-wait timed out (~1s stall + per dependency edge). +""" + +import asyncio +import time + +import pytest +import pytest_asyncio + +from radical.asyncflow import LocalExecutionBackend, WorkflowEngine + + +class TestSchedulerResponsiveness: + @pytest_asyncio.fixture + async def flow(self): + backend = await LocalExecutionBackend() + flow = await WorkflowEngine.create(backend=backend) + yield flow + await flow.shutdown() + + @pytest.mark.asyncio + async def test_dependent_starts_promptly_after_dependency(self, flow): + """A chain behind a 50ms task must not stall in the 1s event-wait.""" + + @flow.function_task + async def slow(): + await asyncio.sleep(0.05) + return 1 + + @flow.function_task + async def fast(dep): + return dep + 1 + + start = time.perf_counter() + result = await fast(slow()) + elapsed = time.perf_counter() - start + + assert result == 2 + # poll-based loop needed ~1.05s here; allow generous CI margin + assert elapsed < 0.5, f"dependent task stalled: {elapsed:.3f}s" + + @pytest.mark.asyncio + async def test_sequential_latency_below_poll_interval(self, flow): + """Per-task round-trip must beat the former 10ms poll floor.""" + + @flow.function_task + async def noop(): + return None + + n = 20 + await noop() # warmup + start = time.perf_counter() + for _ in range(n): + await noop() + avg = (time.perf_counter() - start) / n + + # poll-based loop could not go below 10ms per task + assert avg < 0.008, f"avg task round-trip too slow: {avg * 1000:.1f}ms"