updates to logging

Author: apoorvkh

src/torchrunx/agent.py

@@ -19,7 +19,7 @@
     LauncherAgentGroup,
     get_open_port,
 )
-from .utils.logging import log_records_to_socket, redirect_stdio_to_logger
+from .utils.logs import log_records_to_socket, redirect_stdio_to_logger
 from .worker import WorkerArgs, worker_entrypoint
 
 
@@ -48,23 +48,19 @@ def main(
         hostname: Hostname of this agent.
     """
     # Stream logs to logging server
-
-    logger = logging.getLogger()
-    redirect_stdio_to_logger(logger)
+    logger = logging.getLogger(f"{__package__}.{hostname}")
 
     log_records_to_socket(
-        logger=logger,
-        hostname=hostname,
-        local_rank=None,
-        logger_hostname=logger_hostname,
-        logger_port=logger_port,
+        hostname=hostname, local_rank=None, logger_hostname=logger_hostname, logger_port=logger_port
     )
 
-    logging.debug("Agent logging setup.")
+    redirect_stdio_to_logger(logger)
+
+    logger.debug("Agent logging setup.")
 
     # Set up launcher-agent group
 
-    logging.debug("Initializing launcher-agent group.")
+    logger.debug("Initializing launcher-agent group.")
 
     launcher_agent_group = LauncherAgentGroup(
         launcher_hostname=launcher_hostname,
@@ -77,7 +73,7 @@ def main(
 
     # Communicate initial payloads between launcher/agents
 
-    logging.debug("Sending agent details to launcher.")
+    logger.debug("Sending agent details to launcher.")
 
     payload = AgentPayload(
         hostname=socket.getfqdn(),
@@ -86,7 +82,6 @@ def main(
     )
 
     launcher_payload, agent_payloads = launcher_agent_group.sync_payloads(payload=payload)
-    main_agent_payload = agent_payloads[0]
 
     hostname = launcher_payload.hostnames[agent_rank]
     worker_world_size = launcher_payload.worker_world_size
@@ -95,7 +90,7 @@ def main(
 
     # Spawn worker processes
 
-    logging.debug("Launching worker processes.")
+    logger.debug("Launching worker processes.")
 
     ctx = dist_mp.start_processes(
         name=f"{hostname}_",
@@ -106,8 +101,8 @@ def main(
                     function=launcher_payload.fn,
                     logger_hostname=logger_hostname,
                     logger_port=logger_port,
-                    main_agent_hostname=main_agent_payload.hostname,
-                    main_agent_port=main_agent_payload.port,
+                    master_hostname=agent_payloads[0].hostname,
+                    master_port=agent_payloads[0].port,
                     backend=launcher_payload.backend,
                     rank=worker_global_ranks[i],
                     local_rank=i,
@@ -146,12 +141,12 @@ def main(
             all_done = all(s.state == "done" for s in agent_statuses)
             any_failed = any(s.state == "failed" for s in agent_statuses)
             if all_done or any_failed:
-                logging.debug("Workers exiting %s.", "cleanly" if not any_failed else "with errors")
+                logger.debug("Workers exiting %s.", "cleanly" if not any_failed else "with errors")
                 break
     finally:
         ctx.close()
         sys.stdout.flush()
         sys.stderr.flush()
         launcher_agent_group.shutdown()
 
-    logging.debug("Agent exiting.")
+    logger.debug("Agent exiting.")

src/torchrunx/integrations/parsing.py renamed to src/torchrunx/integrations/cli.py

@@ -29,7 +29,7 @@
     resolve_environment,
 )
 from .utils.errors import ExceptionFromWorker, WorkerFailedError
-from .utils.logging import LoggingServerArgs, default_handlers, start_logging_server
+from .utils.logs import LoggingServerArgs, default_handlers, start_logging_server
 
 DEFAULT_ENV_VARS_FOR_COPY = (
     "PATH",
@@ -45,8 +45,6 @@
 FunctionP = ParamSpec("FunctionP")
 FunctionR = TypeVar("FunctionR")
 
-logger = logging.getLogger(__name__)
-
 
 @dataclass
 class Launcher:
@@ -104,6 +102,8 @@ def run(  # noqa: C901, PLR0912, PLR0915
             WorkerFailedError: If a worker fails (e.g. from a segmentation fault).
             AgentFailedError: If an agent fails, e.g. from an OS signal.
         """
+        logger = logging.getLogger(__package__)
+
         if not dist.is_available():
             msg = "The torch.distributed package is not available."
             raise RuntimeError(msg)
@@ -249,11 +249,6 @@ def handler_factory() -> list[logging.Handler]:
             if log_process is not None:
                 log_process.kill()
 
-            logger.debug("Killing launcher-agent group.")
-
-            if launcher_agent_group is not None:
-                launcher_agent_group.shutdown()
-
             # cleanup: SIGTERM all agents
             if agent_payloads is not None:
                 for agent_payload, agent_hostname in zip(agent_payloads, hostnames):
@@ -265,6 +260,10 @@ def handler_factory() -> list[logging.Handler]:
                         ssh_config_file=ssh_config_file,
                     )
 
