From 8185f210e6be2d5de11b4f0213040e07f8d36118 Mon Sep 17 00:00:00 2001 From: Abhishek Enaguthi Date: Wed, 15 Jul 2026 18:21:46 -0700 Subject: [PATCH] Fix get_shm_prefix follower race until locals content is ready Attach-only FileNotFound retries still fail when the shm name is visible but packed locals are not yet consistent across processes. Retry attach plus unpack/match with jitter, then emit a clear permanent error. Adds a regression test for delayed content visibility (issue #901). --- streaming/base/shared/prefix.py | 69 ++++++++++++++++++++++++--------- tests/test_shared.py | 42 +++++++++++++++++++- 2 files changed, 92 insertions(+), 19 deletions(-) diff --git a/streaming/base/shared/prefix.py b/streaming/base/shared/prefix.py index 9decae99d..5b4cf8aaa 100644 --- a/streaming/base/shared/prefix.py +++ b/streaming/base/shared/prefix.py @@ -22,6 +22,15 @@ from streaming.base.world import World +class _SharedMemoryNotReady(Exception): + """Follower attached before leader shm content was fully visible. + + Dist barriers do not wait for OS SharedMemory propagation. Attach may + succeed while the packed locals buffer is still empty or inconsistent; + callers should retry with backoff instead of failing the job. + """ + + def _each_prefix_int() -> Iterator[int]: """Get each possible prefix int to check in order. @@ -230,28 +239,52 @@ def get_shm_prefix(streams_local: list[str], dist.barrier() # Non-local leaders go next, searching for match. + # Retry both name visibility (FileNotFoundError) and content readiness + # (_SharedMemoryNotReady): barrier sync does not imply OS shm propagation + # (see https://github.com/mosaicml/streaming/issues/901). if not world.is_local_leader: + name = _get_path(prefix_int, LOCALS) - @retry_decorator(exc_class=FileNotFoundError, + @retry_decorator(exc_class=(FileNotFoundError, _SharedMemoryNotReady), num_attempts=100, initial_backoff=TICK, - max_jitter=0.0) - def _attach_to_shm() -> SharedMemory: - """Attach to shared memory created by local leader.""" - name = _get_path(prefix_int, LOCALS) - return SharedMemory(name, False) + max_jitter=TICK) + def _attach_and_validate_locals() -> SharedMemory: + """Attach to leader shm and wait until packed locals are consistent.""" + follower_shm = SharedMemory(name, False) + try: + their_locals, their_prefix_int = _unpack_locals(bytes(follower_shm.buf)) + except (ValueError, UnicodeDecodeError, IndexError, OverflowError) as err: + raise _SharedMemoryNotReady( + f'Shared memory prefix={prefix_int} attached but locals are not readable yet' + ) from err + if streams_local != their_locals or prefix_int != their_prefix_int: + # Treat mismatch as not-ready during the race window; after retries + # exhaust we re-read once below to emit a clear permanent error. + raise _SharedMemoryNotReady( + f'Shared memory prefix={prefix_int} content not yet consistent with leader') + return follower_shm try: - shm = _attach_to_shm() - except FileNotFoundError: - raise RuntimeError(f'Internal error: shared memory prefix={prefix_int} was not ' + - f'registered by local leader. This may be because you specified ' + - f'different ``local`` parameters from different ranks.') - - their_locals, their_prefix_int = _unpack_locals(bytes(shm.buf)) - if streams_local != their_locals or prefix_int != their_prefix_int: - raise RuntimeError(f'Internal error: shared memory registered does not match ' + - f'local leader as streams_local or prefix_int not match. ' + - f'local leader: {their_locals} and {their_prefix_int}. ' + - f'expected: {streams_local} and {prefix_int}.') + shm = _attach_and_validate_locals() + except (FileNotFoundError, _SharedMemoryNotReady): + # Final probe: distinguish "never created" vs permanent locals mismatch. + try: + shm = SharedMemory(name, False) + except FileNotFoundError as err: + raise RuntimeError( + f'Internal error: shared memory prefix={prefix_int} was not ' + + f'registered by local leader. This may be because you specified ' + + f'different ``local`` parameters from different ranks.') from err + try: + their_locals, their_prefix_int = _unpack_locals(bytes(shm.buf)) + except (ValueError, UnicodeDecodeError, IndexError, OverflowError) as err: + raise RuntimeError( + f'Internal error: shared memory prefix={prefix_int} was registered but ' + + f'locals never became readable.') from err + if streams_local != their_locals or prefix_int != their_prefix_int: + raise RuntimeError(f'Internal error: shared memory registered does not match ' + + f'local leader as streams_local or prefix_int not match. ' + + f'local leader: {their_locals} and {their_prefix_int}. ' + + f'expected: {streams_local} and {prefix_int}.') return prefix_int, shm # pyright: ignore diff --git a/tests/test_shared.py b/tests/test_shared.py index 882261e85..f4a9f2bdd 100644 --- a/tests/test_shared.py +++ b/tests/test_shared.py @@ -16,7 +16,8 @@ from streaming.base.constant import LOCALS from streaming.base.shared import SharedArray, get_shm_prefix from streaming.base.shared.memory import SharedMemory -from streaming.base.shared.prefix import _check_and_find +from streaming.base.shared.prefix import (_SharedMemoryNotReady, _check_and_find, _get_path, + _pack_locals, _unpack_locals) from streaming.base.util import clean_stale_shared_memory from streaming.base.world import World from tests.common.utils import convert_to_mds @@ -196,6 +197,45 @@ def test_shared_memory_permission_error(mock_shared_memory_class: MagicMock): assert next_prefix == 1 +def test_unpack_locals_roundtrip(): + packed = _pack_locals(['/tmp/a', '/tmp/b'], 7) + assert _unpack_locals(packed) == (['/tmp/a', '/tmp/b'], 7) + + +def test_get_shm_prefix_follower_retries_until_locals_readable(local_remote_dir: tuple[str, str]): + """Non-leaders must retry when shm attaches before packed locals are consistent (#901).""" + local, _ = local_remote_dir + clean_stale_shared_memory() + + prefix_int = 0 + name = _get_path(prefix_int, LOCALS) + packed = _pack_locals([local], prefix_int) + leader_shm = SharedMemory(name, True, len(packed)) + leader_shm.buf[:len(packed)] = packed + + real_unpack = _unpack_locals + attempts = {'n': 0} + + def flaky_unpack(data: bytes) -> tuple[list[str], int]: + attempts['n'] += 1 + # Simulate OS visibility of the name before the packed payload is consistent. + if attempts['n'] <= 2: + raise ValueError('simulated incomplete shm payload') + return real_unpack(data) + + follower_world = MagicMock() + follower_world.is_local_leader = False + + with patch('streaming.base.shared.prefix._check_and_find_retrying', return_value=prefix_int), \ + patch('streaming.base.shared.prefix._unpack_locals', side_effect=flaky_unpack): + got_prefix, follower_shm = get_shm_prefix([local], [None], follower_world) + + assert got_prefix == prefix_int + assert attempts['n'] >= 3 + assert _unpack_locals(bytes(follower_shm.buf)) == ([local], prefix_int) + assert issubclass(_SharedMemoryNotReady, Exception) + + # Global counter to track attach attempts (per process) attach_attempts = 0