Last active
June 9, 2025 19:08
-
-
Save finsberg/5447d830165aae8235a2d4163f61f510 to your computer and use it in GitHub Desktop.
Logging to file with MPI
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import logging | |
| from mpi4py import MPI | |
| mode2mpi_mode = { | |
| "w": MPI.MODE_WRONLY | MPI.MODE_CREATE | MPI.MODE_EXCL, | |
| "a": MPI.MODE_WRONLY | MPI.MODE_CREATE | MPI.MODE_APPEND, | |
| } | |
| class MPIFileHandler(logging.FileHandler): | |
| def __init__( | |
| self, | |
| filename: Path, | |
| mode: str = "a", | |
| comm=MPI.COMM_WORLD, | |
| delay: bool = False, | |
| ): | |
| self.comm = comm | |
| self.mpi_mode = mode2mpi_mode[mode] | |
| super().__init__(filename=filename, mode=mode, delay=delay) | |
| def _open(self): | |
| stream = MPI.File.Open(self.comm, self.baseFilename, self.mpi_mode) | |
| stream.Set_atomicity(True) | |
| return stream | |
| def emit(self, record): | |
| msg = self.format(record) | |
| self.stream.Write_shared((msg + self.terminator).encode(self.encoding)) | |
| def close(self): | |
| self.stream.Sync() | |
| self.stream.Close() | |
| if __name__ == "__main__": | |
| comm = MPI.COMM_WORLD | |
| rank = comm.rank | |
| size = comm.size | |
| FORMAT = ( | |
| "%(asctime)s %(rank)s%(name)s - %(levelname)s - " "%(message)s (%(filename)s:%(lineno)d)" | |
| ) | |
| class Formatter(logging.Formatter): | |
| def format(self, record): | |
| record.rank = f"CPU {rank}: " if size > 1 else "" | |
| return super().format(record) | |
| class MPIFilter(logging.Filter): | |
| def filter(self, record): | |
| if rank == 0: | |
| return 1 | |
| else: | |
| return 0 | |
| file_handler = MPIFileHandler("output.log", comm=comm) | |
| file_handler.setLevel(logging.INFO) | |
| # Add filter if you only want to log on rank 0 | |
| # file_handler.addFilter(MPIFilter()) | |
| file_handler.setFormatter(logging.Formatter(FORMAT)) | |
| logging.getLogger().addHandler(file_handler) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment