(
self,
ip: str,
port: int,
num_processes: int,
process_id: int,
local_device_ids: Optional[Union[int, Sequence[int]]] = None,
)
| 18 | visible_devices: Optional[str] = "all" |
| 19 | |
| 20 | def initialize( |
| 21 | self, |
| 22 | ip: str, |
| 23 | port: int, |
| 24 | num_processes: int, |
| 25 | process_id: int, |
| 26 | local_device_ids: Optional[Union[int, Sequence[int]]] = None, |
| 27 | ): |
| 28 | coordinator_address = ip + ":" + str(port) |
| 29 | if local_device_ids is None: |
| 30 | local_device_ids = [process_id] |
| 31 | elif isinstance(local_device_ids, int): |
| 32 | local_device_ids = [local_device_ids] |
| 33 | else: |
| 34 | local_device_ids = list(local_device_ids) |
| 35 | |
| 36 | self.ip = ip |
| 37 | self.port = port |
| 38 | self.visible_devices = ",".join(str(x) for x in local_device_ids) |
| 39 | self.process_id = process_id |
| 40 | |
| 41 | if process_id == 0: |
| 42 | if self.service is not None: |
| 43 | raise RuntimeError("distributed.initialize should only be called once.") |
| 44 | self.service = xe.get_distributed_runtime_service( |
| 45 | coordinator_address, num_processes, use_coordination_service=True |
| 46 | ) |
| 47 | |
| 48 | if self.client is not None: |
| 49 | raise RuntimeError("distributed.initialize should only be called once.") |
| 50 | |
| 51 | # Set init_timeout to 5 min to leave time for all the processes to connect |
| 52 | self.client = xe.get_distributed_runtime_client( |
| 53 | coordinator_address, |
| 54 | process_id, |
| 55 | use_coordination_service=True, |
| 56 | init_timeout=300, |
| 57 | ) |
| 58 | self.client.connect() |
| 59 | self.initialize_preemption_sync_manager() |
| 60 | |
| 61 | def shutdown(self): |
| 62 | if self.client: |
no test coverage detected