Spaces:
Sleeping
Sleeping
| """Branch — parallel fork/join step.""" | |
| from __future__ import annotations | |
| import asyncio | |
| import dataclasses | |
| from concurrent.futures import ThreadPoolExecutor | |
| from enum import Enum | |
| from types import MappingProxyType | |
| from typing import Callable | |
| from .context import StepContext | |
| from .errors import BranchError | |
| class MergeStrategy(Enum): | |
| """Built-in merge strategies for Branch outputs.""" | |
| RAISE_ON_CONFLICT = "raise_on_conflict" | |
| LAST_WRITE_WINS = "last_write_wins" | |
| NAMESPACED = "namespaced" | |
| # --------------------------------------------------------------------------- | |
| # Built-in merge functions | |
| # --------------------------------------------------------------------------- | |
| def _merge_raise_on_conflict(ctxs: list[StepContext]) -> StepContext: | |
| """Raise if any two branches wrote different values for the same field. | |
| Metadata is always merged (union across all branches; last writer wins | |
| within metadata — there is no named-field semantic there). | |
| Uses ``type(ctxs[0])`` so subclass fields are included in the comparison. | |
| """ | |
| if len(ctxs) == 1: | |
| return ctxs[0] | |
| conflicts: set[str] = set() | |
| for f in dataclasses.fields(type(ctxs[0])): | |
| if f.name == "metadata": | |
| continue | |
| first_val = getattr(ctxs[0], f.name) | |
| if any(getattr(ctx, f.name) != first_val for ctx in ctxs[1:]): | |
| conflicts.add(f.name) | |
| if conflicts: | |
| raise ValueError( | |
| f"Branch outputs conflict on fields {conflicts!r}. " | |
| "Use a different merge strategy or ensure branches write disjoint fields." | |
| ) | |
| merged_meta: dict = {} | |
| for ctx in ctxs: | |
| merged_meta.update(ctx.metadata) | |
| return dataclasses.replace(ctxs[0], metadata=MappingProxyType(merged_meta)) | |
| def _merge_last_write_wins(ctxs: list[StepContext]) -> StepContext: | |
| """Last branch's value wins for every conflicting field. | |
| Uses ``type(ctxs[0])`` so subclass fields are included in the comparison. | |
| """ | |
| if len(ctxs) == 1: | |
| return ctxs[0] | |
| # Start from first context, overlay with each subsequent one | |
| result = ctxs[0] | |
| ctx_type = type(ctxs[0]) | |
| for ctx in ctxs[1:]: | |
| changes: dict = {} | |
| for f in dataclasses.fields(ctx_type): | |
| if f.name == "metadata": | |
| continue | |
| val = getattr(ctx, f.name) | |
| if val != getattr(result, f.name): | |
| changes[f.name] = val | |
| if changes: | |
| result = dataclasses.replace(result, **changes) | |
| merged_meta: dict = {} | |
| for ctx in ctxs: | |
| merged_meta.update(ctx.metadata) | |
| return dataclasses.replace(result, metadata=MappingProxyType(merged_meta)) | |
| def _merge_namespaced(ctxs: list[StepContext]) -> StepContext: | |
| """Each branch's output is stored at ``ctx.metadata["branch_N"]``. | |
| Named fields are taken from the first branch; no conflict is possible | |
| because branch outputs are kept in separate metadata keys. | |
| """ | |
| base = ctxs[0] | |
| extra: dict = {f"branch_{i}": ctx for i, ctx in enumerate(ctxs)} | |
| merged_meta = MappingProxyType({**base.metadata, **extra}) | |
| return dataclasses.replace(base, metadata=merged_meta) | |
| _BUILTIN_MERGES: dict[MergeStrategy, Callable] = { | |
| MergeStrategy.RAISE_ON_CONFLICT: _merge_raise_on_conflict, | |
| MergeStrategy.LAST_WRITE_WINS: _merge_last_write_wins, | |
| MergeStrategy.NAMESPACED: _merge_namespaced, | |
| } | |
| # --------------------------------------------------------------------------- | |
| # Branch | |
| # --------------------------------------------------------------------------- | |
| class Branch: | |
| """Runs multiple pipelines in parallel, then merges their outputs. | |
| ``Branch`` satisfies ``StepProtocol`` — it can be used wherever a step | |
| is expected. ``requires`` and ``provides`` are inferred from the union | |
| of the child pipelines' contracts. | |
| In sync contexts (called directly), fan-out is via | |
| ``ThreadPoolExecutor``. In async contexts (awaited), fan-out is via | |
| ``asyncio.gather``. | |
| All branches always run to completion before any failure is raised — | |
| ``BranchError`` carries the full list of failures. | |
| """ | |
| def __init__( | |
| self, | |
| *pipelines: object, | |
| merge: MergeStrategy | Callable = MergeStrategy.RAISE_ON_CONFLICT, | |
| ) -> None: | |
| if not pipelines: | |
| raise ValueError("Branch requires at least one child pipeline.") | |
| self.pipelines = list(pipelines) | |
| if callable(merge) and not isinstance(merge, MergeStrategy): | |
| self._merge_fn: Callable = merge | |
| else: | |
| self._merge_fn = _BUILTIN_MERGES[merge] # type: ignore[index] | |
| # Infer requires/provides from the union of child contracts | |
| all_requires: set[str] = set() | |
| all_provides: set[str] = set() | |
| for p in self.pipelines: | |
| all_requires |= set(getattr(p, "requires", frozenset())) | |
| all_provides |= set(getattr(p, "provides", frozenset())) | |
| self.requires: frozenset[str] = frozenset(all_requires) | |
| self.provides: frozenset[str] = frozenset(all_provides) | |
| # ------------------------------------------------------------------ | |
| # Sync execution | |
| # ------------------------------------------------------------------ | |
| def __call__(self, ctx: StepContext) -> StepContext: | |
| """Sync fan-out via ThreadPoolExecutor. | |
| All branches receive the same (frozen) context — no copy needed. | |
| All branches run to completion before any failure is raised. | |
| """ | |
| with ThreadPoolExecutor(max_workers=len(self.pipelines)) as executor: | |
| futures: list = [executor.submit(p, ctx) for p in self.pipelines] # type: ignore[arg-type] | |
| results: list[StepContext] = [] | |
| failures: list[BaseException] = [] | |
| for f in futures: | |
| try: | |
| results.append(f.result()) | |
| except BaseException as exc: # noqa: BLE001 | |
| failures.append(exc) | |
| if failures: | |
| raise BranchError(failures) | |
| return self._merge_fn(results) | |
| # ------------------------------------------------------------------ | |
| # Async execution | |
| # ------------------------------------------------------------------ | |
| async def __call_async__(self, ctx: StepContext) -> StepContext: | |
| """Async fan-out via asyncio.gather. | |
| ``return_exceptions=True`` guarantees all branches run to completion | |
| even when one fails; the full failure list is surfaced via | |
| ``BranchError``. | |
| Sync child pipelines are wrapped with ``asyncio.to_thread`` so they | |
| run in a thread pool rather than blocking the event loop. | |
| """ | |
| async def _run_child(child: object) -> StepContext: | |
| if asyncio.iscoroutinefunction(getattr(child, "__call__", None)): | |
| return await child(ctx) # type: ignore[operator] | |
| if hasattr(child, "__call_async__"): | |
| return await child.__call_async__(ctx) # type: ignore[union-attr] | |
| # Sync callable — run in thread pool so it doesn't block the loop | |
| return await asyncio.to_thread(child, ctx) # type: ignore[arg-type] | |
| raw = await asyncio.gather( | |
| *[_run_child(p) for p in self.pipelines], | |
| return_exceptions=True, | |
| ) | |
| failures = [r for r in raw if isinstance(r, BaseException)] | |
| if failures: | |
| raise BranchError(failures) | |
| return self._merge_fn([r for r in raw if not isinstance(r, BaseException)]) | |