def _get_engine_config(self):
op_def = op_def_pb2.OpDef(op=types_pb2.GET_ENGINE_CONFIG)
dag_def = op_def_pb2.DagDef()
dag_def.op.extend([op_def])
fetch_request = message_pb2.RunStepRequest(
session_id=self._session_id, dag_def=dag_def
)
fetch_response = self._analytical_engine_stub.RunStep(fetch_request)
config = json.loads(fetch_response.result.decode("utf-8"))
if self._launcher_type == types_pb2.K8S:
config["vineyard_service_name"] = self._launcher.get_vineyard_service_name()
config["vineyard_rpc_endpoint"] = self._launcher.get_vineyard_rpc_endpoint()
return config
def generate_runstep_requests(self, session_id, dag_def):
runstep_requests = []
chunks_list = self.split(dag_def)
# head
runstep_request = message_pb2.RunStepRequest(
head=message_pb2.RunStepRequestHead(session_id=session_id,
dag_def=dag_def))
runstep_requests.append(runstep_request)
# bodies
for chunks, op_key in chunks_list:
for i, chunk in enumerate(chunks):
# check the last element
has_next = True
if i + 1 == len(chunks):
has_next = False
runstep_request = message_pb2.RunStepRequest(
body=message_pb2.RunStepRequestBody(
chunk=chunk, op_key=op_key, has_next=has_next))
runstep_requests.append(runstep_request)
# return a generator for stream request
for item in runstep_requests:
yield item
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 _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_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 _run_step_impl(self, dag_def):
request = message_pb2.RunStepRequest(session_id=self._session_id,
dag_def=dag_def)
response = self._stub.RunStep(request)
return check_grpc_response(response)