+            if launcher_agent_group is not None:
+                logger.debug("Killing launcher-agent group.")
+                launcher_agent_group.shutdown()
+
 
 @dataclass
 class LaunchResult(Generic[FunctionR]):

src/torchrunx/launcher.py

@@ -1,3 +0,0 @@
-from .logging import add_filter_to_handler, file_handler, stream_handler
-
-__all__ = ["add_filter_to_handler", "file_handler", "stream_handler"]

src/torchrunx/utils/__init__.py

@@ -119,7 +119,7 @@ class LauncherPayload:
     hostnames: list[str]
     worker_global_ranks: list[list[int]]
     worker_world_size: int
-    backend: Literal["nccl", "gloo", "mpi", "ucc", "auto"] | None
+    backend: Literal["nccl", "gloo", "mpi", "ucc"] | None
     timeout: int
 
 

src/torchrunx/utils/comm.py

@@ -1,4 +1,4 @@
-"""Utilities for intercepting logs in worker processes and handling these in the Launcher."""  # noqa: A005
+"""Utilities for intercepting logs in worker processes and handling these in the Launcher."""
 
 from __future__ import annotations
 
@@ -24,7 +24,7 @@
 from contextlib import redirect_stderr, redirect_stdout
 from dataclasses import dataclass
 from io import StringIO
-from logging import Handler, Logger
+from logging import Handler, Logger, LogRecord
 from logging.handlers import SocketHandler
 from multiprocessing.synchronize import Event as EventClass
 from pathlib import Path
@@ -55,40 +55,40 @@ def _filter(record: WorkerLogRecord) -> bool:
     handler.addFilter(_filter)  # pyright: ignore [reportArgumentType]
 
 
-def default_handlers(
-    hostnames: list[str],
-    workers_per_host: list[int],
-    log_level: int = logging.INFO,
-) -> list[logging.Handler]:
+def default_handlers(hostnames: list[str], workers_per_host: list[int]) -> list[logging.Handler]:
     """Constructs default :obj:`logging.Handler` objects.
 
     Logs for the rank 0 agent and worker are written to launcher process stdout.
     Logs for all hosts/workers are written to files in ``$TORCHRUNX_LOG_DIR`` (named by timestamp,
     hostname, local_rank).
     """
     log_dir = Path(os.environ.get("TORCHRUNX_LOG_DIR", "torchrunx_logs"))
-    log_level = logging._nameToLevel[os.environ.get("TORCHRUNX_LOG_LEVEL", "INFO")]  # noqa: SLF001
+    file_log_level = logging._nameToLevel[os.environ.get("TORCHRUNX_LOG_LEVEL", "INFO")]  # noqa: SLF001
+
     return [
-        stream_handler(hostname=hostnames[0], local_rank=None, log_level=log_level),
-        stream_handler(hostname=hostnames[0], local_rank=0, log_level=log_level),
-        *file_handlers(hostnames, workers_per_host, log_dir=log_dir, log_level=log_level),
+        RedirectHandler(hostname=hostnames[0], local_rank=None),
+        RedirectHandler(hostname=hostnames[0], local_rank=0),
+        *file_handlers(hostnames, workers_per_host, log_dir=log_dir, log_level=file_log_level),
     ]
 
 
+class RedirectHandler(logging.Handler):
+    def __init__(self, hostname: str, local_rank: int | None) -> None:
+        super().__init__()
+        add_filter_to_handler(self, hostname=hostname, local_rank=local_rank)
+
+    def emit(self, record: LogRecord) -> None:
+        logger = logging.getLogger(record.name)
+        if logger.isEnabledFor(record.levelno):
+            logger.handle(record)
+
+
 def stream_handler(
     hostname: str, local_rank: int | None, log_level: int = logging.NOTSET
 ) -> logging.Handler:
     """Handler builder function for writing logs from specified hostname/rank to stdout."""
     handler = logging.StreamHandler(stream=sys.stdout)
     add_filter_to_handler(handler, hostname, local_rank, log_level=log_level)
