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)