# How to do CUDA traces callback only for uncaught OOM errors?

**URL:** <https://discuss.pytorch.org/t/how-to-do-cuda-traces-callback-only-for-uncaught-oom-errors/183983>\
**Category:** C++\
**Created:** [July 11, 2023, 6:21pm UTC](https://discuss.pytorch.org/t/how-to-do-cuda-traces-callback-only-for-uncaught-oom-errors/183983 "2023-07-11T18:21:34Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![oren.learn](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/oren.learn/32/61586_2.png) [@oren.learn](https://discuss.pytorch.org/u/oren.learn)\
**Post date:** [July 11, 2023, 6:21pm UTC](https://discuss.pytorch.org/t/how-to-do-cuda-traces-callback-only-for-uncaught-oom-errors/183983/1 "2023-07-11T18:21:34Z")

</div>

From this blog post by @zdevito, we see that we can add a callback every time it OOMs but this callback happens for caught OOMs and retries. Is there a way to do callback only when it is a uncaught OOM?

```python
def oom_observer(device, alloc, device_alloc, device_free):
    # snapshot right after an OOM happened
    print('saving allocated state during OOM')
    snapshot = torch.cuda.memory._snapshot()
    dump(snapshot, open('oom_snapshot.pickle', 'wb'))

torch._C._cuda_attach_out_of_memory_observer(oom_observer)

```

looking at the C++ code, i don’t see an obvious way to add a flag for only when it is uncaught OOMs [cuda/module.cpp](https://github.com/pytorch/pytorch/blob/eb03af44eec31ab11b8fcef4607d44a16024062a/torch/csrc/cuda/Module.cpp#L811C36-L811C61) [CUDACachingAllocator.cpp](https://github.com/pytorch/pytorch/blob/eb03af44eec31ab11b8fcef4607d44a16024062a/c10/cuda/CUDACachingAllocator.cpp#L1304-L1306)

---

<div class="post-metadata">

**Author:** ![zdevito](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/zdevito/32/62124_2.png) [@zdevito](https://discuss.pytorch.org/u/zdevito)\
**Post date:** [July 11, 2023, 9:14pm UTC](https://discuss.pytorch.org/t/how-to-do-cuda-traces-callback-only-for-uncaught-oom-errors/183983/2 "2023-07-11T21:14:13Z")

</div>

There isn’t an easy way to do it because the callback needs to run when the OOM happens but before the stack unwinds, otherwise a lot of memory will be released before we can take the snapshot. It is possible to record multiple snapshots and then only save the last one. Each snapshot save will take some time but it will eventually get the actual OOM snapshot:

```auto
def oom_observer(device, alloc, device_alloc, device_free):
    # snapshot right after an OOM happened
    global snapshot
    snapshot = torch.cuda.memory._snapshot()
    
torch._C._cuda_attach_out_of_memory_observer(oom_observer)
try:
  <application code>
except torch.cuda.OutOfMemoryError:
  print('saving allocated state during OOM')  
  dump(snapshot, open('oom_snapshot.pickle', 'wb'))

```