wassname/ray
1
0
Fork 0
mirror of https://github.com/wassname/ray.git synced 2026-08-11 05:51:40 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Streaming] Streaming Cross-Lang API (#7464)

Browse Source
This commit is contained in:
chaokunyang
2020-04-29 13:42:08 +08:00
committed by GitHub
parent 101255f782
commit 91f630f709
72 changed files with 1612 additions and 408 deletions
Download Patch File Download Diff File Expand all files Collapse all files
streaming/python/collector.py
+38 -8
View File
@@ -1,10 +1,13 @@
import logging
import pickle
import typing
from abc import ABC, abstractmethod
from ray import Language
from ray.actor import ActorHandle
from ray.streaming import function
from ray.streaming import message
from ray.streaming import partition
from ray.streaming.runtime import serialization
from ray.streaming.runtime.transfer import ChannelID, DataWriter
logger = logging.getLogger(__name__)
@@ -31,19 +34,46 @@ class CollectionCollector(Collector):
class OutputCollector(Collector):
def __init__(self, channel_ids: typing.List[str], writer: DataWriter,
def __init__(self, writer: DataWriter, channel_ids: typing.List[str],
target_actors: typing.List[ActorHandle],
partition_func: partition.Partition):
self._channel_ids = [ChannelID(id_str) for id_str in channel_ids]
self._writer = writer
self._channel_ids = [ChannelID(id_str) for id_str in channel_ids]
self._target_languages = []
for actor in target_actors:
if actor._ray_actor_language == Language.PYTHON:
self._target_languages.append(function.Language.PYTHON)
elif actor._ray_actor_language == Language.JAVA:
self._target_languages.append(function.Language.JAVA)
else:
raise Exception("Unsupported language {}"
.format(actor._ray_actor_language))
self._partition_func = partition_func
self.python_serializer = serialization.PythonSerializer()
self.cross_lang_serializer = serialization.CrossLangSerializer()
logger.info(
"Create OutputCollector, channel_ids {}, partition_func {}".format(
channel_ids, partition_func))
def collect(self, record):
partitions = self._partition_func.partition(record,
len(self._channel_ids))
serialized_message = pickle.dumps(record)
partitions = self._partition_func \
.partition(record, len(self._channel_ids))
python_buffer = None
cross_lang_buffer = None
for partition_index in partitions:
self._writer.write(self._channel_ids[partition_index],
serialized_message)
if self._target_languages[partition_index] == \
function.Language.PYTHON:
# avoid repeated serialization
if python_buffer is None:
python_buffer = self.python_serializer.serialize(record)
self._writer.write(
self._channel_ids[partition_index],
serialization._PYTHON_TYPE_ID + python_buffer)
else:
# avoid repeated serialization
if cross_lang_buffer is None:
cross_lang_buffer = self.cross_lang_serializer.serialize(
record)
self._writer.write(
self._channel_ids[partition_index],
serialization._CROSS_LANG_TYPE_ID + cross_lang_buffer)