mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
9a8f71f23e
Codemod done with https://gist.github.com/ezyang/2e8b0463cdc6be278478495b23ff0530 with assistance from ChatGPT. Signed-off-by: Edward Z. Yang <ezyang@meta.com> Pull Request resolved: https://github.com/pytorch/pytorch/pull/98697 Approved by: https://github.com/voznesenskym
38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
from contextlib import contextmanager
|
|
from typing import cast
|
|
import logging
|
|
from . import api
|
|
from . import TensorPipeAgent
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
@contextmanager
|
|
def _group_membership_management(store, name, is_join):
|
|
token_key = "RpcGroupManagementToken"
|
|
join_or_leave = "join" if is_join else "leave"
|
|
my_token = f"Token_for_{name}_{join_or_leave}"
|
|
while True:
|
|
# Retrieve token from store to signal start of rank join/leave critical section
|
|
returned = store.compare_set(token_key, "", my_token).decode()
|
|
if returned == my_token:
|
|
# Yield to the function this context manager wraps
|
|
yield
|
|
# Finished, now exit and release token
|
|
# Update from store to signal end of rank join/leave critical section
|
|
store.set(token_key, "")
|
|
# Other will wait for this token to be set before they execute
|
|
store.set(my_token, "Done")
|
|
break
|
|
else:
|
|
# Store will wait for the token to be released
|
|
try:
|
|
store.wait([returned])
|
|
except RuntimeError:
|
|
logger.error("Group membership token %s timed out waiting for %s to be released.", my_token, returned)
|
|
raise
|
|
|
|
def _update_group_membership(worker_info, my_devices, reverse_device_map, is_join):
|
|
agent = cast(TensorPipeAgent, api._get_current_rpc_agent())
|
|
ret = agent._update_group_membership(worker_info, my_devices, reverse_device_map, is_join)
|
|
return ret
|