INNER CODE UNIT · Python
printflock
huggingface/large_language_model_training_playbook · debug/printflock.py:22
def printflock(*args, **kwargs):
"""
This is a wrapper around the built-in Python `print` which calls `flock` before calling
`print` and unlocks it immediately after. This wrapper is useful for when each rank needs to
print a message without getting it interleaved with prints from other ranks.
The lock file is the file this wrapper is defined in.
The output order will be random per rank.
Example:
>>> # assuming 4 GPUs
>>> world_size = dist.get_world_size()
>>> rank = dist.get_rank()
>>> printflock(f"This is a very long message from rank {rank}/{world_size}")
This is a very long message from rank 0/4
This is a very long message from rank 2/4
This is a very long message from rank 3/4
This is a very long message from rank 1/4