100 lines
3.1 KiB
Python
100 lines
3.1 KiB
Python
"""Timestamped logging and phase timing for the cloud worker runtime."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from contextlib import contextmanager
|
|
import logging
|
|
import os
|
|
import sys
|
|
import time
|
|
from typing import Any, Iterator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_TARGET_LOGGER: logging.Logger = logger
|
|
_HANDLER_NAME = "comfyui-modal-sync-cloud-timestamped"
|
|
|
|
|
|
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
|
|
|
|
|
|
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,
|
|
)
|