From b57e260b5400439e7e83f65cb25eca560bddea33 Mon Sep 17 00:00:00 2001 From: Andrew Liu Date: Wed, 29 Jan 2025 05:53:13 +0000 Subject: [PATCH] use window_rows and window_cols --- modal/container_process.py | 2 +- modal/io_streams.py | 8 +++----- 2 files changed, 4 insertions(+), 6 deletions(-) diff --git a/modal/container_process.py b/modal/container_process.py index 46a1e2df70..ab575f56e4 100644 --- a/modal/container_process.py +++ b/modal/container_process.py @@ -158,7 +158,7 @@ async def _send_window_resize(rows: int, cols: int): # - 2 bytes for the number of columns (big-endian) dims = struct.pack(">HH", rows, cols) self.stdin.write(dims) - await self.stdin.drain(_is_resize=True) + await self.stdin.drain(_terminal_size=(rows, cols)) async with TaskContext() as tc: stdout_task = tc.create_task(_write_to_fd_loop(self.stdout)) diff --git a/modal/io_streams.py b/modal/io_streams.py index d076e906b5..d4dca8064c 100644 --- a/modal/io_streams.py +++ b/modal/io_streams.py @@ -386,7 +386,7 @@ def write_eof(self) -> None: async def drain( self, - _is_resize: bool = False, # internal option to send terminal window resize events + _terminal_size: Optional[tuple[int, int]] = None, # internal option to send terminal window resize events ) -> None: """Flush the write buffer and send data to the running process. @@ -407,9 +407,6 @@ async def drain( ``` """ data = bytes(self._buffer) - if _is_resize: - assert len(data) == 4, "unexpected buffer size for resize event" # rrww - assert self._object_type == "container_process", "resize event can only be sent for container processes" self._buffer.clear() index = self._get_next_index() @@ -431,7 +428,8 @@ async def drain( message=data, message_index=index, eof=self._is_closed, - is_resize=_is_resize, + window_rows=_terminal_size[0] if _terminal_size else None, + window_cols=_terminal_size[1] if _terminal_size else None, ), ), )