Add support for multi task operators and stream context (#59)

* Add support for multi-task operators and stream contexts.

* Update comment.

* Rename method.

Author: doru1004

examples/multitask/multitask.py

@@ -0,0 +1,81 @@
+#
+# Copyright IBM Corporation 2021
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import ray
+import rayvens
+import sys
+
+# Send message to Slack sink using a multi-tasking operator. The operator
+# will be a Ray Task which spawns three other Ray Tasks. Each of the three
+# spawned tasks will send their event to the sinks attached to the stream.
+
+# Command line arguments and validation:
+if len(sys.argv) < 4:
+    print(f'usage: {sys.argv[0]} <slack_channel> <slack_webhook> <run_mode>')
+    sys.exit(1)
+slack_channel = sys.argv[1]
+slack_webhook = sys.argv[2]
+run_mode = sys.argv[3]
+if run_mode not in ['local', 'mixed', 'operator']:
+    raise RuntimeError(f'Invalid run mode provided: {run_mode}')
+
+# Initialize ray either on the cluster or locally otherwise.
+if run_mode == 'operator':
+    ray.init(address='auto')
+else:
+    ray.init()
+
+# Start rayvens in operator mode.
+rayvens.init(mode=run_mode)
+
+# Create stream.
+stream = rayvens.Stream('slack')
+
+
+# Operator sub-task:
+@ray.remote
+def sub_task(context, intermediate_data):
+    sub_task_outgoing_event = "sub-task " + intermediate_data
+    print(sub_task_outgoing_event)
+    context.publish(sub_task_outgoing_event)
+
+
+# Operator task:
+@ray.remote
+def multi_part_task(context, incoming_event):
+    print("multi-part-task:", incoming_event)
+    for i in range(3):
+        sub_task.remote(context, "sub-event" + str(i))
+
+
+# Event sink config.
+sink_config = dict(kind='slack-sink',
+                   route='/toslack',
+                   channel=slack_channel,
+                   webhook_url=slack_webhook)
+
+# Add sink to stream.
+sink = stream.add_sink(sink_config)
+
+# Add multi-task operator to stream.
+stream.add_multitask_operator(multi_part_task)
+
+# Sends message to all sinks attached to this stream.
+stream << f'Sending message to Slack sink in run mode {run_mode}.'
+
+# Disconnect any sources or sinks attached to the stream 2 seconds after
+# the stream is idle (i.e. no events were propagated by the stream).
+stream.disconnect_all(after_idle_for=2)

rayvens/api.py

@@ -56,6 +56,9 @@ def append(self, data):
     def add_operator(self, operator):
         return ray.get(self.actor.add_operator.remote(operator))
 
+    def add_multi_operator(self, operator):
+        return ray.get(self.actor.add_multi_operator.remote(operator))
+
     def add_source(self, source_config):
         return ray.get(self.actor.add_source.remote(self, source_config))
 
@@ -112,45 +115,72 @@ def _idle_time(self):
         return time.time() - latest_timestamp
 
 
+class StreamContext:
+    def __init__(self):
+        self.sink_restrictions = {}
+        self.subscribers = {}
+        self.latest_sent_event_timestamp = None
+        self.event_counter = 0
+        self.limit_subscribers = False
+        self.is_multi_operator = False
+
+    def publish(self, data):
+        if data is not None:
+            for name, subscriber in self.subscribers.items():
+                if name in self.sink_restrictions:
+                    type_restrictions = self.sink_restrictions[name]
+                    if not self._accepts_data_type(data, type_restrictions):
+                        continue
+                _eval(self, subscriber, data)
+
+    # Check if the sink we are routing the message to has any restrictions
+    # in terms of message type. A message will only be routed to a sink
+    # if the sink accepts its type.
+    def _accepts_data_type(self, data, type_restrictions):
+        # If there are no restrictions return immediately:
+        if len(type_restrictions) == 0:
+            return True
+        for restricted_type in type_restrictions:
+            if isinstance(data, restricted_type):
+                return True
+        return False
+
+
 @ray.remote(num_cpus=0)
 class StreamActor:
     def __init__(self, name, operator=None):
         self.name = name
-        self._subscribers = {}
         self._operator = operator
         self._sources = {}
         self._sinks = {}
-        self._latest_sent_event_timestamp = None
-        self._limit_subscribers = False
-        self._event_counter = 0
+        self.context = StreamContext()
 
     def send_to(self, subscriber, name=None):
-        if self._limit_subscribers:
+        if self.context.limit_subscribers:
             return
-        if name in self._subscribers:
+        if name in self.context.subscribers:
             raise RuntimeError(
                 f'Stream {self.name} already has a subscriber named {name}.')
         if name is None:
             name = object()
-        self._subscribers[name] = subscriber
+        self.context.subscribers[name] = subscriber
 
     def append(self, data):
         if data is None:
             return
         if self._operator is not None:
-            data = _eval(self._operator, data)
-        for name, subscriber in self._subscribers.items():
-            if name in self._sinks:
-                integration = self._sinks[name]
-                if not integration.accepts_data_type(data):
-                    continue
-            _eval(subscriber, data)
-        self._latest_sent_event_timestamp = time.time()
-        self._event_counter += 1
+            data = _eval(self.context, self._operator, data)
+        self.context.publish(data)
+        self.context.latest_sent_event_timestamp = time.time()
+        self.context.event_counter += 1
 
     def add_operator(self, operator):
         self._operator = operator
 
+    def add_multitask_operator(self, operator):
+        self._operator = operator
+        self.context.is_multi_operator = True
+
     def add_source(self, stream, source_config):
         source_config["integration_type"] = 'source'
         source_name = name_source(source_config)
@@ -170,13 +200,15 @@ def add_sink(self, stream, sink_config):
                 f'Stream {self.name} already has a sink named {sink_name}.')
         self._sinks[sink_name] = _global_camel.add_sink(
             stream, sink_config, sink_name)
