| 81 | |
| 82 | |
| 83 | def initialize( |
| 84 | ip: str, |
| 85 | port: int, |
| 86 | num_processes: int, |
| 87 | process_id: int, |
| 88 | local_device_ids: Optional[Union[int, Sequence[int]]] = None, |
| 89 | ): |
| 90 | ip = "127.0.0.1" if ip == "localhost" else ip |
| 91 | if global_state.service == None and global_state.client == None: |
| 92 | global_state.initialize(ip, port, num_processes, process_id, local_device_ids) |
| 93 | atexit.register(shutdown) |
| 94 | else: |
| 95 | assert ( |
| 96 | global_state.client != None |
| 97 | ), "global_state.client should not be None if server is created" |
| 98 | if global_state.ip == ip and global_state.port == port: |
| 99 | return |
| 100 | else: |
| 101 | msg = ( |
| 102 | f"xla distribute server/client have been created on {global_state.ip}:{global_state.port}. " |
| 103 | f"so ignore the request to create on {ip}:{port}" |
| 104 | ) |
| 105 | warnings.warn(msg, category=RuntimeWarning) |
| 106 | |
| 107 | |
| 108 | def shutdown(): |