Race Condition in base.shared.get_shm_prefix (possible general SharedMemory issue)

[nqgl]

Environment

• OS: Ubuntu
• Hardware (GPU, or instance type): 4xH200, 4xH100, 8xH100 (runpod instances)

Issue

Race condition in get_shm_prefix causes both the following errors to get raised due to OS not having finished synchronizing shared memory

streaming/streaming/base/shared/prefix.py

Lines 234 to 246 in 223be8f

	try:
	shm = SharedMemory(name, False)
	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}.')

Related: #717 #332 #824 #884 possibly #767

Pasting the contents of my comment from #824 (this contains most of the context and info)

I've been running into the same issue, but I'm only using a single StreamingDataset. I think it's a race condition with the SharedMemory that gets created by the leader here and then accessed by non-leader processes in the next if block:

streaming/streaming/base/shared/prefix.py

Lines 222 to 230 in 223be8f

	if world.is_local_leader:
	name = _get_path(prefix_int, LOCALS)
	data = _pack_locals(streams_local, prefix_int)
	shm = SharedMemory(name, True, len(data))
	shm.buf[:len(data)] = data
	if dist.is_available() and dist.is_initialized():
	dist.barrier()

I believe this race condition happens because there is some latency for the operating system to fully synchronize shared memory across different processes. The barrier doesn't fix this, because the leader isn't doing anything to wait for the OS-level propagation to happen.

The reason why I think this is the cause is I was able to fix the issue by inserting the following sleep call (at line 230) to give time to synchronize:

    if world.is_local_leader:
        name = _get_path(prefix_int, LOCALS)
        data = _pack_locals(streams_local, prefix_int)
        shm = SharedMemory(name, True, len(data))
        shm.buf[: len(data)] = data

    if dist.is_available() and dist.is_initialized():
        dist.barrier()
    sleep(1) # this sleep gives time for the OS to update the shm on all processes

More ideal solutions to this would not sleep here and instead would just wait for the OS to propagate to all processes. We could also retry the non-local-leaders block up to some timeout if it fails.

Edit to add:
I may have been hitting the same exceptions as @schopra8 but different reasons, as in my case the error was happening with a single StreamingDataset and commenting out the dist.destroy_process_group() had no impact.
However, I think it's still decently likely they're related issues -- the second dataset or the process group destruction may have been increasing the latency of shared memory propagation.

Also:
This is on Linux with Streaming 0.12.0. Downgrading to 0.9.0 did not affect the issue for me.

Note on a possible general SharedMemory issue

There are other open issues having to do with similar issues happening in places other than get_shm_prefix. I don't have time to look at these so I'm not sure, but it seems likely that they all have a common cause if other code is using SharedMemory similarly to how it is used in get_shm_prefix: creating and writing to SharedMemory does not block until it has synchronized to all processes. In get_shm_prefix it is used as though it is blocking and I suspect it may be used the same way in other parts of the codebase causing similar race conditions. If so, making writes to SharedMemory block until the data are fully propagated would solve multiple issues at once.

[XiaohanZhangCMU]

@nqgl thanks for flagging it. Yes, I have noticed that get_shm_prefix does not work properly on certain OS. The change in 0.12.0 was mostly a patch to enable an internal use case. We didn't see such issue when testing internally, so I cannot reproduce. Also, occasionally dist.barrier() does not "block" on the desired command, which is very rare but happens.

What you proposed sounds reasonable to me. Can you help create a PR that properly fix the issue? I can help you with the CI and get it merged in (sorry we are really short of hands this quarter).

[nqgl]

Thank you for the reply!
Do you have a suggestion for how we could best wait until the SharedMemory has been synchronized? I don't know a good way to check for that off the top of my head but if I did I might but up for making a PR to fix this.

[XiaohanZhangCMU]

@nqgl we also don't have a good idea since torch.distributed is not so reliable.
I think the backoff and retry logic you proposed is solid enough. It may not cover all cases, but 90% is good enough!

If you go with that approach, consider leveraging streaming's retry operator and adjust the jitter to find a conservative value that's safe for the SharedMemory context to propagate.

[somay-jalan]

hey, did anyone come up with a fix? my code did not work with sleep(1)
I am trying to run nvidia/nemo with the streaming and getting stuck in this error.

[nqgl]

@somay-jalan
What error are you getting?
Are you on multi gpu/multi node?

Do you have reason to expect slow updating of the filesystem on the instances you're using? If so, you could try increasing the time in the sleep to like 10 seconds.

It's been a while since I hit this issue and we ended up not using Streaming, but I still have some memories about the issue and I did end up with a workaround to this particular problem so I'm happy to see if I still remember enough to help you with this.

[somay-jalan]

