Rewrite Choice step to make it usable from mlrun (#537)

* Rewrite Choice step to make it usable from mlrun

[ML-7818](https://iguazio.atlassian.net/browse/ML-7818)

* Add missing space

* Hack to avoid issue with mlrun preview

* Improve docs

* Remove accidental kwargs, add type annotation

storey/flow.py

@@ -252,15 +252,16 @@ def _event_string(event):
  def _should_terminate(self):
  return self._termination_received == len(self._inlets)
 
- async def _do_downstream(self, event):
- if not self._outlets:
+ async def _do_downstream(self, event, outlets=None):
+ outlets = self._outlets if outlets is None else outlets
+ if not outlets:
  return
  if event is _termination_obj:
  # Only propagate the termination object once we received one per inlet
- self._outlets[0]._termination_received += 1
- if self._outlets[0]._should_terminate():
- self._termination_result = await self._outlets[0]._do(_termination_obj)
- for outlet in self._outlets[1:] + self._get_recovery_steps():
+ outlets[0]._termination_received += 1
+ if outlets[0]._should_terminate():
+ self._termination_result = await outlets[0]._do(_termination_obj)
+ for outlet in outlets[1:] + self._get_recovery_steps():
  outlet._termination_received += 1
  if outlet._should_terminate():
  self._termination_result = self._termination_result_fn(
@@ -269,28 +270,28 @@ async def _do_downstream(self, event):
  return self._termination_result
  # If there is more than one outlet, allow concurrent execution.
  tasks = []
- if len(self._outlets) > 1:
+ if len(outlets) > 1:
  awaitable_result = event._awaitable_result
  event._awaitable_result = None
  original_events = getattr(event, "_original_events", None)
  # Temporarily delete self-reference to avoid deepcopy getting stuck in an infinite loop
  event._original_events = None
- for i in range(1, len(self._outlets)):
+ for i in range(1, len(outlets)):
  event_copy = copy.deepcopy(event)
  event_copy._awaitable_result = awaitable_result
  event_copy._original_events = original_events
- tasks.append(asyncio.get_running_loop().create_task(self._outlets[i]._do_and_recover(event_copy)))
+ tasks.append(asyncio.get_running_loop().create_task(outlets[i]._do_and_recover(event_copy)))
  # Set self-reference back after deepcopy
  event._original_events = original_events
  event._awaitable_result = awaitable_result
  if self.verbose and self.logger:
  step_name = self.name
  event_string = self._event_string(event)
- self.logger.debug(f"{step_name} -> {self._outlets[0].name} | {event_string}")
- await self._outlets[0]._do_and_recover(event) # Optimization - avoids creating a task for the first outlet.
+ self.logger.debug(f"{step_name} -> {outlets[0].name} | {event_string}")
+ await outlets[0]._do_and_recover(event) # Optimization - avoids creating a task for the first outlet.
  for i, task in enumerate(tasks, start=1):
  if self.verbose and self.logger:
- self.logger.debug(f"{step_name} -> {self._outlets[i].name} | {event_string}")
+ self.logger.debug(f"{step_name} -> {outlets[i].name} | {event_string}")
  await task
 
  def _get_event_or_body(self, event):
@@ -347,46 +348,48 @@ def _get_uuid(self):
 
 
 class Choice(Flow):
- """Redirects each input element into at most one of multiple downstreams.
-
- :param choice_array: a list of (downstream, condition) tuples, where downstream is a step and condition is a
- function. The first condition in the list to evaluate as true for an input element causes that element to
- be redirected to that downstream step.
- :type choice_array: tuple of (Flow, Function (Event=>boolean))
- :param default: a default step for events that did not match any condition in choice_array. If not set, elements
- that don't match any condition will be discarded.
- :type default: Flow
- :param name: Name of this step, as it should appear in logs. Defaults to class name (Choice).
- :type name: string
- :param full_event: Whether user functions should receive and return Event objects (when True),
- or only the payload (when False). Defaults to False.
- :type full_event: boolean
+ """
+ Redirects each input element into any number of predetermined downstream steps. Override select_outlets()
+ to route events to any number of downstream steps.
  """
 
- def __init__(self, choice_array, default=None, **kwargs):
- Flow.__init__(self, **kwargs)
-
- self._choice_array = choice_array
- for outlet, _ in choice_array:
- self.to(outlet)
-
- if default:
- self.to(default)
- self._default = default
+ def _init(self):
+ super()._init()
+ self._name_to_outlet = {}
+ for outlet in self._outlets:
+ if outlet.name in self._name_to_outlet:
+ raise ValueError(f"Ambiguous outlet name '{outlet.name}' in Choice step")
+ self._name_to_outlet[outlet.name] = outlet
+ # TODO: hacky way of supporting mlrun preview, which replaces targets with a DFTarget
+ self._passthrough_for_preview = list(self._name_to_outlet) == ["dataframe"]
+
+ def select_outlets(self, event) -> List[str]:
+ """
+ Override this method to route events based on a customer logic. The default implementation will route all
+ events to all outlets.
+ """
+ return list(self._name_to_outlet.keys())
 
  async def _do(self, event):
- if not self._outlets or event is _termination_obj:
- return await super()._do_downstream(event)
- chosen_outlet = None
- element = self._get_event_or_body(event)
- for outlet, condition in self._choice_array:
- if condition(element):
- chosen_outlet = outlet
- break
- if chosen_outlet:
- await chosen_outlet._do(event)
- elif self._default:
- await self._default._do(event)
+ if event is _termination_obj:
+ return await self._do_downstream(_termination_obj)
+ else:
+ event_body = event if self._full_event else event.body
+ outlet_names = self.select_outlets(event_body)
+ outlets = []
+ if self._passthrough_for_preview:
+ outlet = self._name_to_outlet["dataframe"]
+ outlets.append(outlet)
+ else:
+ for outlet_name in outlet_names:
+ if outlet_name not in self._name_to_outlet:
+ raise ValueError(
+ f"select_outlets() returned outlet name '{outlet_name}', which is not one of the "
+ f"defined outlets: " + ", ".join(self._name_to_outlet)
+ )
+ outlet = self._name_to_outlet[outlet_name]
+ outlets.append(outlet)
+ return await self._do_downstream(event, outlets=outlets)
 
 
 class Recover(Flow):

storey/sources.py

@@ -313,6 +313,7 @@ async def _run_loop(self):
  await _commit_handled_events(self._outstanding_offsets, committer, commit_all=True)
  self._termination_future.set_result(termination_result)
  except BaseException as ex:
+ traceback.print_exc()
  if self.logger:
  message = "An error was raised"
  raised_by = getattr(ex, "_raised_by_storey_step", None)

tests/test_flow.py

@@ -1690,26 +1690,42 @@ def boom(_):
 
 
 def test_choice():
- small_reduce = Reduce(0, lambda acc, x: acc + x)
+ class MyChoice(Choice):
+ def select_outlets(self, event):
+ outlets = ["all_events"]
+ if event > 5:
+ outlets.append("more_than_five")
+ else:
+ outlets.append("up_to_five")
+ return outlets
 
- big_reduce = build_flow([Map(lambda x: x * 100), Reduce(0, lambda acc, x: acc + x)])
+ source = SyncEmitSource()
+ my_choice = MyChoice(termination_result_fn=lambda x, y: x + y)
+ all_events = Map(lambda x: x, name="all_events")
+ more_than_five = Map(lambda x: x * 10, name="more_than_five")
+ up_to_five = Map(lambda x: x * 100, name="up_to_five")
+ sum_up_all_events = Reduce(0, lambda acc, x: acc + x)
+ sum_up_more_than_five = Reduce(0, lambda acc, x: acc + x)
+ sum_up_up_to_five = Reduce(0, lambda acc, x: acc + x)
+
+ source.to(my_choice)
+ my_choice.to(all_events)
+ my_choice.to(more_than_five)
+ my_choice.to(up_to_five)
+ all_events.to(sum_up_all_events)
+ more_than_five.to(sum_up_more_than_five)
+ up_to_five.to(sum_up_up_to_five)
 
- controller = build_flow(
- [
- SyncEmitSource(),
- Choice(
- [(big_reduce, lambda x: x % 2 == 0)],
- default=small_reduce,
- termination_result_fn=lambda x, y: x + y,
- ),
- ]
- ).run()
+ controller = source.run()
 
- for i in range(10):
+ for i in range(4, 8):
  controller.emit(i)
+
  controller.terminate()
  termination_result = controller.await_termination()
- assert termination_result == 2025
+
+ expected = sum(range(4, 8)) + sum(range(6, 8)) * 10 + sum(range(4, 6)) * 100
+ assert termination_result == expected
 
 
 def test_metadata():