Extract Modal container administration
This commit is contained in:
Ken Simpson
+67
-623
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
@@ -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."""
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user