(self, cuda_visible_devices=None)
| 117 | os.environ.update(rank_zero_info) |
| 118 | |
| 119 | def __init__(self, cuda_visible_devices=None) -> None: |
| 120 | # construct a meta from envrionment variable. Note that the import must be inside the class because it is executed remotely |
| 121 | import os |
| 122 | world_size = int(os.environ['WORLD_SIZE']) |
| 123 | rank = int(os.environ['RANK']) |
| 124 | self._rank = rank |
| 125 | self._world_size = world_size |
| 126 | |
| 127 | master_addr = os.environ["MASTER_ADDR"] |
| 128 | master_port = os.environ["MASTER_PORT"] |
| 129 | |
| 130 | local_world_size = int(os.getenv("LOCAL_WORLD_SIZE", "1")) |
| 131 | local_rank = int(os.getenv("LOCAL_RANK", "0")) |
| 132 | |
| 133 | store = { |
| 134 | '_world_size': world_size, |
| 135 | '_rank': rank, |
| 136 | '_local_world_size': local_world_size, |
| 137 | '_local_rank': local_rank, |
| 138 | '_master_addr': master_addr, |
| 139 | '_master_port': master_port |
| 140 | } |
| 141 | if cuda_visible_devices is not None: |
| 142 | store['_cuda_visible_devices'] = cuda_visible_devices |
| 143 | |
| 144 | meta = WorkerMeta(store=store) |
| 145 | self._configure_with_meta(meta=meta) |
| 146 | |
| 147 | def _configure_with_meta(self, meta: WorkerMeta): |
| 148 | """ |
nothing calls this directly
no test coverage detected