"""Flow node that collects one message per input."""
import numbers
from collections import deque
from collections.abc import Mapping, Sequence
from typing import Any, ClassVar
from heros import HERORef
from herostools.flow.node import _FLOW_PATH_KEY, Node, _flow_path
from herostools.helper import log, merge_dicts
[docs]
class Join(Node):
"""Node that collects one message per input and then processes them together.
Messages are grouped into rounds by their identifier. `_process` runs once per completed
round, with the data of every input keyed by input name, in config order. The default
`_process` emits that dict. The rules:
- A round completes when every configured input has delivered. The metadata of the inputs is
merged in config order. Their flow paths are joined in config order without duplicates.
Completed rounds are processed in completion order.
- At most `max_open_rounds` incomplete rounds are open. A message with a new identifier at the
limit drops the oldest incomplete round and logs a warning that names its identifier and
missing inputs. The warning is logged once per crossing of the limit, and again after the
number of open rounds fell below the limit.
- The last `max_open_rounds` finished or dropped identifiers are remembered. Later messages
for them are ignored.
- A repeated message from an input for an open round replaces the earlier one.
- A message that fails validation does not count.
- `feed` processes its message alone, as input `"manual"`. It never joins a round.
Rounds are only touched on the worker, so they need no lock.
"""
_hero_implements: ClassVar[list[str]] = [*Node._hero_implements, "herostools.flow.Join"]
def __init__(
self,
inputs: Mapping[str, HERORef | str] | Sequence[HERORef | str],
*,
max_open_rounds: int = 10,
**kwargs,
):
"""Initialize the join.
Args:
inputs: Sources to collect, as for `Node`. Must not be empty.
max_open_rounds: Number of incomplete rounds kept open, and of finished identifiers
remembered.
**kwargs: Passed on to `Node`.
Raises:
ValueError: If `inputs` is empty or `max_open_rounds` is not positive.
TypeError: If `max_open_rounds` is not an integer (bool is refused).
"""
super().__init__(inputs, **kwargs)
if not self._inputs:
raise ValueError("a join needs at least one input")
if not isinstance(max_open_rounds, numbers.Integral) or isinstance(max_open_rounds, bool):
raise TypeError(f"max_open_rounds must be an int, got {max_open_rounds!r}")
if max_open_rounds <= 0:
raise ValueError(f"max_open_rounds must be positive, got {max_open_rounds!r}")
self._max_open_rounds = int(max_open_rounds)
self._rounds: dict[Any, dict[str, tuple[Any, dict]]] = {}
self._finished: deque = deque(maxlen=self._max_open_rounds)
self._warned = False
[docs]
def _process(self, inputs: dict[str, Any], metadata: dict) -> tuple[Any, dict] | None:
"""Emit the dict of all inputs. Override in subclasses to combine them."""
return dict(inputs), {}
[docs]
def _collect(self, input_name: str, data: Any, metadata: dict) -> tuple[dict, dict] | None:
"""Add the message to its round and return the round once every input delivered."""
identifier = metadata[self._identifier_key]
if identifier in self._finished:
log.debug(f"{self.name}: Ignoring late data for finished round {identifier!r}")
return None
if len(self._rounds) < self._max_open_rounds:
self._warned = False
if identifier not in self._rounds:
if len(self._rounds) >= self._max_open_rounds:
self._drop_oldest_round()
self._rounds[identifier] = {}
collected = self._rounds[identifier]
collected[input_name] = (data, metadata)
if len(collected) < len(self._inputs):
return None
del self._rounds[identifier]
self._finished.append(identifier)
ordered = [collected[name] for name in self._inputs]
merged: dict = {}
for _, input_metadata in ordered:
merged = merge_dicts(merged, input_metadata)
merged[_FLOW_PATH_KEY] = list(dict.fromkeys(name for _, m in ordered for name in _flow_path(m)))
return {name: item[0] for name, item in zip(self._inputs, ordered, strict=True)}, merged
[docs]
def _drop_oldest_round(self) -> None:
"""Drop the oldest incomplete round, with a warning once per crossing of the limit."""
identifier = next(iter(self._rounds))
dropped = self._rounds.pop(identifier)
missing = [name for name in self._inputs if name not in dropped]
self._finished.append(identifier)
if not self._warned:
self._warned = True
log.warning(f"{self.name}: Dropped incomplete round {identifier!r}, missing inputs {missing}")