(message, ranks=None)
| 84 | |
| 85 | |
| 86 | def get_dist_msg(message, ranks=None): |
| 87 | from deepspeed import comm as dist |
| 88 | """Return a message with rank prefix when one of following conditions is met: |
| 89 | |
| 90 | + not dist.is_initialized() |
| 91 | + dist.get_rank() in ranks if ranks is not None or ranks = [-1] |
| 92 | |
| 93 | If neither is met, `None` is returned. |
| 94 | |
| 95 | Example: "hello" => "[Rank 0] hello" |
| 96 | |
| 97 | Args: |
| 98 | message (str) |
| 99 | ranks (list) |
| 100 | """ |
| 101 | should_log = not dist.is_initialized() |
| 102 | ranks = ranks or [] |
| 103 | my_rank = dist.get_rank() if dist.is_initialized() else -1 |
| 104 | if ranks and not should_log: |
| 105 | should_log = ranks[0] == -1 |
| 106 | should_log = should_log or (my_rank in set(ranks)) |
| 107 | if should_log: |
| 108 | return "[Rank {}] {}".format(my_rank, message) |
| 109 | else: |
| 110 | return None |
| 111 | |
| 112 | |
| 113 | def log_dist(message, ranks=None, level=logging.INFO): |
no test coverage detected
searching dependent graphs…