updates to logging hierarchy

Author: apoorvkh

src/torchrunx/agent.py

@@ -49,12 +49,11 @@ def main(
     """
     # Setup logging & stream logs to server
 
-    logger = logging.getLogger(f"{__package__}.{hostname}")
-
     log_records_to_socket(
         hostname=hostname, local_rank=None, logger_hostname=logger_hostname, logger_port=logger_port
     )
 
+    logger = logging.getLogger()
     redirect_stdio_to_logger(logger)
 
     logger.debug("Initializing launcher-agent group.")

src/torchrunx/utils/logs.py

@@ -169,7 +169,19 @@ def handle(self) -> None:
                     while len(chunk) < slen:
                         chunk = chunk + self.connection.recv(slen - len(chunk))
                     obj = pickle.loads(chunk)
-                    record = logging.makeLogRecord(obj)
+
+                    ## Transform log record
+
+                    record: WorkerLogRecord = logging.makeLogRecord(obj)  # pyright: ignore [reportAssignmentType]
+
+                    if record.name != "root":
+                        record.msg = f"{record.name}:{record.msg}"
+
+                    record.name = f"torchrunx.{record.hostname}"
+                    if record.local_rank is not None:
+                        record.name += f".{record.local_rank}"
+
+                    ## Handle log record
 
                     for handler in handlers:
                         handler.handle(record)

src/torchrunx/worker.py

@@ -59,15 +59,14 @@ def worker_entrypoint(serialized_worker_args: bytes) -> Any | ExceptionFromWorke
 
     # Start logging to the logging server (i.e. the launcher)
 
-    logger = logging.getLogger(f"{__package__}.{worker_args.hostname}.{worker_args.local_rank}")
-
     log_records_to_socket(
         hostname=worker_args.hostname,
         local_rank=worker_args.local_rank,
         logger_hostname=worker_args.logger_hostname,
         logger_port=worker_args.logger_port,
     )
 
+    logger = logging.getLogger()
     redirect_stdio_to_logger(logger)
 
     # Set rank/world environment variables