ttulttul-ComfyUI-Modal/cloud_runtime_logging.py

"""Timestamped logging and phase timing for the cloud worker runtime."""

from __future__ import annotations

from contextlib import contextmanager
import logging
import os
from pathlib import Path
import sys
import time
from typing import Any, Iterable, Iterator

logger = logging.getLogger(__name__)

_TARGET_LOGGER: logging.Logger = logger
_HANDLER_NAME = "comfyui-modal-sync-cloud-timestamped"
_MODULE_HANDLER_NAME = "comfyui-modal-sync-cloud-modules"


def _build_cloud_log_formatter() -> logging.Formatter:
    """Return the default formatter for remote Modal-Sync logs with timestamps."""
    return logging.Formatter(
        fmt="%(asctime)s.%(msecs)03d +%(relativeCreated)07.0fms %(levelname)s [%(name)s] %(message)s",
        datefmt="%H:%M:%S",
    )


def configure_cloud_runtime_logging(
    target_logger: logging.Logger,
    handler_name: str,
) -> logging.Logger:
    """Configure and retain the stable entrypoint logger used by cloud helpers."""
    global _TARGET_LOGGER
    global _HANDLER_NAME
    _TARGET_LOGGER = target_logger
    _HANDLER_NAME = handler_name
    target_logger.setLevel(logging.INFO)
    for existing_handler in target_logger.handlers:
        if getattr(existing_handler, "name", "") == handler_name:
            return target_logger

    handler = logging.StreamHandler(sys.stdout)
    handler.set_name(handler_name)
    handler.setLevel(logging.INFO)
    handler.setFormatter(_build_cloud_log_formatter())
    target_logger.addHandler(handler)
    target_logger.propagate = False
    return target_logger


class _CurrentStdoutHandler(logging.Handler):
    """Write each record to whatever ``sys.stdout`` is when the record is emitted.

    Modal memory-snapshot restores can replace ``sys.stdout`` after import, so a
    stream captured at import time may no longer reach the container log.
    """

    def emit(self, record: logging.LogRecord) -> None:
        """Format one record and print it to the live stdout stream."""
        try:
            print(self.format(record), file=sys.stdout, flush=True)
        except (OSError, ValueError):
            self.handleError(record)


def cloud_project_module_names(repo_root: Path) -> tuple[str, ...]:
    """Return the flat module names of the project files shipped to the container."""
    return tuple(
        sorted(
            module_path.stem
            for module_path in repo_root.glob("*.py")
            if module_path.stem != "__init__"
        )
    )


def _cloud_module_handler(
    module_loggers: Iterable[logging.Logger],
) -> logging.Handler:
    """Return the shared module handler, reusing one already installed."""
    for module_logger in module_loggers:
        for existing_handler in module_logger.handlers:
            if getattr(existing_handler, "name", "") == _MODULE_HANDLER_NAME:
                return existing_handler
    handler = _CurrentStdoutHandler(level=logging.INFO)
    handler.set_name(_MODULE_HANDLER_NAME)
    handler.setFormatter(_cloud_formatter())
    return handler


def configure_cloud_module_loggers(module_names: Iterable[str]) -> tuple[str, ...]:
    """Emit INFO logs from sibling project modules inside the Modal container.

    Only the entrypoint logger used to get a handler, so INFO records from modules
    such as ``cloud_durable_invocation`` fell through to Python's WARNING-only
    fallback and were silently dropped.
    """
    module_loggers = [
        logging.getLogger(module_name)
        for module_name in module_names
        if module_name != _TARGET_LOGGER.name
    ]
    handler = _cloud_module_handler(module_loggers)
    for module_logger in module_loggers:
        module_logger.setLevel(logging.INFO)
        if handler not in module_logger.handlers:
            module_logger.addHandler(handler)
        module_logger.propagate = False
    configured_names = tuple(module_logger.name for module_logger in module_loggers)
    logger.debug("Configured cloud module loggers: %s", configured_names)
    return configured_names


def _is_modal_container_runtime() -> bool:
    """Return whether the current process is executing inside a Modal container."""
    return os.getenv("MODAL_IS_REMOTE") == "1" or bool(os.getenv("MODAL_TASK_ID"))


def _cloud_formatter() -> logging.Formatter:
    """Return the configured formatter used for cloud phase trace lines."""
    for existing_handler in _TARGET_LOGGER.handlers:
        if getattr(existing_handler, "name", "") == _HANDLER_NAME:
            formatter = existing_handler.formatter
            if formatter is not None:
                return formatter
    return _build_cloud_log_formatter()


def _emit_cloud_info(message: str, *args: Any) -> None:
    """Emit an info line and mirror it to stdout inside Modal containers."""
    if not _is_modal_container_runtime():
        _TARGET_LOGGER.info(message, *args)
        return

    record = _TARGET_LOGGER.makeRecord(
        _TARGET_LOGGER.name,
        logging.INFO,
        __file__,
        0,
        message,
        args,
        exc_info=None,
    )
    print(_cloud_formatter().format(record), file=sys.stdout, flush=True)


@contextmanager
def _timed_phase(phase: str, **fields: Any) -> Iterator[None]:
    """Log a start/finish pair with elapsed time for a named execution phase."""
    field_suffix = ""
    if fields:
        rendered_fields = " ".join(f"{key}={value}" for key, value in fields.items())
        field_suffix = f" {rendered_fields}"
    phase_started_at = time.perf_counter()
    _emit_cloud_info("Starting %s%s", phase, field_suffix)
    try:
        yield
    finally:
        _emit_cloud_info(
            "Finished %s in %.3fs%s",
            phase,
            time.perf_counter() - phase_started_at,
            field_suffix,
        )