diff --git a/remote/modal_app.py b/remote/modal_app.py index 6482d30..be0f82f 100644 --- a/remote/modal_app.py +++ b/remote/modal_app.py @@ -12,9 +12,7 @@ import json import logging import os import queue -import select import shutil -import socket import statistics import subprocess import sys @@ -115,6 +113,19 @@ from .modal_billing import ( _prune_modal_hourly_billing_cache, get_hourly_modal_app_billing, ) +from .modal_container_logs import ( + MODAL_GPU_ESTIMATED_USD_PER_SECOND, + MODAL_GPU_PRICING_EFFECTIVE_DATE, + ModalContainerStatus, + ModalContainerStatusError, + _coerce_modal_task_id, + _is_remote_container_log_stream_enabled, + _release_remote_container_log_stream, + _retain_remote_container_log_stream, + _stop_modal_task_synchronously, + list_active_modal_containers, + stop_managed_modal_container, +) logger = logging.getLogger(__name__) @@ -163,12 +174,6 @@ _ROOT_LOADER_PREWARM_CLASS_TYPES = frozenset( ) _SPECULATIVE_PREWARM_TARGET_KEY = "speculative_remote_prewarm_target" -# Remote container log and billing state. -_REMOTE_CONTAINER_LOG_STREAMS_LOCK = threading.Lock() -_REMOTE_CONTAINER_LOG_STREAMS: dict[str, "_RemoteContainerLogStreamState"] = {} -_REMOTE_CONTAINER_LOG_STDERR_LOCK = threading.Lock() -_REMOTE_CONTAINER_LOG_STREAM_IDLE_GRACE_SECONDS = 30.0 - # Active invocation and local durable-session state. _IMPLICIT_BATCH_PRESERVING_TARGETS = frozenset({("CreateVideo", "images")}) _ACTIVE_REMOTE_INVOCATIONS_LOCK = threading.Lock() @@ -203,26 +208,6 @@ _DURABLE_BRIDGE_REHYDRATION_IO_TYPES = frozenset( {"CLIP", "MODEL", "NOISE", "SAMPLER", "VAE"} ) -# Published Modal GPU pricing metadata. -MODAL_GPU_PRICING_EFFECTIVE_DATE = "2026-08-13" -MODAL_GPU_ESTIMATED_USD_PER_SECOND: dict[str, float] = { - "T4": 0.000164, - "L4": 0.000222, - "A10": 0.000306, - "L40S": 0.000542, - "A100": 0.000583, - "A100-40GB": 0.000583, - "A100-80GB": 0.000694, - "RTX-PRO-6000": 0.000842, - "H100": 0.001097, - "H100!": 0.001097, - "H200": 0.001261, - "B200": 0.001736, - "B200+": 0.001736, - "B300": 0.001972, -} - - def _remote_modal_call_worker_count() -> int: """Return the number of local worker threads reserved for blocking Modal calls.""" return max(1, int(get_settings().max_inflight_calls)) @@ -309,23 +294,8 @@ class _ModalAutoDeployState: last_error: BaseException | None = None -@dataclass -class _RemoteContainerLogLineBuffer: - """Buffer partial remote log lines so stderr mirroring stays line-oriented.""" - - task_id: str - buffered_text: str = "" -@dataclass -class _RemoteContainerLogStreamState: - """Track one active local watcher mirroring logs for a Modal container.""" - - task_id: str - stop_event: threading.Event - thread: threading.Thread - refcount: int = 0 - idle_stop_timer: threading.Timer | None = None @dataclass @@ -372,586 +342,60 @@ class ModalRemoteAppOutOfDateError(ModalRemoteInvocationError): """Raised when a deployed Modal app is incompatible with the local client.""" -class ModalContainerStatusError(RuntimeError): - """Raised when active Modal container status cannot be queried.""" - - - -@dataclass(frozen=True) -class ModalContainerStatus: - """Describe one active Modal container owned by this ComfyUI instance.""" - - container_id: str - app_id: str - app_name: str - modal_gpu: str - estimated_gpu_cost_per_second: float - state: str - enqueued_at: float | None - started_at: float | None - - def as_dict(self) -> dict[str, str | float | None]: - """Return a JSON-serializable representation for the frontend.""" - return { - "container_id": self.container_id, - "app_id": self.app_id, - "app_name": self.app_name, - "modal_gpu": self.modal_gpu, - "estimated_gpu_cost_per_second": self.estimated_gpu_cost_per_second, - "state": self.state, - "enqueued_at": self.enqueued_at, - "started_at": self.started_at, - } - -def _is_remote_container_log_stream_enabled() -> bool: - """Return whether remote Modal container logs should be mirrored locally.""" - return bool(get_settings().stream_remote_container_logs) - - -def _coerce_modal_task_id(value: Any) -> str | None: - """Normalize one streamed Modal task id into a non-empty string.""" - if value is None: - return None - task_id = str(value).strip() - return task_id or None - - -def _write_remote_container_log_line( - task_id: str, - line: str, - *, - stream: Any = None, -) -> None: - """Write one complete remote container log line to the local stderr stream.""" - target_stream = stream if stream is not None else sys.stderr - with _REMOTE_CONTAINER_LOG_STDERR_LOCK: - target_stream.write(f"[modal:{task_id}] {line}") - target_stream.flush() - - -def _write_remote_container_log_chunk( - line_buffer: _RemoteContainerLogLineBuffer, - chunk: str, - *, - stream: Any = None, -) -> None: - """Buffer one remote log chunk and mirror complete lines to local stderr.""" - line_buffer.buffered_text += chunk - while True: - newline_index = line_buffer.buffered_text.find("\n") - if newline_index < 0: - return - next_line = line_buffer.buffered_text[: newline_index + 1] - line_buffer.buffered_text = line_buffer.buffered_text[newline_index + 1 :] - _write_remote_container_log_line(line_buffer.task_id, next_line, stream=stream) - - -def _flush_remote_container_log_chunk( - line_buffer: _RemoteContainerLogLineBuffer, - *, - stream: Any = None, -) -> None: - """Flush any partial remote log text remaining in the local line buffer.""" - if not line_buffer.buffered_text: - return - _write_remote_container_log_line( - line_buffer.task_id, - f"{line_buffer.buffered_text}\n", - stream=stream, - ) - line_buffer.buffered_text = "" - - -def _managed_modal_app_gpus(settings: ModalSyncSettings) -> dict[str, str]: - """Map every GPU-specific app name owned by this ComfyUI instance to its GPU.""" - return { - modal_deployment_app_name( - settings_for_modal_gpu(settings, modal_gpu) - ): modal_gpu - for modal_gpu in MODAL_GPU_TYPES - } - - -def _optional_modal_timestamp(value: Any) -> float | None: - """Normalize an unset Modal protobuf timestamp to ``None``.""" - timestamp = float(value or 0.0) - return timestamp if timestamp > 0 else None - - - -def _resolved_modal_environment_name() -> str | None: - """Return the explicitly selected Modal environment, if one is configured.""" - object_module = importlib.import_module("modal._object") - environments_module = importlib.import_module("modal.environments") - environment = environments_module.ensure_env(_modal_environment_name()) - environment_name = object_module._get_environment_name(environment) - if environment_name is None: - return None - normalized_environment_name = str(environment_name).strip() - return normalized_environment_name or None - - -@lru_cache(maxsize=None) -def _synchronized_modal_callable( - async_callable: Callable[..., Any], -) -> Callable[..., Any]: - """Wrap one async Modal SDK operation on Modal's managed event loop.""" - async_utils_module = importlib.import_module("modal._utils.async_utils") - return async_utils_module.synchronize_api(async_callable) - - -async def _list_modal_tasks_on_sdk_loop( - client_module: Any, - api_pb2: Any, - environment_name: str | None, -) -> Any: - """List Modal tasks while running on Modal's managed event loop.""" - client = await client_module._Client.from_env() - return await client.stub.TaskList( - api_pb2.TaskListRequest(environment_name=environment_name) - ) - - -def _list_modal_tasks_synchronously( - client_module: Any, - api_pb2: Any, - environment_name: str | None, -) -> Any: - """Run one Modal task-list query through the synchronized SDK bridge.""" - blocking_callable = _synchronized_modal_callable(_list_modal_tasks_on_sdk_loop) - return blocking_callable(client_module, api_pb2, environment_name) - - -async def _stop_modal_task_on_sdk_loop( - client_module: Any, - api_pb2: Any, - container_id: str, -) -> Any: - """Stop one exact Modal task while running on Modal's managed event loop.""" - client = await client_module._Client.from_env() - return await client.stub.ContainerStop( - api_pb2.ContainerStopRequest(task_id=container_id) - ) - - -def _stop_modal_task_synchronously( - client_module: Any, - api_pb2: Any, - container_id: str, -) -> Any: - """Run one Modal task-stop request through the synchronized SDK bridge.""" - blocking_callable = _synchronized_modal_callable(_stop_modal_task_on_sdk_loop) - return blocking_callable(client_module, api_pb2, container_id) - - - -async def list_active_modal_containers( - settings: ModalSyncSettings | None = None, -) -> list[ModalContainerStatus]: - """List active Modal containers belonging to this ComfyUI instance.""" - if modal is None: - raise ModalContainerStatusError("The Modal SDK is unavailable.") - - resolved_settings = settings or get_settings() - managed_app_gpus = _managed_modal_app_gpus(resolved_settings) - try: - client_module = importlib.import_module("modal.client") - exception_module = importlib.import_module("modal.exception") - api_pb2 = importlib.import_module("modal_proto.api_pb2") - except ModuleNotFoundError as exc: - raise ModalContainerStatusError( - "The installed Modal SDK does not expose the container list API." - ) from exc - - modal_error_type = getattr(exception_module, "Error", RuntimeError) - try: - environment_name = _resolved_modal_environment_name() - response = await asyncio.to_thread( - _list_modal_tasks_synchronously, - client_module, - api_pb2, - environment_name, - ) - except (modal_error_type, OSError, AttributeError, RuntimeError) as exc: - raise ModalContainerStatusError( - f"Unable to list Modal containers: {exc}" - ) from exc - - containers: list[ModalContainerStatus] = [] - for task in response.tasks: - app_name = str(task.app_description) - modal_gpu = managed_app_gpus.get(app_name) - if modal_gpu is None: - continue - started_at = _optional_modal_timestamp(task.started_at) - containers.append( - ModalContainerStatus( - container_id=str(task.task_id), - app_id=str(task.app_id), - app_name=app_name, - modal_gpu=modal_gpu, - estimated_gpu_cost_per_second=MODAL_GPU_ESTIMATED_USD_PER_SECOND[ - modal_gpu - ], - state="running" if started_at is not None else "starting", - enqueued_at=_optional_modal_timestamp(task.enqueued_at), - started_at=started_at, - ) - ) - - containers.sort( - key=lambda container: ( - container.enqueued_at or container.started_at or 0.0, - container.container_id, - ) - ) - logger.debug( - "Listed %d active Modal container(s) for the status UI.", len(containers) - ) - return containers - - -async def stop_managed_modal_container( - container_id: str, - settings: ModalSyncSettings | None = None, -) -> bool: - """Stop one active container after verifying it belongs to this installation.""" - normalized_container_id = container_id.strip() - if not normalized_container_id: - raise ValueError("Modal container_id must not be empty.") - - active_containers = await list_active_modal_containers(settings) - if not any( - container.container_id == normalized_container_id - for container in active_containers - ): - return False - - try: - client_module = importlib.import_module("modal.client") - exception_module = importlib.import_module("modal.exception") - api_pb2 = importlib.import_module("modal_proto.api_pb2") - except ModuleNotFoundError as exc: - raise ModalContainerStatusError( - "The installed Modal SDK does not expose the container stop API." - ) from exc - - modal_error_type = getattr(exception_module, "Error", RuntimeError) - try: - await asyncio.to_thread( - _stop_modal_task_synchronously, - client_module, - api_pb2, - normalized_container_id, - ) - except (modal_error_type, OSError, AttributeError, RuntimeError) as exc: - raise ModalContainerStatusError( - f"Unable to stop Modal container {normalized_container_id}: {exc}" - ) from exc - logger.warning("Stopped managed Modal container %s.", normalized_container_id) - return True - - -async def _stream_remote_container_logs_via_modal_sdk_async( - task_id: str, - stop_event: threading.Event, -) -> None: - """Follow one Modal container log stream through the Python SDK internals.""" - if modal is None: - raise ModuleNotFoundError("Modal SDK is unavailable.") - - client_module = importlib.import_module("modal.client") - exception_module = importlib.import_module("modal.exception") - api_pb2 = importlib.import_module("modal_proto.api_pb2") - grpclib_exceptions = importlib.import_module("grpclib.exceptions") - client = await client_module._Client.from_env() - line_buffer = _RemoteContainerLogLineBuffer(task_id=task_id) - last_entry_id = "" - - try: - while not stop_event.is_set(): - request = api_pb2.AppGetLogsRequest( - task_id=task_id, - timeout=5, - last_entry_id=last_entry_id, - ) - try: - async for log_batch in client.stub.AppGetLogs.unary_stream(request): - if stop_event.is_set(): - break - if log_batch.entry_id: - last_entry_id = str(log_batch.entry_id) - if bool(log_batch.app_done): - logger.info( - "Modal SDK log stream finished for task_id=%s.", task_id - ) - return - for log_item in log_batch.items: - log_data = getattr(log_item, "data", "") - if not log_data: - continue - _write_remote_container_log_chunk(line_buffer, str(log_data)) - except ( - exception_module.ServiceError, - exception_module.InternalError, - grpclib_exceptions.StreamTerminatedError, - socket.gaierror, - ) as exc: - if stop_event.is_set(): - break - logger.warning( - "Retrying Modal SDK log stream for task_id=%s after transient failure: %s", - task_id, - exc, - ) - continue - except AttributeError as exc: - if stop_event.is_set(): - break - if "_write_appdata" in str(exc): - logger.warning( - "Retrying Modal SDK log stream for task_id=%s after connection loss: %s", - task_id, - exc, - ) - continue - raise - finally: - _flush_remote_container_log_chunk(line_buffer) - - -def _stream_remote_container_logs_via_modal_sdk( - task_id: str, - stop_event: threading.Event, -) -> bool: - """Try to mirror one Modal container log stream through the installed SDK.""" - if modal is None: - return False - blocking_callable = _synchronized_modal_callable( - _stream_remote_container_logs_via_modal_sdk_async - ) - blocking_callable(task_id, stop_event) - return True - - -def _stream_remote_container_logs_via_modal_cli( - task_id: str, - stop_event: threading.Event, -) -> bool: - """Try to mirror one Modal container log stream through the Modal CLI.""" - modal_cli = shutil.which("modal") - if modal_cli is None: - return False - - command = [modal_cli, "container", "logs", task_id, "-f"] - line_buffer = _RemoteContainerLogLineBuffer(task_id=task_id) - process = subprocess.Popen( - command, - stdout=subprocess.PIPE, - stderr=subprocess.STDOUT, - bufsize=0, - ) - line_buffer = _RemoteContainerLogLineBuffer(task_id=task_id) - try: - stdout_stream = process.stdout - if stdout_stream is None: - raise RuntimeError("Modal CLI log process did not expose a stdout stream.") - while not stop_event.is_set(): - ready_streams, _, _ = select.select([stdout_stream], [], [], 0.25) - if ready_streams: - next_chunk = stdout_stream.read(4096) - if next_chunk: - _write_remote_container_log_chunk( - line_buffer, - next_chunk.decode("utf-8", errors="replace"), - ) - continue - if process.poll() is not None: - trailing_chunk = stdout_stream.read() - if trailing_chunk: - _write_remote_container_log_chunk( - line_buffer, - trailing_chunk.decode("utf-8", errors="replace"), - ) - break - if stop_event.is_set() and process.poll() is None: - process.terminate() - finally: - if process.poll() is None: - try: - process.wait(timeout=2.0) - except subprocess.TimeoutExpired: - process.kill() - process.wait(timeout=2.0) - _flush_remote_container_log_chunk(line_buffer) - - if process.returncode not in {0, None} and not stop_event.is_set(): - raise RuntimeError( - f"Modal CLI exited with status {process.returncode} while streaming logs for task_id={task_id}." - ) - return True - - -def _run_remote_container_log_stream(task_id: str, stop_event: threading.Event) -> None: - """Run the best available remote log streaming backend for one Modal container.""" - logger.info("Starting remote Modal container log stream for task_id=%s.", task_id) - try: - if _stream_remote_container_logs_via_modal_cli(task_id, stop_event): - logger.info( - "Stopped remote Modal container log stream for task_id=%s.", task_id - ) - return - except (OSError, RuntimeError, subprocess.SubprocessError, ValueError) as exc: - logger.warning( - "Falling back from Modal CLI log streaming for task_id=%s after failure: %s", - task_id, - exc, - ) - - try: - if _stream_remote_container_logs_via_modal_sdk(task_id, stop_event): - logger.info( - "Stopped remote Modal container log stream for task_id=%s.", task_id - ) - return - except ( - AttributeError, - ImportError, - ModuleNotFoundError, - OSError, - RuntimeError, - ValueError, - ) as exc: - logger.warning( - "Modal SDK log streaming failed for task_id=%s: %s", - task_id, - exc, - ) - - logger.warning( - "Unable to mirror remote Modal container logs for task_id=%s because neither the Modal SDK nor CLI is available locally.", - task_id, - ) - - -def _new_remote_container_log_stream_state( - task_id: str, -) -> _RemoteContainerLogStreamState: - """Create an unstarted container-log watcher state.""" - stop_event = threading.Event() - return _RemoteContainerLogStreamState( - task_id=task_id, - stop_event=stop_event, - thread=threading.Thread( - target=_run_remote_container_log_stream, - args=(task_id, stop_event), - name=f"modal-log-stream-{task_id}", - daemon=True, - ), - ) - - -def _cancel_remote_container_log_idle_stop( - stream_state: _RemoteContainerLogStreamState, -) -> None: - """Cancel one pending idle-stop timer before a container is reused.""" - idle_stop_timer = stream_state.idle_stop_timer - if idle_stop_timer is None: - return - idle_stop_timer.cancel() - stream_state.idle_stop_timer = None - - -def _stop_remote_container_log_stream_if_idle( - task_id: str, - expected_state: _RemoteContainerLogStreamState, -) -> None: - """Stop one watcher after its grace period if no payload reused it.""" - with _REMOTE_CONTAINER_LOG_STREAMS_LOCK: - stream_state = _REMOTE_CONTAINER_LOG_STREAMS.get(task_id) - if stream_state is not expected_state or stream_state.refcount != 0: - return - stream_state.idle_stop_timer = None - _REMOTE_CONTAINER_LOG_STREAMS.pop(task_id, None) - - logger.info( - "Stopping idle remote Modal container log stream for task_id=%s.", - task_id, - ) - stream_state.stop_event.set() - stream_state.thread.join(timeout=0.2) - - -def _schedule_remote_container_log_idle_stop( - stream_state: _RemoteContainerLogStreamState, -) -> None: - """Schedule one daemon timer to stop an unreferenced log watcher.""" - idle_stop_timer = threading.Timer( - _REMOTE_CONTAINER_LOG_STREAM_IDLE_GRACE_SECONDS, - _stop_remote_container_log_stream_if_idle, - args=(stream_state.task_id, stream_state), - ) - idle_stop_timer.daemon = True - stream_state.idle_stop_timer = idle_stop_timer - idle_stop_timer.start() - - -def _retain_remote_container_log_stream(task_id: str) -> str: - """Increment one shared Modal container log watcher and start it if needed.""" - with _REMOTE_CONTAINER_LOG_STREAMS_LOCK: - stream_state = _REMOTE_CONTAINER_LOG_STREAMS.get(task_id) - if stream_state is not None and not stream_state.thread.is_alive(): - _cancel_remote_container_log_idle_stop(stream_state) - _REMOTE_CONTAINER_LOG_STREAMS.pop(task_id, None) - stream_state = None - - if stream_state is None: - stream_state = _new_remote_container_log_stream_state(task_id) - _REMOTE_CONTAINER_LOG_STREAMS[task_id] = stream_state - logger.info( - "Creating remote Modal container log stream for task_id=%s.", task_id - ) - stream_state.thread.start() - else: - _cancel_remote_container_log_idle_stop(stream_state) - logger.info( - "Reusing remote Modal container log stream for task_id=%s.", task_id - ) - stream_state.refcount += 1 - logger.info( - "Remote Modal container log stream retain task_id=%s refcount=%d.", - task_id, - stream_state.refcount, - ) - - return task_id - - -def _release_remote_container_log_stream(task_id: str) -> None: - """Release one watcher while allowing prompt-local container reuse.""" - with _REMOTE_CONTAINER_LOG_STREAMS_LOCK: - stream_state = _REMOTE_CONTAINER_LOG_STREAMS.get(task_id) - if stream_state is None: - return - stream_state.refcount = max(0, stream_state.refcount - 1) - logger.info( - "Remote Modal container log stream release task_id=%s refcount=%d.", - task_id, - stream_state.refcount, - ) - if stream_state.refcount != 0: - return - if not stream_state.thread.is_alive(): - _REMOTE_CONTAINER_LOG_STREAMS.pop(task_id, None) - return - if stream_state.idle_stop_timer is None: - logger.info( - "Keeping remote Modal container log stream alive for %.1fs " - "to allow task reuse without replaying history task_id=%s.", - _REMOTE_CONTAINER_LOG_STREAM_IDLE_GRACE_SECONDS, - task_id, - ) - _schedule_remote_container_log_idle_stop(stream_state) + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + def _close_remote_payload_stream(stream_events: Iterator[dict[str, Any]]) -> None: diff --git a/remote/modal_container_logs.py b/remote/modal_container_logs.py new file mode 100644 index 0000000..d82ab21 --- /dev/null +++ b/remote/modal_container_logs.py @@ -0,0 +1,665 @@ +"""Managed Modal container inventory, termination, and log streaming.""" + +from __future__ import annotations + +import asyncio +from dataclasses import dataclass +from functools import lru_cache +import importlib +import logging +import os +import select +import shutil +import socket +import subprocess +import sys +import threading +from typing import Any, Callable + +from ..settings import ( + MODAL_GPU_TYPES, + ModalSyncSettings, + get_settings, + modal_deployment_app_name, + settings_for_modal_gpu, +) + +logger = logging.getLogger(__name__) + +try: + import modal # type: ignore +except ModuleNotFoundError: # pragma: no cover - local fallback environments. + modal = None + +_REMOTE_CONTAINER_LOG_STREAMS_LOCK = threading.Lock() +_REMOTE_CONTAINER_LOG_STREAMS: dict[str, "_RemoteContainerLogStreamState"] = {} +_REMOTE_CONTAINER_LOG_STDERR_LOCK = threading.Lock() +_REMOTE_CONTAINER_LOG_STREAM_IDLE_GRACE_SECONDS = 30.0 + +MODAL_GPU_PRICING_EFFECTIVE_DATE = "2026-08-13" +MODAL_GPU_ESTIMATED_USD_PER_SECOND: dict[str, float] = { + "T4": 0.000164, + "L4": 0.000222, + "A10": 0.000306, + "L40S": 0.000542, + "A100": 0.000583, + "A100-40GB": 0.000583, + "A100-80GB": 0.000694, + "RTX-PRO-6000": 0.000842, + "H100": 0.001097, + "H100!": 0.001097, + "H200": 0.001261, + "B200": 0.001736, + "B200+": 0.001736, + "B300": 0.001972, +} + + +def _modal_environment_name() -> str | None: + """Return the active Modal environment name when explicitly configured.""" + environment_name = os.getenv("MODAL_ENVIRONMENT") + if environment_name is None: + return None + normalized = environment_name.strip() + return normalized or None + + +@dataclass +class _RemoteContainerLogLineBuffer: + """Buffer partial remote log lines so stderr mirroring stays line-oriented.""" + + task_id: str + buffered_text: str = "" + + +@dataclass +class _RemoteContainerLogStreamState: + """Track one active local watcher mirroring logs for a Modal container.""" + + task_id: str + stop_event: threading.Event + thread: threading.Thread + refcount: int = 0 + idle_stop_timer: threading.Timer | None = None + + +class ModalContainerStatusError(RuntimeError): + """Raised when active Modal container status cannot be queried.""" + + + +@dataclass(frozen=True) +class ModalContainerStatus: + """Describe one active Modal container owned by this ComfyUI instance.""" + + container_id: str + app_id: str + app_name: str + modal_gpu: str + estimated_gpu_cost_per_second: float + state: str + enqueued_at: float | None + started_at: float | None + + def as_dict(self) -> dict[str, str | float | None]: + """Return a JSON-serializable representation for the frontend.""" + return { + "container_id": self.container_id, + "app_id": self.app_id, + "app_name": self.app_name, + "modal_gpu": self.modal_gpu, + "estimated_gpu_cost_per_second": self.estimated_gpu_cost_per_second, + "state": self.state, + "enqueued_at": self.enqueued_at, + "started_at": self.started_at, + } + +def _is_remote_container_log_stream_enabled() -> bool: + """Return whether remote Modal container logs should be mirrored locally.""" + return bool(get_settings().stream_remote_container_logs) + + +def _coerce_modal_task_id(value: Any) -> str | None: + """Normalize one streamed Modal task id into a non-empty string.""" + if value is None: + return None + task_id = str(value).strip() + return task_id or None + + +def _write_remote_container_log_line( + task_id: str, + line: str, + *, + stream: Any = None, +) -> None: + """Write one complete remote container log line to the local stderr stream.""" + target_stream = stream if stream is not None else sys.stderr + with _REMOTE_CONTAINER_LOG_STDERR_LOCK: + target_stream.write(f"[modal:{task_id}] {line}") + target_stream.flush() + + +def _write_remote_container_log_chunk( + line_buffer: _RemoteContainerLogLineBuffer, + chunk: str, + *, + stream: Any = None, +) -> None: + """Buffer one remote log chunk and mirror complete lines to local stderr.""" + line_buffer.buffered_text += chunk + while True: + newline_index = line_buffer.buffered_text.find("\n") + if newline_index < 0: + return + next_line = line_buffer.buffered_text[: newline_index + 1] + line_buffer.buffered_text = line_buffer.buffered_text[newline_index + 1 :] + _write_remote_container_log_line(line_buffer.task_id, next_line, stream=stream) + + +def _flush_remote_container_log_chunk( + line_buffer: _RemoteContainerLogLineBuffer, + *, + stream: Any = None, +) -> None: + """Flush any partial remote log text remaining in the local line buffer.""" + if not line_buffer.buffered_text: + return + _write_remote_container_log_line( + line_buffer.task_id, + f"{line_buffer.buffered_text}\n", + stream=stream, + ) + line_buffer.buffered_text = "" + + +def _managed_modal_app_gpus(settings: ModalSyncSettings) -> dict[str, str]: + """Map every GPU-specific app name owned by this ComfyUI instance to its GPU.""" + return { + modal_deployment_app_name( + settings_for_modal_gpu(settings, modal_gpu) + ): modal_gpu + for modal_gpu in MODAL_GPU_TYPES + } + + +def _optional_modal_timestamp(value: Any) -> float | None: + """Normalize an unset Modal protobuf timestamp to ``None``.""" + timestamp = float(value or 0.0) + return timestamp if timestamp > 0 else None + + + +def _resolved_modal_environment_name() -> str | None: + """Return the explicitly selected Modal environment, if one is configured.""" + object_module = importlib.import_module("modal._object") + environments_module = importlib.import_module("modal.environments") + environment = environments_module.ensure_env(_modal_environment_name()) + environment_name = object_module._get_environment_name(environment) + if environment_name is None: + return None + normalized_environment_name = str(environment_name).strip() + return normalized_environment_name or None + + +@lru_cache(maxsize=None) +def _synchronized_modal_callable( + async_callable: Callable[..., Any], +) -> Callable[..., Any]: + """Wrap one async Modal SDK operation on Modal's managed event loop.""" + async_utils_module = importlib.import_module("modal._utils.async_utils") + return async_utils_module.synchronize_api(async_callable) + + +async def _list_modal_tasks_on_sdk_loop( + client_module: Any, + api_pb2: Any, + environment_name: str | None, +) -> Any: + """List Modal tasks while running on Modal's managed event loop.""" + client = await client_module._Client.from_env() + return await client.stub.TaskList( + api_pb2.TaskListRequest(environment_name=environment_name) + ) + + +def _list_modal_tasks_synchronously( + client_module: Any, + api_pb2: Any, + environment_name: str | None, +) -> Any: + """Run one Modal task-list query through the synchronized SDK bridge.""" + blocking_callable = _synchronized_modal_callable(_list_modal_tasks_on_sdk_loop) + return blocking_callable(client_module, api_pb2, environment_name) + + +async def _stop_modal_task_on_sdk_loop( + client_module: Any, + api_pb2: Any, + container_id: str, +) -> Any: + """Stop one exact Modal task while running on Modal's managed event loop.""" + client = await client_module._Client.from_env() + return await client.stub.ContainerStop( + api_pb2.ContainerStopRequest(task_id=container_id) + ) + + +def _stop_modal_task_synchronously( + client_module: Any, + api_pb2: Any, + container_id: str, +) -> Any: + """Run one Modal task-stop request through the synchronized SDK bridge.""" + blocking_callable = _synchronized_modal_callable(_stop_modal_task_on_sdk_loop) + return blocking_callable(client_module, api_pb2, container_id) + + + +async def list_active_modal_containers( + settings: ModalSyncSettings | None = None, +) -> list[ModalContainerStatus]: + """List active Modal containers belonging to this ComfyUI instance.""" + if modal is None: + raise ModalContainerStatusError("The Modal SDK is unavailable.") + + resolved_settings = settings or get_settings() + managed_app_gpus = _managed_modal_app_gpus(resolved_settings) + try: + client_module = importlib.import_module("modal.client") + exception_module = importlib.import_module("modal.exception") + api_pb2 = importlib.import_module("modal_proto.api_pb2") + except ModuleNotFoundError as exc: + raise ModalContainerStatusError( + "The installed Modal SDK does not expose the container list API." + ) from exc + + modal_error_type = getattr(exception_module, "Error", RuntimeError) + try: + environment_name = _resolved_modal_environment_name() + response = await asyncio.to_thread( + _list_modal_tasks_synchronously, + client_module, + api_pb2, + environment_name, + ) + except (modal_error_type, OSError, AttributeError, RuntimeError) as exc: + raise ModalContainerStatusError( + f"Unable to list Modal containers: {exc}" + ) from exc + + containers: list[ModalContainerStatus] = [] + for task in response.tasks: + app_name = str(task.app_description) + modal_gpu = managed_app_gpus.get(app_name) + if modal_gpu is None: + continue + started_at = _optional_modal_timestamp(task.started_at) + containers.append( + ModalContainerStatus( + container_id=str(task.task_id), + app_id=str(task.app_id), + app_name=app_name, + modal_gpu=modal_gpu, + estimated_gpu_cost_per_second=MODAL_GPU_ESTIMATED_USD_PER_SECOND[ + modal_gpu + ], + state="running" if started_at is not None else "starting", + enqueued_at=_optional_modal_timestamp(task.enqueued_at), + started_at=started_at, + ) + ) + + containers.sort( + key=lambda container: ( + container.enqueued_at or container.started_at or 0.0, + container.container_id, + ) + ) + logger.debug( + "Listed %d active Modal container(s) for the status UI.", len(containers) + ) + return containers + + +async def stop_managed_modal_container( + container_id: str, + settings: ModalSyncSettings | None = None, +) -> bool: + """Stop one active container after verifying it belongs to this installation.""" + normalized_container_id = container_id.strip() + if not normalized_container_id: + raise ValueError("Modal container_id must not be empty.") + + active_containers = await list_active_modal_containers(settings) + if not any( + container.container_id == normalized_container_id + for container in active_containers + ): + return False + + try: + client_module = importlib.import_module("modal.client") + exception_module = importlib.import_module("modal.exception") + api_pb2 = importlib.import_module("modal_proto.api_pb2") + except ModuleNotFoundError as exc: + raise ModalContainerStatusError( + "The installed Modal SDK does not expose the container stop API." + ) from exc + + modal_error_type = getattr(exception_module, "Error", RuntimeError) + try: + await asyncio.to_thread( + _stop_modal_task_synchronously, + client_module, + api_pb2, + normalized_container_id, + ) + except (modal_error_type, OSError, AttributeError, RuntimeError) as exc: + raise ModalContainerStatusError( + f"Unable to stop Modal container {normalized_container_id}: {exc}" + ) from exc + logger.warning("Stopped managed Modal container %s.", normalized_container_id) + return True + + +async def _stream_remote_container_logs_via_modal_sdk_async( + task_id: str, + stop_event: threading.Event, +) -> None: + """Follow one Modal container log stream through the Python SDK internals.""" + if modal is None: + raise ModuleNotFoundError("Modal SDK is unavailable.") + + client_module = importlib.import_module("modal.client") + exception_module = importlib.import_module("modal.exception") + api_pb2 = importlib.import_module("modal_proto.api_pb2") + grpclib_exceptions = importlib.import_module("grpclib.exceptions") + client = await client_module._Client.from_env() + line_buffer = _RemoteContainerLogLineBuffer(task_id=task_id) + last_entry_id = "" + + try: + while not stop_event.is_set(): + request = api_pb2.AppGetLogsRequest( + task_id=task_id, + timeout=5, + last_entry_id=last_entry_id, + ) + try: + async for log_batch in client.stub.AppGetLogs.unary_stream(request): + if stop_event.is_set(): + break + if log_batch.entry_id: + last_entry_id = str(log_batch.entry_id) + if bool(log_batch.app_done): + logger.info( + "Modal SDK log stream finished for task_id=%s.", task_id + ) + return + for log_item in log_batch.items: + log_data = getattr(log_item, "data", "") + if not log_data: + continue + _write_remote_container_log_chunk(line_buffer, str(log_data)) + except ( + exception_module.ServiceError, + exception_module.InternalError, + grpclib_exceptions.StreamTerminatedError, + socket.gaierror, + ) as exc: + if stop_event.is_set(): + break + logger.warning( + "Retrying Modal SDK log stream for task_id=%s after transient failure: %s", + task_id, + exc, + ) + continue + except AttributeError as exc: + if stop_event.is_set(): + break + if "_write_appdata" in str(exc): + logger.warning( + "Retrying Modal SDK log stream for task_id=%s after connection loss: %s", + task_id, + exc, + ) + continue + raise + finally: + _flush_remote_container_log_chunk(line_buffer) + + +def _stream_remote_container_logs_via_modal_sdk( + task_id: str, + stop_event: threading.Event, +) -> bool: + """Try to mirror one Modal container log stream through the installed SDK.""" + if modal is None: + return False + blocking_callable = _synchronized_modal_callable( + _stream_remote_container_logs_via_modal_sdk_async + ) + blocking_callable(task_id, stop_event) + return True + + +def _stream_remote_container_logs_via_modal_cli( + task_id: str, + stop_event: threading.Event, +) -> bool: + """Try to mirror one Modal container log stream through the Modal CLI.""" + modal_cli = shutil.which("modal") + if modal_cli is None: + return False + + command = [modal_cli, "container", "logs", task_id, "-f"] + line_buffer = _RemoteContainerLogLineBuffer(task_id=task_id) + process = subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + bufsize=0, + ) + line_buffer = _RemoteContainerLogLineBuffer(task_id=task_id) + try: + stdout_stream = process.stdout + if stdout_stream is None: + raise RuntimeError("Modal CLI log process did not expose a stdout stream.") + while not stop_event.is_set(): + ready_streams, _, _ = select.select([stdout_stream], [], [], 0.25) + if ready_streams: + next_chunk = stdout_stream.read(4096) + if next_chunk: + _write_remote_container_log_chunk( + line_buffer, + next_chunk.decode("utf-8", errors="replace"), + ) + continue + if process.poll() is not None: + trailing_chunk = stdout_stream.read() + if trailing_chunk: + _write_remote_container_log_chunk( + line_buffer, + trailing_chunk.decode("utf-8", errors="replace"), + ) + break + if stop_event.is_set() and process.poll() is None: + process.terminate() + finally: + if process.poll() is None: + try: + process.wait(timeout=2.0) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=2.0) + _flush_remote_container_log_chunk(line_buffer) + + if process.returncode not in {0, None} and not stop_event.is_set(): + raise RuntimeError( + f"Modal CLI exited with status {process.returncode} while streaming logs for task_id={task_id}." + ) + return True + + +def _run_remote_container_log_stream(task_id: str, stop_event: threading.Event) -> None: + """Run the best available remote log streaming backend for one Modal container.""" + logger.info("Starting remote Modal container log stream for task_id=%s.", task_id) + try: + if _stream_remote_container_logs_via_modal_cli(task_id, stop_event): + logger.info( + "Stopped remote Modal container log stream for task_id=%s.", task_id + ) + return + except (OSError, RuntimeError, subprocess.SubprocessError, ValueError) as exc: + logger.warning( + "Falling back from Modal CLI log streaming for task_id=%s after failure: %s", + task_id, + exc, + ) + + try: + if _stream_remote_container_logs_via_modal_sdk(task_id, stop_event): + logger.info( + "Stopped remote Modal container log stream for task_id=%s.", task_id + ) + return + except ( + AttributeError, + ImportError, + ModuleNotFoundError, + OSError, + RuntimeError, + ValueError, + ) as exc: + logger.warning( + "Modal SDK log streaming failed for task_id=%s: %s", + task_id, + exc, + ) + + logger.warning( + "Unable to mirror remote Modal container logs for task_id=%s because neither the Modal SDK nor CLI is available locally.", + task_id, + ) + + +def _new_remote_container_log_stream_state( + task_id: str, +) -> _RemoteContainerLogStreamState: + """Create an unstarted container-log watcher state.""" + stop_event = threading.Event() + return _RemoteContainerLogStreamState( + task_id=task_id, + stop_event=stop_event, + thread=threading.Thread( + target=_run_remote_container_log_stream, + args=(task_id, stop_event), + name=f"modal-log-stream-{task_id}", + daemon=True, + ), + ) + + +def _cancel_remote_container_log_idle_stop( + stream_state: _RemoteContainerLogStreamState, +) -> None: + """Cancel one pending idle-stop timer before a container is reused.""" + idle_stop_timer = stream_state.idle_stop_timer + if idle_stop_timer is None: + return + idle_stop_timer.cancel() + stream_state.idle_stop_timer = None + + +def _stop_remote_container_log_stream_if_idle( + task_id: str, + expected_state: _RemoteContainerLogStreamState, +) -> None: + """Stop one watcher after its grace period if no payload reused it.""" + with _REMOTE_CONTAINER_LOG_STREAMS_LOCK: + stream_state = _REMOTE_CONTAINER_LOG_STREAMS.get(task_id) + if stream_state is not expected_state or stream_state.refcount != 0: + return + stream_state.idle_stop_timer = None + _REMOTE_CONTAINER_LOG_STREAMS.pop(task_id, None) + + logger.info( + "Stopping idle remote Modal container log stream for task_id=%s.", + task_id, + ) + stream_state.stop_event.set() + stream_state.thread.join(timeout=0.2) + + +def _schedule_remote_container_log_idle_stop( + stream_state: _RemoteContainerLogStreamState, +) -> None: + """Schedule one daemon timer to stop an unreferenced log watcher.""" + idle_stop_timer = threading.Timer( + _REMOTE_CONTAINER_LOG_STREAM_IDLE_GRACE_SECONDS, + _stop_remote_container_log_stream_if_idle, + args=(stream_state.task_id, stream_state), + ) + idle_stop_timer.daemon = True + stream_state.idle_stop_timer = idle_stop_timer + idle_stop_timer.start() + + +def _retain_remote_container_log_stream(task_id: str) -> str: + """Increment one shared Modal container log watcher and start it if needed.""" + with _REMOTE_CONTAINER_LOG_STREAMS_LOCK: + stream_state = _REMOTE_CONTAINER_LOG_STREAMS.get(task_id) + if stream_state is not None and not stream_state.thread.is_alive(): + _cancel_remote_container_log_idle_stop(stream_state) + _REMOTE_CONTAINER_LOG_STREAMS.pop(task_id, None) + stream_state = None + + if stream_state is None: + stream_state = _new_remote_container_log_stream_state(task_id) + _REMOTE_CONTAINER_LOG_STREAMS[task_id] = stream_state + logger.info( + "Creating remote Modal container log stream for task_id=%s.", task_id + ) + stream_state.thread.start() + else: + _cancel_remote_container_log_idle_stop(stream_state) + logger.info( + "Reusing remote Modal container log stream for task_id=%s.", task_id + ) + stream_state.refcount += 1 + logger.info( + "Remote Modal container log stream retain task_id=%s refcount=%d.", + task_id, + stream_state.refcount, + ) + + return task_id + + +def _release_remote_container_log_stream(task_id: str) -> None: + """Release one watcher while allowing prompt-local container reuse.""" + with _REMOTE_CONTAINER_LOG_STREAMS_LOCK: + stream_state = _REMOTE_CONTAINER_LOG_STREAMS.get(task_id) + if stream_state is None: + return + stream_state.refcount = max(0, stream_state.refcount - 1) + logger.info( + "Remote Modal container log stream release task_id=%s refcount=%d.", + task_id, + stream_state.refcount, + ) + if stream_state.refcount != 0: + return + if not stream_state.thread.is_alive(): + _REMOTE_CONTAINER_LOG_STREAMS.pop(task_id, None) + return + if stream_state.idle_stop_timer is None: + logger.info( + "Keeping remote Modal container log stream alive for %.1fs " + "to allow task reuse without replaying history task_id=%s.", + _REMOTE_CONTAINER_LOG_STREAM_IDLE_GRACE_SECONDS, + task_id, + ) + _schedule_remote_container_log_idle_stop(stream_state) diff --git a/tests/conftest.py b/tests/conftest.py index 18197e0..feae952 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -449,6 +449,12 @@ def modal_billing_module(extension_package: object) -> object: return importlib.import_module(f"{PACKAGE_NAME}.remote.modal_billing") +@pytest.fixture(scope="session") +def modal_container_logs_module(extension_package: object) -> object: + """Return managed Modal container and log-stream state.""" + return importlib.import_module(f"{PACKAGE_NAME}.remote.modal_container_logs") + + @pytest.fixture(scope="session") def serialization_module(extension_package: object) -> object: """Return the serialization helpers module.""" diff --git a/tests/test_modal_executor_node.py b/tests/test_modal_executor_node.py index 655e8d9..f55d550 100644 --- a/tests/test_modal_executor_node.py +++ b/tests/test_modal_executor_node.py @@ -8403,13 +8403,14 @@ def test_emit_local_mapped_lane_progress_start_marks_lane_as_setup_only( def test_remote_modal_consumes_remote_log_stream_events_with_retain_release( remote_modal_app_module: Any, + modal_container_logs_module: Any, monkeypatch: Any, ) -> None: """Stream consumers should retain and release one container-log watcher per remote payload.""" log_stream_calls: list[tuple[str, str]] = [] monkeypatch.setattr( - remote_modal_app_module, + modal_container_logs_module, "get_settings", lambda: types.SimpleNamespace(stream_remote_container_logs=True), ) @@ -8444,13 +8445,13 @@ def test_remote_modal_consumes_remote_log_stream_events_with_retain_release( def test_list_active_modal_containers_filters_and_classifies_managed_apps( - remote_modal_app_module: Any, + modal_container_logs_module: Any, monkeypatch: Any, ) -> None: """Container status should include every active GPU app for this ComfyUI instance.""" - settings = remote_modal_app_module.get_settings() - b300_app_name = remote_modal_app_module.modal_deployment_app_name( - remote_modal_app_module.settings_for_modal_gpu(settings, "B300") + settings = modal_container_logs_module.get_settings() + b300_app_name = modal_container_logs_module.modal_deployment_app_name( + modal_container_logs_module.settings_for_modal_gpu(settings, "B300") ) request_environment_names: list[str] = [] request_thread_ids: list[int] = [] @@ -8524,7 +8525,7 @@ def test_list_active_modal_containers_filters_and_classifies_managed_apps( return blocking_callable - original_import_module = remote_modal_app_module.importlib.import_module + original_import_module = modal_container_logs_module.importlib.import_module def fake_import_module(name: str) -> Any: """Supply the Modal SDK modules used by the container status query.""" @@ -8545,15 +8546,15 @@ def test_list_active_modal_containers_filters_and_classifies_managed_apps( ) return original_import_module(name) - monkeypatch.setattr(remote_modal_app_module, "modal", object()) - monkeypatch.setattr(remote_modal_app_module.importlib, "import_module", fake_import_module) + monkeypatch.setattr(modal_container_logs_module, "modal", object()) + monkeypatch.setattr(modal_container_logs_module.importlib, "import_module", fake_import_module) first_containers = asyncio.run( - remote_modal_app_module.list_active_modal_containers(settings) + modal_container_logs_module.list_active_modal_containers(settings) ) - containers = asyncio.run(remote_modal_app_module.list_active_modal_containers(settings)) + containers = asyncio.run(modal_container_logs_module.list_active_modal_containers(settings)) stopped = asyncio.run( - remote_modal_app_module.stop_managed_modal_container("ta-running", settings) + modal_container_logs_module.stop_managed_modal_container("ta-running", settings) ) assert first_containers == containers @@ -8572,7 +8573,7 @@ def test_list_active_modal_containers_filters_and_classifies_managed_apps( def test_stop_managed_modal_container_verifies_ownership_before_exact_stop( - remote_modal_app_module: Any, + modal_container_logs_module: Any, monkeypatch: Any, ) -> None: """Container termination must target only a currently listed managed task.""" @@ -8581,7 +8582,7 @@ def test_stop_managed_modal_container_verifies_ownership_before_exact_stop( async def fake_list_active(_settings: Any) -> list[Any]: """Return one container already filtered to this ComfyUI installation.""" return [ - remote_modal_app_module.ModalContainerStatus( + modal_container_logs_module.ModalContainerStatus( container_id="ta-managed", app_id="ap-managed", app_name="comfy-modal-sync-B300", @@ -8601,7 +8602,7 @@ def test_stop_managed_modal_container_verifies_ownership_before_exact_stop( """Record the exact task passed to the synchronized SDK bridge.""" stopped_container_ids.append(container_id) - original_import_module = remote_modal_app_module.importlib.import_module + original_import_module = modal_container_logs_module.importlib.import_module def fake_import_module(name: str) -> Any: """Supply the minimal modules needed after ownership verification.""" @@ -8614,43 +8615,43 @@ def test_stop_managed_modal_container_verifies_ownership_before_exact_stop( return original_import_module(name) monkeypatch.setattr( - remote_modal_app_module, + modal_container_logs_module, "list_active_modal_containers", fake_list_active, ) monkeypatch.setattr( - remote_modal_app_module, + modal_container_logs_module, "_stop_modal_task_synchronously", fake_stop, ) - monkeypatch.setattr(remote_modal_app_module.importlib, "import_module", fake_import_module) + monkeypatch.setattr(modal_container_logs_module.importlib, "import_module", fake_import_module) assert asyncio.run( - remote_modal_app_module.stop_managed_modal_container( + modal_container_logs_module.stop_managed_modal_container( "ta-managed", - remote_modal_app_module.get_settings(), + modal_container_logs_module.get_settings(), ) ) assert not asyncio.run( - remote_modal_app_module.stop_managed_modal_container( + modal_container_logs_module.stop_managed_modal_container( "ta-unrelated", - remote_modal_app_module.get_settings(), + modal_container_logs_module.get_settings(), ) ) assert stopped_container_ids == ["ta-managed"] def test_modal_gpu_estimated_rates_cover_supported_aliases( - remote_modal_app_module: Any, + modal_container_logs_module: Any, ) -> None: """Published GPU estimates should cover every selectable billed GPU identity.""" - rates = remote_modal_app_module.MODAL_GPU_ESTIMATED_USD_PER_SECOND + rates = modal_container_logs_module.MODAL_GPU_ESTIMATED_USD_PER_SECOND - assert set(rates) == set(remote_modal_app_module.MODAL_GPU_TYPES) + assert set(rates) == set(modal_container_logs_module.MODAL_GPU_TYPES) assert rates["A100"] == rates["A100-40GB"] assert rates["H100!"] == rates["H100"] assert rates["B200+"] == rates["B200"] - assert remote_modal_app_module.MODAL_GPU_PRICING_EFFECTIVE_DATE == "2026-08-13" + assert modal_container_logs_module.MODAL_GPU_PRICING_EFFECTIVE_DATE == "2026-08-13" def test_completed_modal_billing_interval_uses_hourly_resolution_and_buffer( @@ -9026,7 +9027,7 @@ def test_remote_modal_stops_consuming_stream_after_terminal_result( def test_remote_modal_cli_log_stream_mirrors_lines_to_stderr( - remote_modal_app_module: Any, + modal_container_logs_module: Any, monkeypatch: Any, ) -> None: """CLI-backed container log streaming should mirror complete prefixed lines into local stderr.""" @@ -9076,10 +9077,10 @@ def test_remote_modal_cli_log_stream_mirrors_lines_to_stderr( self.returncode = 0 stderr_buffer = StringIO() - monkeypatch.setattr(remote_modal_app_module.shutil, "which", lambda name: "/usr/bin/modal") - monkeypatch.setattr(remote_modal_app_module.subprocess, "Popen", lambda *args, **kwargs: FakeProcess()) + monkeypatch.setattr(modal_container_logs_module.shutil, "which", lambda name: "/usr/bin/modal") + monkeypatch.setattr(modal_container_logs_module.subprocess, "Popen", lambda *args, **kwargs: FakeProcess()) monkeypatch.setattr( - remote_modal_app_module.select, + modal_container_logs_module.select, "select", lambda streams, _write, _error, _timeout: ( streams if streams[0]._index < len(streams[0]._chunks) else [], @@ -9087,9 +9088,9 @@ def test_remote_modal_cli_log_stream_mirrors_lines_to_stderr( [], ), ) - monkeypatch.setattr(remote_modal_app_module.sys, "stderr", stderr_buffer) + monkeypatch.setattr(modal_container_logs_module.sys, "stderr", stderr_buffer) - result = remote_modal_app_module._stream_remote_container_logs_via_modal_cli( + result = modal_container_logs_module._stream_remote_container_logs_via_modal_cli( "ta-123", threading.Event(), ) @@ -9102,30 +9103,30 @@ def test_remote_modal_cli_log_stream_mirrors_lines_to_stderr( def test_remote_modal_log_stream_prefers_cli_before_sdk( - remote_modal_app_module: Any, + modal_container_logs_module: Any, monkeypatch: Any, ) -> None: """The local log watcher should avoid the SDK path when the Modal CLI is available.""" backend_calls: list[str] = [] monkeypatch.setattr( - remote_modal_app_module, + modal_container_logs_module, "_stream_remote_container_logs_via_modal_cli", lambda task_id, stop_event: backend_calls.append(f"cli:{task_id}") or True, ) monkeypatch.setattr( - remote_modal_app_module, + modal_container_logs_module, "_stream_remote_container_logs_via_modal_sdk", lambda task_id, stop_event: backend_calls.append(f"sdk:{task_id}") or True, ) - remote_modal_app_module._run_remote_container_log_stream("ta-123", threading.Event()) + modal_container_logs_module._run_remote_container_log_stream("ta-123", threading.Event()) assert backend_calls == ["cli:ta-123"] def test_remote_modal_log_stream_survives_short_container_reuse_gap( - remote_modal_app_module: Any, + modal_container_logs_module: Any, monkeypatch: Any, ) -> None: """A reused task id should keep one watcher and avoid replaying container history.""" @@ -9185,32 +9186,32 @@ def test_remote_modal_log_stream_survives_short_container_reuse_gap( if not self.cancelled: self.function(*self.args) - monkeypatch.setattr(remote_modal_app_module, "_REMOTE_CONTAINER_LOG_STREAMS", {}) - monkeypatch.setattr(remote_modal_app_module.threading, "Thread", FakeThread) - monkeypatch.setattr(remote_modal_app_module.threading, "Timer", FakeTimer) + monkeypatch.setattr(modal_container_logs_module, "_REMOTE_CONTAINER_LOG_STREAMS", {}) + monkeypatch.setattr(modal_container_logs_module.threading, "Thread", FakeThread) + monkeypatch.setattr(modal_container_logs_module.threading, "Timer", FakeTimer) - assert remote_modal_app_module._retain_remote_container_log_stream("ta-reused") == "ta-reused" - stream_state = remote_modal_app_module._REMOTE_CONTAINER_LOG_STREAMS["ta-reused"] - remote_modal_app_module._release_remote_container_log_stream("ta-reused") + assert modal_container_logs_module._retain_remote_container_log_stream("ta-reused") == "ta-reused" + stream_state = modal_container_logs_module._REMOTE_CONTAINER_LOG_STREAMS["ta-reused"] + modal_container_logs_module._release_remote_container_log_stream("ta-reused") assert len(created_threads) == 1 assert len(created_timers) == 1 assert created_timers[0].daemon is True assert stream_state.stop_event.is_set() is False - remote_modal_app_module._retain_remote_container_log_stream("ta-reused") + modal_container_logs_module._retain_remote_container_log_stream("ta-reused") assert len(created_threads) == 1 assert created_timers[0].cancelled is True assert stream_state.refcount == 1 - remote_modal_app_module._release_remote_container_log_stream("ta-reused") + modal_container_logs_module._release_remote_container_log_stream("ta-reused") assert len(created_timers) == 2 created_timers[1].fire() assert stream_state.stop_event.is_set() is True assert stream_state.thread.join_calls == [0.2] - assert "ta-reused" not in remote_modal_app_module._REMOTE_CONTAINER_LOG_STREAMS + assert "ta-reused" not in modal_container_logs_module._REMOTE_CONTAINER_LOG_STREAMS def test_remote_modal_consumes_streamed_executed_outputs_and_previews(