def _cleanup(self, cleanup_instance=True, is_dangling=False):
# clean up session resources.
for key in self._object_manager.keys():
obj = self._object_manager.get(key)
obj_type = obj.type
unload_type = None
if obj_type == "app":
unload_type = types_pb2.UNLOAD_APP
config = {
types_pb2.APP_NAME: attr_value_pb2.AttrValue(
s=obj.key.encode("utf-8")
)
}
elif obj_type == "graph":
unload_type = types_pb2.UNLOAD_GRAPH
config = {
types_pb2.GRAPH_NAME: attr_value_pb2.AttrValue(
s=obj.key.encode("utf-8")
)
}
# dynamic graph doesn't have a vineyard id
if obj.vineyard_id != -1:
config[types_pb2.VINEYARD_ID] = attr_value_pb2.AttrValue(
i=obj.vineyard_id
)
if unload_type:
dag_def = create_single_op_dag(unload_type, config)
request = message_pb2.RunStepRequest(
session_id=self._session_id, dag_def=dag_def
)
self._analytical_engine_stub.RunStep(request)
self._object_manager.clear()
self._request = None
# cancel dangling detect timer
if self._dangling_detecting_timer:
self._dangling_detecting_timer.cancel()
# close engines
if cleanup_instance:
self._analytical_engine_stub = None
self._analytical_engine_endpoint = None
self._launcher.stop(is_dangling=is_dangling)
self._session_id = None
def _get_config(
edges: Sequence[EdgeLabel],
vertices: Sequence[VertexLabel],
directed: bool,
oid_type: str,
generate_eid: bool,
) -> Dict:
config = {}
attr = attr_value_pb2.AttrValue()
for label in chain(edges, vertices):
label.finish(oid_type)
for edge in edges:
attr.list.func.extend([process_edge(edge)])
attr.list.func.extend([process_vertex(vertex) for vertex in vertices])
directed_attr = utils.b_to_attr(directed)
generate_eid_attr = utils.b_to_attr(generate_eid)
config[types_pb2.ARROW_PROPERTY_DEFINITION] = attr
config[types_pb2.DIRECTED] = directed_attr
config[types_pb2.OID_TYPE] = utils.s_to_attr(oid_type)
config[types_pb2.GENERATE_EID] = generate_eid_attr
# vid_type is fixed
config[types_pb2.VID_TYPE] = utils.s_to_attr("uint64_t")
config[types_pb2.IS_FROM_VINEYARD_ID] = utils.b_to_attr(False)
return config
def get_attr(self):
if not self.finished:
self.finish()
attr = attr_value_pb2.AttrValue()
attr.func.name = "loader"
attr.func.attr[types_pb2.PROTOCOL].CopyFrom(utils.s_to_attr(self.protocol))
# Let graphscope handle local files cause it's implemented in c++ and
# doesn't add an additional stream layer.
# Maybe handled by vineyard in the near future
if self.protocol == "file":
source = "{}#{}".format(self.source, self.options)
attr.func.attr[types_pb2.VALUES].CopyFrom(
utils.bytes_to_attr(source.encode("utf-8"))
)
elif self.protocol in ("numpy", "pandas"):
attr.func.attr[types_pb2.ROW_NUM].CopyFrom(utils.i_to_attr(self.row_num))
attr.func.attr[types_pb2.COLUMN_NUM].CopyFrom(
utils.i_to_attr(self.column_num)
)
# Use key start from 10000 + col_index to store raw bytes.
for i in range(len(self.property_bytes)):
attr.func.attr[10000 + i].CopyFrom(
utils.bytes_to_attr(self.property_bytes[i])
)
else: # Let vineyard handle other data source.
attr.func.attr[types_pb2.VALUES].CopyFrom(
utils.bytes_to_attr(self.source.encode("utf-8"))
)
return attr
def list_str_to_attr(list_of_str):
attr = attr_value_pb2.AttrValue()
attr.list.s[:] = [
item.encode("utf-8") if isinstance(item, str) else item
for item in list_of_str
]
return attr
def _maybe_register_graph(self, op, session_id):
graph_sig = get_graph_sha256(op.attr)
space = self._builtin_workspace
graph_lib_path = get_lib_path(os.path.join(space, graph_sig), graph_sig)
if not os.path.isfile(graph_lib_path):
compiled_path = self._compile_lib_and_distribute(
compile_graph_frame, graph_sig, op
)
if graph_lib_path != compiled_path:
raise RuntimeError("Computed path not equal to compiled path.")
if graph_sig not in self._object_manager:
# register graph
op_def = op_def_pb2.OpDef(op=types_pb2.REGISTER_GRAPH_TYPE)
op_def.attr[types_pb2.GRAPH_LIBRARY_PATH].CopyFrom(
attr_value_pb2.AttrValue(s=graph_lib_path.encode("utf-8"))
)
op_def.attr[types_pb2.TYPE_SIGNATURE].CopyFrom(
attr_value_pb2.AttrValue(s=graph_sig.encode("utf-8"))
)
op_def.attr[types_pb2.GRAPH_TYPE].CopyFrom(
attr_value_pb2.AttrValue(
graph_type=op.attr[types_pb2.GRAPH_TYPE].graph_type
)
)
dag_def = op_def_pb2.DagDef()
dag_def.op.extend([op_def])
register_request = message_pb2.RunStepRequest(
session_id=session_id, dag_def=dag_def
)
register_response = self._analytical_engine_stub.RunStep(register_request)
if register_response.status.code == error_codes_pb2.OK:
self._object_manager.put(
graph_sig,
LibMeta(register_response.result, "graph_frame", graph_lib_path),
)
else:
raise RuntimeError("Error occur when register graph")
op.attr[types_pb2.TYPE_SIGNATURE].CopyFrom(
attr_value_pb2.AttrValue(s=graph_sig.encode("utf-8"))
)
return op
def _maybe_compile_app(self, op):
app_sig = self._generate_app_sig(op.attr)
if app_sig in self._object_manager:
app_lib_path = self._object_manager.get(app_sig).lib_path
else:
app_lib_path = self._compile_lib_and_distribute(compile_app, app_sig, op)
op.attr[types_pb2.APP_LIBRARY_PATH].CopyFrom(
attr_value_pb2.AttrValue(s=app_lib_path.encode("utf-8"))
)
return op, app_sig, app_lib_path
def _construct_op_of_empty_graph(self):
config = {}
config[types_pb2.ARROW_PROPERTY_DEFINITION] = attr_value_pb2.AttrValue()
config[types_pb2.DIRECTED] = utils.b_to_attr(self._directed)
config[types_pb2.GENERATE_EID] = utils.b_to_attr(self._generate_eid)
config[types_pb2.OID_TYPE] = utils.s_to_attr(self._oid_type)
config[types_pb2.VID_TYPE] = utils.s_to_attr("uint64_t")
config[types_pb2.IS_FROM_VINEYARD_ID] = utils.b_to_attr(False)
return dag_utils.create_graph(
self.session_id, graph_def_pb2.ARROW_PROPERTY, inputs=None, attrs=config
)
def _maybe_register_graph(self, op, session_id):
graph_sig = self._generate_graph_sig(op.attr)
if graph_sig in self._object_manager:
lib_meta = self._object_manager.get(graph_sig)
graph_lib_path = lib_meta.lib_path
else:
graph_lib_path = self._compile_lib_and_distribute(
compile_graph_frame, graph_sig, op
)
# register graph
op_def = op_def_pb2.OpDef(op=types_pb2.REGISTER_GRAPH_TYPE)
op_def.attr[types_pb2.GRAPH_LIBRARY_PATH].CopyFrom(
attr_value_pb2.AttrValue(s=graph_lib_path.encode("utf-8"))
)
op_def.attr[types_pb2.TYPE_SIGNATURE].CopyFrom(
attr_value_pb2.AttrValue(s=graph_sig.encode("utf-8"))
)
op_def.attr[types_pb2.GRAPH_TYPE].CopyFrom(
attr_value_pb2.AttrValue(
graph_type=op.attr[types_pb2.GRAPH_TYPE].graph_type
)
)
dag_def = op_def_pb2.DagDef()
dag_def.op.extend([op_def])
register_request = message_pb2.RunStepRequest(
session_id=session_id, dag_def=dag_def
)
register_response = self._analytical_engine_stub.RunStep(register_request)
if register_response.status.code == error_codes_pb2.OK:
self._object_manager.put(
graph_sig,
LibMeta(register_response.result, "graph_frame", graph_lib_path),
)
else:
raise RuntimeError("Error occur when register graph")
op.attr[types_pb2.TYPE_SIGNATURE].CopyFrom(
attr_value_pb2.AttrValue(s=graph_sig.encode("utf-8"))
)
return op
def _cleanup(self, is_dangling=False):
# clean up session resources.
for key in self._object_manager.keys():
obj = self._object_manager.get(key)
obj_type = obj.type
unload_type = None
if obj_type == "app":
unload_type = types_pb2.UNLOAD_APP
config = {
types_pb2.APP_NAME:
attr_value_pb2.AttrValue(s=obj.key.encode("utf-8"))
}
elif obj_type == "graph":
unload_type = types_pb2.UNLOAD_GRAPH
config = {
types_pb2.GRAPH_NAME:
attr_value_pb2.AttrValue(s=obj.key.encode("utf-8"))
}
if unload_type:
dag_def = create_single_op_dag(unload_type, config)
request = message_pb2.RunStepRequest(
session_id=self._session_id, dag_def=dag_def)
self._analytical_engine_stub.RunStep(request)
self._object_manager.clear()
self._analytical_engine_stub = None
# cancel dangling detect timer
self._dangling_detecting_timer.cancel()
# close engines
self._launcher.stop(is_dangling=is_dangling)
self._session_id = None
self._analytical_engine_endpoint = None
def _maybe_compile_app(self, op):
app_sig = get_app_sha256(op.attr)
space = self._builtin_workspace
if types_pb2.GAR in op.attr:
space = self._udf_app_workspace
app_lib_path = get_lib_path(os.path.join(space, app_sig), app_sig)
if not os.path.isfile(app_lib_path):
compiled_path = self._compile_lib_and_distribute(compile_app, app_sig, op)
if app_lib_path != compiled_path:
raise RuntimeError("Computed path not equal to compiled path.")
op.attr[types_pb2.APP_LIBRARY_PATH].CopyFrom(
attr_value_pb2.AttrValue(s=app_lib_path.encode("utf-8"))
)
return op, app_sig, app_lib_path
def get_attr(self):
attr = attr_value_pb2.AttrValue()
attr.func.name = "loader"
attr.func.attr[types_pb2.PROTOCOL].CopyFrom(
utils.s_to_attr(self.protocol))
if self.protocol in ("file", "oss", "vineyard", "mars"):
source = "{}#{}".format(self.source, self.options)
attr.func.attr[types_pb2.VALUES].CopyFrom(
utils.bytes_to_attr(source.encode("utf-8")))
elif self.protocol in ("numpy", "pandas"):
attr.func.attr[types_pb2.ROW_NUM].CopyFrom(
utils.i_to_attr(self.row_num))
attr.func.attr[types_pb2.COLUMN_NUM].CopyFrom(
utils.i_to_attr(self.column_num))
# Use key start from 10000 + col_index to store raw bytes.
for i in range(len(self.property_bytes)):
attr.func.attr[10000 + i].CopyFrom(
utils.bytes_to_attr(self.property_bytes[i]))
else:
raise TypeError("Protocol not recognized " + self.protocol)
return attr
def modify_type_to_attr(t):
return attr_value_pb2.AttrValue(modify_type=t)
def place_holder_to_attr():
return attr_value_pb2.AttrValue(place_holder=types_pb2.PlaceHolder())
def list_i_to_attr(list_i):
attr = attr_value_pb2.AttrValue()
attr.list.i[:] = list_i
return attr
def i_to_attr(i: int) -> attr_value_pb2.AttrValue:
check_argument(isinstance(i, int))
return attr_value_pb2.AttrValue(i=i)
def b_to_attr(b: bool) -> attr_value_pb2.AttrValue:
check_argument(isinstance(b, bool))
return attr_value_pb2.AttrValue(b=b)
def s_to_attr(s: str) -> attr_value_pb2.AttrValue:
check_argument(isinstance(s, str))
return attr_value_pb2.AttrValue(s=s.encode("utf-8"))
def bytes_to_attr(s: bytes) -> attr_value_pb2.AttrValue:
check_argument(isinstance(s, bytes))
return attr_value_pb2.AttrValue(s=s)
def report_type_to_attr(t):
return attr_value_pb2.AttrValue(report_type=t)
def u_to_attr(i: int) -> attr_value_pb2.AttrValue:
check_argument(isinstance(i, int) and i >= 0)
return attr_value_pb2.AttrValue(u=i)
def graph_type_to_attr(t):
return attr_value_pb2.AttrValue(graph_type=t)
def type_to_attr(t):
return attr_value_pb2.AttrValue(type=t)
def f_to_attr(f: float) -> attr_value_pb2.AttrValue:
check_argument(isinstance(f, float))
return attr_value_pb2.AttrValue(f=f)