666 lines
22 KiB
Python
666 lines
22 KiB
Python
"""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)
|