+        self.context.sink_restrictions[sink_name] = self._sinks[
+            sink_name].get_restricted_data_type()
         return sink_name
 
     def unsubscribe(self, subscriber_name):
         if subscriber_name not in self._subscribers:
             raise RuntimeError(f'Stream {self.name} has no subscriber named'
                                f' {subscriber_name}.')
-        self._subscribers.pop(subscriber_name)
+        self.context.subscribers.pop(subscriber_name)
 
     def disconnect_source(self, source_name):
         if source_name not in self._sources:
@@ -191,7 +223,8 @@ def disconnect_sink(self, sink_name):
                 f'Stream {self.name} has no sink named {sink_name}.')
         _global_camel.disconnect(self._sinks[sink_name])
         self._sinks.pop(sink_name)
-        self._subscribers.pop(sink_name)
+        self.context.sink_restrictions.pop(sink_name)
+        self.context.subscribers.pop(sink_name)
 
     def disconnect_all(self, stream_drain_timeout):
         for source_name in dict(self._sources):
@@ -201,29 +234,35 @@ def disconnect_all(self, stream_drain_timeout):
             self.disconnect_sink(sink_name)
 
     def event_count(self):
-        return self.event_count
+        return self.context.event_counter
 
     def _meta(self, action, *args, **kwargs):
         return verify_do(self, _global_camel, action, *args, **kwargs)
 
     def _get_latest_timestamp(self):
-        return self._latest_sent_event_timestamp
+        return self.context.latest_sent_event_timestamp
 
     def _fetch_processors(self):
-        self._limit_subscribers = True
-        return self._subscribers, self._operator
+        self.context.limit_subscribers = True
+        return self.context.subscribers, self._operator
 
     def _update_timestamp(self, timestamp):
-        self._latest_sent_event_timestamp = timestamp
+        self.context.latest_sent_event_timestamp = timestamp
 
 
-def _eval(f, data):
+def _eval(context, f, data):
     if isinstance(f, Stream):
         return f.append(data)
     elif isinstance(f, ray.actor.ActorHandle):
         return f.append.remote(data)
-    elif isinstance(f, ray.actor.ActorMethod) or isinstance(
-            f, ray.remote_function.RemoteFunction):
+    elif isinstance(f, ray.actor.ActorMethod):
+        return f.remote(data)
+    elif isinstance(f, ray.remote_function.RemoteFunction):
+        if context.is_multi_operator:
+            if context.subscribers is None:
+                raise RuntimeError('No subscribers or sinks provided.')
+            f.remote(context, data)
+            return None
         return f.remote(data)
     else:
         return f(data)

rayvens/core/integration.py

@@ -132,19 +132,10 @@ def disconnect(self, mode):
         # process that runs it.
         self.invocation.kill()
 
-    # Check if the sink we are routing the message to has any restrictions
-    # in terms of message type. A message will only be routed to a sink
-    # if the sink accepts its type.
-    def accepts_data_type(self, data):
+    # Get any type restrictions:
+    def get_restricted_data_type(self):
         # If there are no restrictions return immediately:
-        restricted_message_types = self.input_restrictions[
-            'restricted_message_types']
-        if len(restricted_message_types) == 0:
-            return True
-        for restricted_type in restricted_message_types:
-            if isinstance(data, restricted_type):
-                return True
-        return False
+        return self.input_restrictions['restricted_message_types']
 
     # Method that checks if, based on the configuration, the integration
     # requires something to be run or created before the integration is run.