-    handler.setFormatter(
-        logging.Formatter(
-            "%(asctime)s:%(levelname)s:%(hostname)s[%(local_rank)s]: %(message)s"
-            if local_rank is not None
-            else "%(asctime)s:%(levelname)s:%(hostname)s: %(message)s",
-            datefmt="%Y-%m-%d %H:%M:%S",
-        ),
-    )
     return handler
 
 
@@ -259,14 +259,13 @@ def from_record(cls, record: logging.LogRecord, hostname: str, local_rank: int |
 
 
 def log_records_to_socket(
-    logger: Logger,
     hostname: str,
     local_rank: int | None,  # None indicates agent
     logger_hostname: str,
     logger_port: int,
 ) -> None:
     """Encode LogRecords with hostname/local_rank. Send to TCP socket on Launcher."""
-    logger.setLevel(logging.NOTSET)
+    logging.root.setLevel(logging.NOTSET)
 
     old_factory = logging.getLogRecordFactory()
 
@@ -276,4 +275,4 @@ def record_factory(*args, **kwargs) -> WorkerLogRecord:  # noqa: ANN002, ANN003
 
     logging.setLogRecordFactory(record_factory)
 
-    logger.addHandler(SocketHandler(host=logger_hostname, port=logger_port))
+    logging.root.addHandler(SocketHandler(host=logger_hostname, port=logger_port))

src/torchrunx/utils/logging.py renamed to src/torchrunx/utils/logs.py

@@ -11,12 +11,11 @@
 from typing import Any, Callable, Literal
 
 import cloudpickle
-import torch
 import torch.distributed as dist
 from typing_extensions import Self
 
 from .utils.errors import ExceptionFromWorker
-from .utils.logging import log_records_to_socket, redirect_stdio_to_logger
+from .utils.logs import log_records_to_socket, redirect_stdio_to_logger
 
 __all__ = ["WorkerArgs", "worker_entrypoint"]
 
@@ -28,9 +27,9 @@ class WorkerArgs:
     function: Callable
     logger_hostname: str
     logger_port: int
-    main_agent_hostname: str
-    main_agent_port: int
-    backend: Literal["nccl", "gloo", "mpi", "ucc", "auto"] | None
+    master_hostname: str
+    master_port: int
+    backend: Literal["nccl", "gloo", "mpi", "ucc"] | None
     rank: int
     local_rank: int
     node_rank: int
@@ -60,10 +59,9 @@ def worker_entrypoint(serialized_worker_args: bytes) -> Any | ExceptionFromWorke
 
     # Start logging to the logging server (i.e. the launcher)
 
-    logger = logging.getLogger()
+    logger = logging.getLogger(f"{__package__}.{worker_args.hostname}.{worker_args.local_rank}")
 
     log_records_to_socket(
-        logger=logger,
         hostname=worker_args.hostname,
         local_rank=worker_args.local_rank,
         logger_hostname=worker_args.logger_hostname,
@@ -79,23 +77,21 @@ def worker_entrypoint(serialized_worker_args: bytes) -> Any | ExceptionFromWorke
     os.environ["GROUP_RANK"] = str(worker_args.node_rank)
     os.environ["LOCAL_WORLD_SIZE"] = str(worker_args.local_world_size)
     os.environ["WORLD_SIZE"] = str(worker_args.world_size)
-    os.environ["MASTER_ADDR"] = worker_args.main_agent_hostname
-    os.environ["MASTER_PORT"] = str(worker_args.main_agent_port)
+    os.environ["MASTER_ADDR"] = worker_args.master_hostname
+    os.environ["MASTER_PORT"] = str(worker_args.master_port)
 
     # Prepare the process group (e.g. for communication within the user's function)
 
     if worker_args.backend is not None:
         backend = worker_args.backend
-        if backend == "auto":
-            backend = "nccl" if torch.cuda.is_available() else "gloo"
 
         dist.init_process_group(
             backend=backend,
             world_size=worker_args.world_size,
             rank=worker_args.rank,
             store=dist.TCPStore(  # pyright: ignore [reportPrivateImportUsage]
-                host_name=worker_args.main_agent_hostname,
-                port=worker_args.main_agent_port,
+                host_name=worker_args.master_hostname,
+                port=worker_args.master_port,
                 world_size=worker_args.world_size,
                 is_master=(worker_args.rank == 0),
             ),