| 38 | |
| 39 | |
| 40 | class GRPCClient(object): |
| 41 | def __init__(self, launcher, endpoint, reconnect=False): |
| 42 | """Connect to GRAPE engine at the given :code:`endpoint`.""" |
| 43 | # create the gRPC stub |
| 44 | self._options = [ |
| 45 | ("grpc.max_send_message_length", GS_GRPC_MAX_MESSAGE_LENGTH), |
| 46 | ("grpc.max_receive_message_length", GS_GRPC_MAX_MESSAGE_LENGTH), |
| 47 | ("grpc.max_metadata_size", GS_GRPC_MAX_MESSAGE_LENGTH), |
| 48 | ] |
| 49 | self._endpoint = endpoint |
| 50 | self._launcher = launcher |
| 51 | self._grpc_utils = GRPCUtils() |
| 52 | self._stub = self._get_stub() |
| 53 | self._session_id = None |
| 54 | self._logs_fetching_thread = None |
| 55 | self._reconnect = reconnect |
| 56 | |
| 57 | def _get_stub(self): |
| 58 | channel = grpc.insecure_channel(self._endpoint, options=self._options) |
| 59 | return coordinator_service_pb2_grpc.CoordinatorServiceStub(channel) |
| 60 | |
| 61 | def waiting_service_ready(self, timeout_seconds=60): |
| 62 | begin_time = time.time() |
| 63 | request = message_pb2.HeartBeatRequest() |
| 64 | while True: |
| 65 | if self._launcher: |
| 66 | code = self._launcher.poll() |
| 67 | if code is not None and code != 0: |
| 68 | raise RuntimeError( |
| 69 | f"Start coordinator failed with exit code {code}" |
| 70 | ) |
| 71 | try: |
| 72 | self._stub.HeartBeat(request) |
| 73 | logger.info("GraphScope coordinator service connected.") |
| 74 | break |
| 75 | except grpc.RpcError as e: |
| 76 | # Cannot connect to coordinator for a short time is expected |
| 77 | # as the coordinator takes some time to launch |
| 78 | msg = f"code: {e.code().name}, details: {e.details()}" |
| 79 | if e.code() == grpc.StatusCode.DEADLINE_EXCEEDED: |
| 80 | logger.warning("Heart beat analytical engine failed, %s", msg) |
| 81 | if time.time() - begin_time >= timeout_seconds: |
| 82 | raise ConnectionError(f"Connect coordinator timeout, {msg}") |
| 83 | # refresh the channel incase the server became available |
| 84 | if e.code() == grpc.StatusCode.UNAVAILABLE: |
| 85 | self._stub = self._get_stub() |
| 86 | time.sleep(1) |
| 87 | |
| 88 | def connect(self, cleanup_instance=True, dangling_timeout_seconds=60): |
| 89 | return self._connect_session_impl( |
| 90 | cleanup_instance=cleanup_instance, |
| 91 | dangling_timeout_seconds=dangling_timeout_seconds, |
| 92 | ) |
| 93 | |
| 94 | @property |
| 95 | def session_id(self): |
| 96 | return self._session_id |
| 97 | |