def test_returns_ack_message(self, pybio_model_bytes, grpc_stub):
model = grpc_stub.CreateModelSession(
valid_model_request(pybio_model_bytes))
resp = grpc_stub.GetLogs(inference_pb2.Empty())
record = next(resp)
assert inference_pb2.LogEntry.Level.INFO == record.level
assert "Sending model logs" == record.content
grpc_stub.CloseModelSession(model)
def ping(self):
try:
with grpc.insecure_channel(self.__conn_str) as chan:
client = inference_pb2_grpc.FlightControlStub(chan)
client.Ping(inference_pb2.Empty())
return True
except Exception as e:
return False
def _fetch_devices(config: types.ServerConfig):
try:
port = config.port
if config.autostart:
# in order not to block address for real server todo: remove port hack
port = str(int(config.port) - 20)
addr = socket.gethostbyname(config.address)
conn_conf = ConnConf(addr, port, timeout=10)
if config.autostart:
if addr == "127.0.0.1":
launcher = LocalServerLauncher(conn_conf, path=config.path)
else:
launcher = RemoteSSHServerLauncher(
conn_conf, cred=SSHCred(user=config.username, key_path=config.ssh_key), path=config.path
)
else:
launcher = _NullLauncher()
try:
launcher.start()
with grpc.insecure_channel(f"{addr}:{port}") as chan:
client = inference_pb2_grpc.InferenceStub(chan)
resp = client.ListDevices(inference_pb2.Empty())
return [(d.id, d.id) for d in resp.devices]
except Exception as e:
logger.exception('Failed to fetch devices')
raise
finally:
try:
launcher.stop()
except Exception:
pass
except Exception as e:
logger.error(e)
raise
return []
def shutdown(self):
with grpc.insecure_channel(self.__conn_str) as chan:
client = inference_pb2_grpc.FlightControlStub(chan)
client.Shutdown(inference_pb2.Empty())
return True
def CloseModelSession(self, request: inference_pb2.ModelSession,
context) -> inference_pb2.Empty:
self.__session_manager.close_session(request.id)
return inference_pb2.Empty()
def Shutdown(self, request: inference_pb2.Empty, context):
if self.__done_evt:
self.__done_evt.set()
return inference_pb2.Empty()
def Ping(self, request: inference_pb2.Empty, context):
return inference_pb2.Empty()
def _query_devices(self, grpc_stub):
dev_resp = grpc_stub.ListDevices(inference_pb2.Empty())
device_by_id = {d.id: d for d in dev_resp.devices}
return device_by_id
def test_list_devices(self, grpc_stub):
resp = grpc_stub.ListDevices(inference_pb2.Empty())
device_by_id = {d.id: d for d in resp.devices}
assert "cpu" in device_by_id
assert inference_pb2.Device.Status.AVAILABLE == device_by_id[
"cpu"].status