def create_data_writer(self, session_id, chunk_key, data_size, sender_ref,
ensure_cached=True, timeout=0, callback=None):
from mars.compat import BytesIO
query_key = (session_id, chunk_key)
if query_key in self._data_metas and \
self._data_metas[query_key].status in (ReceiveStatus.RECEIVED, ReceiveStatus.RECEIVING):
self.tell_promise(callback, self.address, self._data_metas[query_key].status)
return
self._data_metas[query_key] = ReceiverDataMeta(chunk_size=data_size, status=ReceiveStatus.RECEIVING)
self._data_writers[query_key] = BytesIO()
self.tell_promise(callback, self.address, None)
def create_data_writers(self, session_id, chunk_keys, data_sizes, sender_ref,
ensure_cached=True, pin_token=None, timeout=0,
use_promise=True, callback=None):
for chunk_key, data_size in zip(chunk_keys, data_sizes):
query_key = (session_id, chunk_key)
if query_key in self._data_metas and \
self._data_metas[query_key].status in (ReceiveStatus.RECEIVED, ReceiveStatus.RECEIVING):
self.tell_promise(callback, self.address, self._data_metas[query_key].status)
return
self._data_metas[query_key] = ReceiverDataMeta(chunk_size=data_size, status=ReceiveStatus.RECEIVING)
self._data_writers[query_key] = BytesIO()
if callback:
self.tell_promise(callback)
def set_status(self, session_id, chunk_key, status):
query_key = (session_id, chunk_key)
try:
self._data_metas[query_key].status = status
except KeyError:
self._data_metas[query_key] = ReceiverDataMeta(status=status)