@nqgl
I am getting the follwoing error
[rank0]: Traceback (most recent call last):
[rank0]: File "/home/somay/coom/main.py", line 15, in
[rank0]: trainer.train()
[rank0]: File "/home/somay/coom/coom/train.py", line 249, in train
[rank0]: self.start_training()
[rank0]: File "/home/somay/coom/coom/train.py", line 233, in start_training
[rank0]: llm.train(
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/nemo/collections/llm/api.py", line 130, in train
[rank0]: trainer.fit(model, datamodule=data)
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 538, in fit
[rank0]: call._call_and_handle_interrupt(
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 46, in _call_and_handle_interrupt
[rank0]: return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/lightning/pytorch/strategies/launchers/subprocess_script.py", line 105, in launch
[rank0]: return function(*args, **kwargs)
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 574, in _fit_impl
[rank0]: self._run(model, ckpt_path=ckpt_path)
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/lightning/pytorch/trainer/trainer.py", line 943, in _run
[rank0]: call._call_setup_hook(self) # allow user to set up LightningModule in accelerator environment
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 102, in _call_setup_hook
[rank0]: _call_lightning_datamodule_hook(trainer, "setup", stage=fn)
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/lightning/pytorch/trainer/call.py", line 189, in _call_lightning_datamodule_hook
[rank0]: return fn(*args, **kwargs)
[rank0]: File "/home/somay/coom/coom/streaming_datamodule/eka_datamodule.py", line 104, in setup
[rank0]: self._train_ds = StreamingDataset(
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/streaming/base/dataset.py", line 541, in init
[rank0]: self._shm_prefix_int, self._locals_shm = get_shm_prefix(streams_local, streams_remote,
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/streaming/base/shared/prefix.py", line 225, in get_shm_prefix
[rank0]: shm = SharedMemory(name, True, len(data))
[rank0]: File "/home/somay/eka_env/lib/python3.10/site-packages/streaming/base/shared/memory.py", line 41, in init
[rank0]: shm = BuiltinSharedMemory(name, create, size)
[rank0]: File "/opt/conda/lib/python3.10/multiprocessing/shared_memory.py", line 104, in init
[rank0]: self._fd = _posixshmem.shm_open(
[rank0]: FileExistsError: [Errno 17] File exists: '/000003_locals'

I am not sure why this does not start with 000000 and 000003, I checked my /dev/shm, there is only 000003_locals no 0,1,2
I am using a 2xA100 setup
Increasing the time to 10 did not fix this.
Also what did you end up using?

[somay-jalan]

hey so i was able to fix the error for myself,
`
class SharedMemory:
"""Improved quiet implementation of shared memory.

  Args:
      name (str, optional): A unique shared memory block name. Defaults to ``None``.
      create (bool, optional): Creates a new shared memory block or attaches to an existing
          shared memory block. Defaults to ``None``.
      size (int, optional): A size of a shared memory block. Defaults to ``0``.
      auto_cleanup (bool, optional): Register atexit handler for cleanup or not. Defaults to
          ``True``.
  """

  def __init__(self,
               name: Optional[str] = None,
               create: Optional[bool] = None,
               size: int = 0,
               auto_cleanup: bool = True):
      self.created_shms = []
      self.opened_shms = []
      shm = None
      # save the original register tracker function
      original_rtracker_reg = resource_tracker.register

      try:
          # if create is True:
          #     # Creates a new shared memory block
          #     shm = BuiltinSharedMemory(name, create, size)
          #     self.created_shms.append(shm)
          if create is False:
              # Avoid tracking shared memory resources in a process who attaches to an existing
              # shared memory block because the process who created the shared memory is
              # responsible for destroying the shared memory block.
              resource_tracker.register = self.fix_register
              # Attaches to an existing shared memory block
              shm = BuiltinSharedMemory(name, create, size)
              self.opened_shms.append(shm)
          else:
              try:
                  # Creates a new shared memory block
                  shm = BuiltinSharedMemory(name, True, size)
                  self.created_shms.append(shm)
              except FileExistsError:
                  sleep(TICK)
                  resource_tracker.register = self.fix_register
                  # Attaches to an existing shared memory block
                  shm = BuiltinSharedMemory(name, False, size)
                  self.opened_shms.append(shm)
          self.shm = shm
      finally:
          resource_tracker.register = original_rtracker_reg

      if auto_cleanup:
          # atexit handler doesn't get called if the program is killed by a signal not
          # handled by python or when os.exit() is called or for any python internal fatal error.
          atexit.register(self.cleanup)

`

in the above class in base/shared/memeory.py, I commented out the if statement on line 39 as it seemed redundent and the the whole pipeline worked after that. I have no idea if this is an actual fix or just works for my case as I don't have complete knowledge as to how this library works but this works for me.

[XiaohanZhangCMU]

@somay-jalan that's a smart finding! I am also not sure why this redundant condition was there in the first place, but it does feel like wrong.
Would you mind creating a PR with your fix?

[somay-jalan]

yes for sure, I will create a PR with my fix.