change logging structure

Author: pmcurtin

src/torchrunx/agent.py

@@ -80,6 +80,7 @@ def entrypoint(serialized_worker_args: bytes):
     logger.setLevel(logging.DEBUG)
     logger.addHandler(socketHandler)
     logger.debug("creating TCPStore for worker group.")
+
     store = dist.TCPStore(  # pyright: ignore[reportPrivateImportUsage]
         host_name=worker_args.master_hostname,
         port=worker_args.master_port,

src/torchrunx/launcher.py

@@ -27,6 +27,7 @@
     LauncherAgentGroup,
     LauncherPayload,
     LogRecordSocketReceiver,
+    default_logging,
     get_open_port,
 )
 
@@ -71,7 +72,7 @@ class Launcher:
     ssh_config_file: str | os.PathLike | None = None
     backend: Literal["mpi", "gloo", "nccl", "ucc", None] = None
     log_dir: os.PathLike | str = "./logs"
-    propagate_logs: bool = True
+    log_spec: dict[str, list[logging.Handler]] | None = None
     env_vars: list[str] = field(
         default_factory=lambda: [
             "PATH",
@@ -106,16 +107,26 @@ def run(
 
         logger = logging.getLogger("torchrunx")
         logger.setLevel(logging.DEBUG)
-        logger.propagate = self.propagate_logs
+        logger.propagate = False
 
         log_dir = Path(self.log_dir)
         log_dir.mkdir(parents=True, exist_ok=True)
-        timestamp = datetime.datetime.now().isoformat(timespec="seconds")
-
-        log_file_formatter = logging.Formatter("%(asctime)s:%(levelname)s:%(name)s:%(message)s")
-        log_file_handler = logging.FileHandler(f"{log_dir}/{timestamp}.log")
-        log_file_handler.setFormatter(log_file_formatter)
-        logger.addHandler(log_file_handler)
+        #timestamp = datetime.datetime.now().isoformat(timespec="seconds")
+
+        if self.log_spec is None:
+            # TODO: this assumes the type of workers_per_host is simply int. We should consider
+            # again whether it's worth supporting inhomogeneous allocations (list[int])
+            self.log_spec = default_logging(num_agents=len(self.hostnames), 
+                                            num_workers=self.workers_per_host, # type: ignore
+                                            log_dir=os.fspath(log_dir))
+
+        log_formatter = logging.Formatter("%(asctime)s:%(levelname)s:%(name)s:%(message)s")
+
+        for lname, handlers in self.log_spec.items(): # type: ignore
+            _logger = logging.getLogger(f"torchrunx.{lname}")
+            for handler in handlers:
+                handler.setFormatter(log_formatter)
+                _logger.addHandler(handler)
 
         log_process = Process(target=monitor_log, args=(), daemon=True)
         log_process.start()
@@ -194,7 +205,7 @@ def run(
 
         worker_log_names = [
             [
-                f"torchrunx.agent-{i}.worker-{local_rank}"
+                f"torchrunx.agent-{i}-worker-{local_rank}"
                 for local_rank in range(workers_per_host[i])  # type: ignore
             ]
             for i in range(len(self.hostnames))
@@ -263,7 +274,7 @@ def launch(
     ssh_config_file: str | os.PathLike | None = None,
     backend: Literal["mpi", "gloo", "nccl", "ucc", None] = None,
     log_dir: os.PathLike | str = "./logs",
-    propagate_logs: bool = True,
+    log_spec: dict[str, list[logging.Handler]] | None = None,
     env_vars: list[str] = [
         "PATH",
         "LD_LIBRARY",
@@ -294,8 +305,8 @@ def launch(
     :type backend: Literal['mpi', 'gloo', 'nccl', 'ucc', None], optional
     :param log_dir: A directory in which logs should be written, defaults to "./logs"
     :type log_dir: os.PathLike | str, optional
-    :param log_level: The logging level, defaults to logging.WARN
-    :type log_level: logging._Level, optional
+    :param log_spec: TODO
+    :type log_spec: TODO
     :param env_vars: A list of environmental variables to be copied from the launcher environment to workers. Allows for bash pattern matching syntax, defaults to ["PATH", "LD_LIBRARY", "LIBRARY_PATH", "PYTHON*", "CUDA*", "TORCH*", "PYTORCH*", "NCCL*"]
     :type env_vars: list[str], optional
     :param env_file: An additional environment file that will be sourced prior to executing ``func``, defaults to None
@@ -312,7 +323,7 @@ def launch(
         ssh_config_file=ssh_config_file,
         backend=backend,
         log_dir=log_dir,
-        propagate_logs=propagate_logs,
+        log_spec=log_spec,
         env_vars=env_vars,
         env_file=env_file,
         timeout=timeout,

src/torchrunx/utils.py

@@ -118,7 +118,8 @@ def sync_payloads(
 
     def sync_agent_statuses(self, status: AgentStatus) -> list[AgentStatus]:
         return self._all_gather(object=status)[1:]
-    
+
+
 class LogRecordStreamHandler(socketserver.StreamRequestHandler):
     """Handler for a streaming logging request.
 
@@ -136,7 +137,7 @@ def handle(self):
             chunk = self.connection.recv(4)
             if len(chunk) < 4:
                 break
-            slen = struct.unpack('>L', chunk)[0]
+            slen = struct.unpack(">L", chunk)[0]
             chunk = self.connection.recv(slen)
             while len(chunk) < slen:
                 chunk = chunk + self.connection.recv(slen - len(chunk))
@@ -150,7 +151,7 @@ def unPickle(self, data):
     def handleLogRecord(self, record):
         # if a name is specified, we use the named logger rather than the one
         # implied by the record.
-        if self.server.logname is not None: # type: ignore
+        if self.server.logname is not None:  # type: ignore
             name = self.server.logname  # type: ignore
         else:
             name = record.name
@@ -162,28 +163,62 @@ def handleLogRecord(self, record):
         if logger.getEffectiveLevel() <= record.levelno:
             logger.handle(record)
 
+
 class LogRecordSocketReceiver(socketserver.ThreadingTCPServer):
     """
     Simple TCP socket-based logging receiver suitable for testing.
     """
 
-    allow_reuse_address = 1 # type: ignore
+    allow_reuse_address = 1  # type: ignore
 
-    def __init__(self, host='localhost',
-                 port=logging.handlers.DEFAULT_TCP_LOGGING_PORT,
-                 handler=LogRecordStreamHandler):
+    def __init__(
+        self,
+        host="localhost",
+        port=logging.handlers.DEFAULT_TCP_LOGGING_PORT,
+        handler=LogRecordStreamHandler,
+    ):
         socketserver.ThreadingTCPServer.__init__(self, (host, port), handler)
         self.abort = 0
         self.timeout = 1
         self.logname = None
-        
 
     def serve_until_stopped(self):
         abort = 0
         while not abort:
-            rd, wr, ex = select.select([self.socket.fileno()],
-                                       [], [],
-                                       self.timeout)
+            rd, wr, ex = select.select([self.socket.fileno()], [], [], self.timeout)
             if rd:
                 self.handle_request()
             abort = self.abort
+
+
+def default_logging(
+    num_agents: int, num_workers: int, log_dir: str
+) -> dict[str, list[logging.Handler]]:
+    """
+    Generates torchrunx's default
+
+    :param num_agents: Number of agents in work group
+    :type num_agents: int
+    :param num_workers: Number of workers per agent
+    :type num_workers: int
+    :return: A logging structure to be passed to :mod:`torchrunx.launch` as the ``log_spec`` argument
+    :rtype: dict[str, list[logging.Handler]]
+    """
+
+    timestamp = datetime.datetime.now().isoformat(timespec="seconds")
+
+    agents: dict[str, list[logging.Handler]] = {
+        f"agent-{i}": [logging.FileHandler(f"{log_dir}/{timestamp}-agent-{i}.log")]
+        for i in range(num_agents)
+    }
+    workers: dict[str, list[logging.Handler]] = {
+        f"agent-{i}-worker-{j}": [
+            logging.FileHandler(f"{log_dir}/{timestamp}-agent-{i}.worker-{j}.log")
+        ]
+        for j in range(num_workers)
+        for i in range(num_agents)
+    }
+
+    workers["agent-0-worker-0"].append(logging.StreamHandler())
+
+    return {**agents, **workers}