Skip to content

Commit

Permalink
address comments
Browse files Browse the repository at this point in the history
  • Loading branch information
azliu0 committed Jan 29, 2025
1 parent b57e260 commit d3b2f1c
Show file tree
Hide file tree
Showing 3 changed files with 20 additions and 22 deletions.
33 changes: 18 additions & 15 deletions modal/_utils/shell_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,13 @@
import asyncio
import contextlib
import errno
import fcntl
import os
import select
import shutil
import signal
import struct
import sys
import termios
import threading
from collections.abc import Coroutine
from queue import Empty, Queue
from types import FrameType
from typing import Callable, Optional

Expand Down Expand Up @@ -90,28 +87,36 @@ class WindowSizeHandler:
"""Handles terminal window resize events."""

def __init__(self):
"""Initialize window size handler. Must be called from the main thread to set signals properly.
"""Initialize window size handler. It is a system limitation that we can only set signals from the main thread.
In case this is invoked from a thread that is not the main thread, e.g. in tests, the context manager
becomes a no-op."""
becomes a no-op.
"""
self._is_main_thread = threading.current_thread() is threading.main_thread()
self._event_queue: Queue[tuple[int, int]] = Queue()

self._current_size: tuple[int, int] = self._get_terminal_size()
self._size_changed = asyncio.Event()

if self._is_main_thread and hasattr(signal, "SIGWINCH"):
signal.signal(signal.SIGWINCH, self._queue_resize_event)

def _get_terminal_size(self) -> tuple[int, int]:
c, r = shutil.get_terminal_size()
return r, c

def _queue_resize_event(self, signum: Optional[int] = None, frame: Optional[FrameType] = None) -> None:
"""Signal handler for SIGWINCH that queues events."""
try:
hw = struct.unpack("hh", fcntl.ioctl(sys.stdout.fileno(), termios.TIOCGWINSZ, b"1234"))
rows, cols = hw
self._event_queue.put((rows, cols))
self._current_size = self._get_terminal_size()
self._size_changed.set()
except Exception:
# ignore failed window size reads
pass

@contextlib.asynccontextmanager
async def watch_window_size(self, handler: Callable[[int, int], Coroutine]):
"""Context manager that processes window resize events from the queue.
Can be run from any thread. If the window manager was initialized from a thread that is not the main thread,
e.g. in tests, this context manager is a no-op.
"""
Expand All @@ -121,11 +126,9 @@ async def watch_window_size(self, handler: Callable[[int, int], Coroutine]):

async def process_events():
while True:
try:
rows, cols = self._event_queue.get_nowait()
await handler(rows, cols)
except Empty:
await asyncio.sleep(0.1)
await self._size_changed.wait()
await handler(*self._current_size)
self._size_changed.clear()

event_task = asyncio.create_task(process_events())
try:
Expand Down
6 changes: 0 additions & 6 deletions modal/container_process.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
# Copyright Modal Labs 2024
import asyncio
import platform
import struct
from typing import Generic, Optional, TypeVar

from modal_proto import api_pb2
Expand Down Expand Up @@ -153,11 +152,6 @@ async def _handle_input(data: bytes, message_index: int):
await self.stdin.drain()

async def _send_window_resize(rows: int, cols: int):
# resize sequence:
# - 2 bytes for the number of rows (big-endian)
# - 2 bytes for the number of columns (big-endian)
dims = struct.pack(">HH", rows, cols)
self.stdin.write(dims)
await self.stdin.drain(_terminal_size=(rows, cols))

async with TaskContext() as tc:
Expand Down
3 changes: 2 additions & 1 deletion modal/io_streams.py
Original file line number Diff line number Diff line change
Expand Up @@ -386,7 +386,8 @@ def write_eof(self) -> None:

async def drain(
self,
_terminal_size: Optional[tuple[int, int]] = None, # internal option to send terminal window resize events
# internal option to send terminal resize events to container exec
_terminal_size: Optional[tuple[int, int]] = None,
) -> None:
"""Flush the write buffer and send data to the running process.
Expand Down

0 comments on commit d3b2f1c

Please sign in to comment.