mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Log the key used for barrier

Browse Source 
PiperOrigin-RevId: 463741926
This commit is contained in:
Yash Katariya 2022-07-27 20:25:07 -07:00 committed by jax authors
parent 27655af6b9
commit 0dbf492cec
1 changed files with 5 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
jax/experimental/gda_serialization/serialization.py
View File

@ -346,13 +346,16 @@ class AsyncManager:
# All processes will wait at the barrier. When all processes are at the
# barrier, the barrier will be satisfied. If not, then it will timeout.
self._client.wait_at_barrier(_get_key(self._count), self._timeout_in_ms)
key_for_barrier = _get_key(self._count)
logging.info('Key used for barrier is %s for process %s',
key_for_barrier, current_process)
self._client.wait_at_barrier(key_for_barrier, self._timeout_in_ms)
logging.info('Finished waiting at barrier for process %s',
current_process)
if current_process == 0:
self._on_commit_callback()
self._client.key_value_set(_get_key(self._count), _CHECKPOINT_SUCCESS)
self._client.key_value_set(key_for_barrier, _CHECKPOINT_SUCCESS)
except Exception as e:
self._exception = e