WorkerLogRecord class

apoorvkh · apoorvkh · commit 9c97d09fb362 · 2024-09-12T18:38:31.000-04:00
diff --git a/src/torchrunx/__init__.py b/src/torchrunx/__init__.py
@@ -1,6 +1,10 @@
 from .launcher import Launcher, launch
+from .logging_utils import add_filter_to_handler, file_handler, stream_handler
 
 __all__ = [
     "Launcher",
     "launch",
+    "add_filter_to_handler",
+    "file_handler",
+    "stream_handler",
 ]
diff --git a/src/torchrunx/agent.py b/src/torchrunx/agent.py
@@ -67,7 +67,7 @@ def entrypoint(serialized_worker_args: SerializedWorkerArgs) -> Any | WorkerExce
 
     redirect_stdio_to_logger(logger)
 
-    store = dist.TCPStore(  # pyright: ignore[reportPrivateImportUsage]
+    store = dist.TCPStore(  # pyright: ignore [reportPrivateImportUsage]
         host_name=worker_args.main_agent_hostname,
         port=worker_args.main_agent_port,
         world_size=worker_args.world_size,
diff --git a/src/torchrunx/logging_utils.py b/src/torchrunx/logging_utils.py
@@ -2,15 +2,115 @@
 
 import datetime
 import logging
-import os  # noqa: TCH003
 import pickle
 import struct
 from contextlib import redirect_stderr, redirect_stdout
+from dataclasses import dataclass
 from io import StringIO
 from logging import Handler, Logger
 from logging.handlers import SocketHandler
 from pathlib import Path
 from socketserver import StreamRequestHandler, ThreadingTCPServer
+from typing import TYPE_CHECKING
+
+from typing_extensions import Self
+
+if TYPE_CHECKING:
+    import os
+
+## Launcher utilities
+
+
+class LogRecordSocketReceiver(ThreadingTCPServer):
+    def __init__(self, host: str, port: int, handlers: list[Handler]) -> None:
+        self.host = host
+        self.port = port
+
+        class _LogRecordStreamHandler(StreamRequestHandler):
+            def handle(self) -> None:
+                while True:
+                    chunk_size = 4
+                    chunk = self.connection.recv(chunk_size)
+                    if len(chunk) < chunk_size:
+                        break
+                    slen = struct.unpack(">L", chunk)[0]
+                    chunk = self.connection.recv(slen)
+                    while len(chunk) < slen:
+                        chunk = chunk + self.connection.recv(slen - len(chunk))
+                    obj = pickle.loads(chunk)
+                    record = logging.makeLogRecord(obj)
+
+                    for handler in handlers:
+                        handler.handle(record)
+
+        super().__init__(
+            server_address=(host, port),
+            RequestHandlerClass=_LogRecordStreamHandler,
+            bind_and_activate=True,
+        )
+        self.daemon_threads = True
+
+    def shutdown(self) -> None:
+        """override BaseServer.shutdown() with added timeout"""
+        self._BaseServer__shutdown_request = True
+        self._BaseServer__is_shut_down.wait(timeout=3)  # pyright: ignore[reportAttributeAccessIssue]
+
+
+## Agent/worker utilities
+
+
+@dataclass
+class WorkerLogRecord(logging.LogRecord):
+    hostname: str
+    worker_rank: int | None
+
+    @classmethod
+    def from_record(cls, record: logging.LogRecord, hostname: str, worker_rank: int | None) -> Self:
+        record.hostname = hostname
+        record.worker_rank = worker_rank
+        record.__class__ = cls
+        return record  # pyright: ignore [reportReturnType]
+
+
+def log_records_to_socket(
+    logger: Logger,
+    hostname: str,
+    worker_rank: int | None,
+    logger_hostname: str,
+    logger_port: int,
+) -> None:
+    logger.setLevel(logging.NOTSET)
+
+    old_factory = logging.getLogRecordFactory()
+
+    def record_factory(*args, **kwargs) -> WorkerLogRecord:  # noqa: ANN002, ANN003
+        record = old_factory(*args, **kwargs)
+        return WorkerLogRecord.from_record(record, hostname, worker_rank)
+
+    logging.setLogRecordFactory(record_factory)
+
+    logger.addHandler(SocketHandler(host=logger_hostname, port=logger_port))
+
+
+def redirect_stdio_to_logger(logger: Logger) -> None:
+    class _LoggingStream(StringIO):
+        def __init__(self, logger: Logger, level: int = logging.NOTSET) -> None:
+            super().__init__()
+            self.logger = logger
+            self.level = level
+
+        def flush(self) -> None:
+            super().flush()
+            value = self.getvalue()
+            if value != "":
+                self.logger.log(self.level, value)
+                self.truncate(0)
+                self.seek(0)
+
+    logging.captureWarnings(capture=True)
+    redirect_stderr(_LoggingStream(logger, level=logging.ERROR)).__enter__()
+    redirect_stdout(_LoggingStream(logger, level=logging.INFO)).__enter__()
+
 
 ## Handler utilities
 
@@ -21,14 +121,27 @@ def add_filter_to_handler(
     worker_rank: int | None,
     log_level: int = logging.NOTSET,
 ) -> None:
-    def _filter(record: logging.LogRecord) -> bool:
+    def _filter(record: WorkerLogRecord) -> bool:
         return (
-            record.hostname == hostname  # pyright: ignore[reportAttributeAccessIssue]
-            and record.worker_rank == worker_rank  # pyright: ignore[reportAttributeAccessIssue]
+            record.hostname == hostname
+            and record.worker_rank == worker_rank
             and record.levelno >= log_level
         )
 
-    handler.addFilter(_filter)
+    handler.addFilter(_filter)  # pyright: ignore [reportArgumentType]
+
+
+def stream_handler(hostname: str, rank: int | None, log_level: int = logging.NOTSET) -> Handler:
+    handler = logging.StreamHandler()
+    add_filter_to_handler(handler, hostname, rank, log_level=log_level)
+    handler.setFormatter(
+        logging.Formatter(
+            "%(asctime)s:%(levelname)s:%(hostname)s[%(worker_rank)s]: %(message)s"
+            if rank is not None
+            else "%(asctime)s:%(levelname)s:%(hostname)s: %(message)s",
+        ),
+    )
+    return handler
 
 
 def file_handler(
@@ -67,19 +180,6 @@ def file_handlers(
     return handlers
 
 
-def stream_handler(hostname: str, rank: int | None, log_level: int = logging.NOTSET) -> Handler:
-    handler = logging.StreamHandler()
-    add_filter_to_handler(handler, hostname, rank, log_level=log_level)
-    handler.setFormatter(
-        logging.Formatter(
-            "%(asctime)s:%(levelname)s:%(hostname)s[%(worker_rank)s]: %(message)s"
-            if rank is not None
-            else "%(asctime)s:%(levelname)s:%(hostname)s: %(message)s",
-        ),
-    )
-    return handler
-
-
 def default_handlers(
     hostnames: list[str],
     workers_per_host: list[int],
@@ -91,86 +191,3 @@ def default_handlers(
         stream_handler(hostname=hostnames[0], rank=0, log_level=log_level),
         *file_handlers(hostnames, workers_per_host, log_dir=log_dir, log_level=log_level),
     ]
-
-
-## Agent/worker utilities
-
-
-def log_records_to_socket(
-    logger: Logger,
-    hostname: str,
-    worker_rank: int | None,
-    logger_hostname: str,
-    logger_port: int,
-) -> None:
-    logger.setLevel(logging.NOTSET)
-
-    old_factory = logging.getLogRecordFactory()
-
-    def record_factory(*args, **kwargs) -> logging.LogRecord:  # noqa: ANN002, ANN003
-        record = old_factory(*args, **kwargs)
-        record.hostname = hostname
-        record.worker_rank = worker_rank
-        return record
-
-    logging.setLogRecordFactory(record_factory)
-
-    logger.addHandler(SocketHandler(host=logger_hostname, port=logger_port))
-
-
-def redirect_stdio_to_logger(logger: Logger) -> None:
-    class _LoggingStream(StringIO):
-        def __init__(self, logger: Logger, level: int = logging.NOTSET) -> None:
-            super().__init__()
-            self.logger = logger
-            self.level = level
-
-        def flush(self) -> None:
-            super().flush()
-            value = self.getvalue()
-            if value != "":
-                self.logger.log(self.level, value)
-                self.truncate(0)
-                self.seek(0)
-
-    logging.captureWarnings(capture=True)
-    redirect_stderr(_LoggingStream(logger, level=logging.ERROR)).__enter__()
-    redirect_stdout(_LoggingStream(logger, level=logging.INFO)).__enter__()
-
-
-## Launcher utilities
-
-
-class LogRecordSocketReceiver(ThreadingTCPServer):
-    def __init__(self, host: str, port: int, handlers: list[Handler]) -> None:
-        self.host = host
-        self.port = port
-
-        class _LogRecordStreamHandler(StreamRequestHandler):
-            def handle(self) -> None:
-                while True:
-                    chunk_size = 4
-                    chunk = self.connection.recv(chunk_size)
-                    if len(chunk) < chunk_size:
-                        break
-                    slen = struct.unpack(">L", chunk)[0]
-                    chunk = self.connection.recv(slen)
-                    while len(chunk) < slen:
-                        chunk = chunk + self.connection.recv(slen - len(chunk))
-                    obj = pickle.loads(chunk)
-                    record = logging.makeLogRecord(obj)
-
-                    for handler in handlers:
-                        handler.handle(record)
-
-        super().__init__(
-            server_address=(host, port),
-            RequestHandlerClass=_LogRecordStreamHandler,
-            bind_and_activate=True,
-        )
-        self.daemon_threads = True
-
-    def shutdown(self) -> None:
-        """override BaseServer.shutdown() with added timeout"""
-        self._BaseServer__shutdown_request = True
-        self._BaseServer__is_shut_down.wait(timeout=3)  # pyright: ignore[reportAttributeAccessIssue]
diff --git a/src/torchrunx/utils.py b/src/torchrunx/utils.py
@@ -77,7 +77,7 @@ def __post_init__(self) -> None:
             backend="gloo",
             world_size=self.world_size,
             rank=self.rank,
-            store=dist.TCPStore(  # pyright: ignore[reportPrivateImportUsage]
+            store=dist.TCPStore(  # pyright: ignore [reportPrivateImportUsage]
                 host_name=self.launcher_hostname,
                 port=self.launcher_port,
                 world_size=self.world_size,
