def _fetch(self, dataobj): "Fetch the object data and update its state." if not dataobj._keep: raise RainException( "Can't fetch object {} without keep flag.".format(dataobj)) if dataobj.state is None: raise RainException("Object {} is not submitted.".format(dataobj)) req = self._datastore.createReader_request() id_to_capnp(dataobj.id, req.id) req.offset = 0 result = req.send().wait() check_result((dataobj.session, ), result) reader = result.reader FETCH_SIZE = 2 << 20 # 2MB eof = False data = [] while not eof: r = reader.read(FETCH_SIZE).wait() data.append(r.data) eof = r.status == "eof" bytedata = b"".join(data) self._get_state((), (dataobj, )) return DataInstance(data=bytedata, data_object=dataobj, data_type=DataType.from_capnp(result.dataType))
def check_result(result): if result.which() == "ok": return # Do nothing elif result.which() == "error": message = result.error.message if result.error.debug: message += "\nDebug:\n" + result.error.debug raise RainException(message) else: raise Exception("Invalid result")
def _fetch(self, dataobj): "Fetch the object data and update its state." if not dataobj._keep: raise RainException( "Can't fetch object {} without keep flag.".format(dataobj)) if dataobj.state is None: raise RainException("Object {} is not submitted.".format(dataobj)) req = self._service.fetch_request() id_to_capnp(dataobj.id, req.id) req.offset = 0 req.size = FETCH_SIZE req.includeMetadata = True result = req.send().wait() check_result((dataobj.session, ), result.status) size = result.metadata.size offset = len(result.data) data = [result.data] while offset < size: req = self._service.fetch_request() id_to_capnp(dataobj.id, req.id) req.offset = offset req.size = FETCH_SIZE req.includeMetadata = False r = req.send().wait() check_result((dataobj.session, ), r.status) data.append(r.data) offset += len(r.data) rawdata = b"".join(data) dataobj.attributes.update( attributes.attributes_from_capnp(result.metadata.attributes)) return DataInstance(data=rawdata, data_object=dataobj, data_type=DataType.from_capnp( result.metadata.dataType))
def split_items(items): """Split items into 'tasks' and 'dataobjects' Throws an error if an item is not task nor object""" tasks = [] dataobjects = [] for item in items: if isinstance(item, Task): tasks.append(item) elif isinstance(item, DataObject): dataobjects.append(item) else: raise RainException( "'{}' is not tasks nor dataobject".format(item)) return tasks, dataobjects
def _wait(self, tasks, dataobjs): req = self._service.wait_request() req.init("taskIds", len(tasks)) for i in range(len(tasks)): task = tasks[i] if task.state is None: raise RainException("Task {} is not submitted".format(task)) id_to_capnp(task.id, req.taskIds[i]) req.init("objectIds", len(dataobjs)) for i in range(len(dataobjs)): id_to_capnp(dataobjs[i].id, req.objectIds[i]) result = req.send().wait() check_result(result)