Skip to content

Commit

Permalink
[pre-commit.ci] auto fixes from pre-commit.com hooks
Browse files Browse the repository at this point in the history
for more information, see https://pre-commit.ci
  • Loading branch information
pre-commit-ci[bot] committed Nov 10, 2022
1 parent 52908f1 commit 286528d
Show file tree
Hide file tree
Showing 2 changed files with 15 additions and 5 deletions.
12 changes: 9 additions & 3 deletions plugins/kernels/fps_kernels/kernel_driver/driver.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,9 @@ async def restart(self, startup_timeout: float = float("inf")) -> None:
msg = create_message("shutdown_request", content={"restart": True})
await send_message(msg, self.control_channel, self.key, change_date_to_str=True)
while True:
msg = cast(Dict[str, Any], await receive_message(self.control_channel, change_str_to_date=True))
msg = cast(
Dict[str, Any], await receive_message(self.control_channel, change_str_to_date=True)
)
if msg["msg_type"] == "shutdown_reply" and msg["content"]["restart"]:
break
await self._wait_for_ready(startup_timeout)
Expand Down Expand Up @@ -162,12 +164,16 @@ async def _wait_for_ready(self, timeout):
)
self.msg_cnt += 1
await send_message(msg, self.shell_channel, self.key, change_date_to_str=True)
msg = await receive_message(self.shell_channel, timeout=new_timeout, change_str_to_date=True)
msg = await receive_message(
self.shell_channel, timeout=new_timeout, change_str_to_date=True
)
if msg is None:
error_message = f"Kernel didn't respond in {timeout} seconds"
raise RuntimeError(error_message)
if msg["msg_type"] == "kernel_info_reply":
msg = await receive_message(self.iopub_channel, timeout=0.2, change_str_to_date=True)
msg = await receive_message(
self.iopub_channel, timeout=0.2, change_str_to_date=True
)
if msg is not None:
break
new_timeout = deadline_to_timeout(deadline)
Expand Down
8 changes: 6 additions & 2 deletions plugins/kernels/fps_kernels/kernel_driver/message.py
Original file line number Diff line number Diff line change
Expand Up @@ -119,11 +119,15 @@ def deserialize(
return message


async def send_message(msg: Dict[str, Any], sock: Socket, key: str, change_date_to_str: bool = False) -> None:
async def send_message(
msg: Dict[str, Any], sock: Socket, key: str, change_date_to_str: bool = False
) -> None:
await sock.send_multipart(serialize(msg, key, change_date_to_str=change_date_to_str), copy=True)


async def receive_message(sock: Socket, timeout: float = float("inf"), change_str_to_date: bool = False) -> Optional[Dict[str, Any]]:
async def receive_message(
sock: Socket, timeout: float = float("inf"), change_str_to_date: bool = False
) -> Optional[Dict[str, Any]]:
timeout *= 1000 # in ms
ready = await sock.poll(timeout)
if ready:
Expand Down

0 comments on commit 286528d

Please sign in to comment.