Skip to content

Commit 856c070

Browse files
committed
Improve error handling in SharedTensor.close() and drop()
Changed exception handling in SharedTensor methods: - close(): Changed from silently logging a warning to properly logging an error with the exception message. This provides better visibility when shared memory cleanup fails, while still not raising exceptions that could cause issues in cleanup paths (like __del__). - drop(): Fixed exception chaining to use 'raise ... from e' for better error traceability and to comply with flake8 B904. Test plan: - Ran pytest tests/unit_tests/util/test_shared_tensor.py - 64 tests passed (1 pre-existing segfault in test_multiprocess_bidirectional)
1 parent db1394e commit 856c070

File tree

1 file changed

+9
-5
lines changed

1 file changed

+9
-5
lines changed

src/forge/util/_shared_tensor.py

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,18 +6,20 @@
66

77
from __future__ import annotations
88

9+
import logging
10+
911
import uuid
1012
from dataclasses import dataclass
1113
from multiprocessing import shared_memory
1214
from typing import Optional, Tuple, Union
13-
import logging
1415

1516
import numpy as np
1617
import torch
1718

1819
logger = logging.getLogger(__name__)
1920
logger.setLevel(logging.INFO)
2021

22+
2123
@dataclass
2224
class SharedTensorHandle:
2325
shm_name: str
@@ -373,8 +375,8 @@ def close(self):
373375

374376
try:
375377
self._shm.close()
376-
except Exception:
377-
logger.warn("Error closing shared memory: ", self._shm_name)
378+
except Exception as e:
379+
logger.error(f"Error closing shared memory {self._shm_name}: {e}")
378380

379381
def drop(self):
380382
"""
@@ -396,8 +398,10 @@ def drop(self):
396398
# Then unlink
397399
try:
398400
self._shm.unlink()
399-
except Exception:
400-
raise RuntimeError("Error unlinking shared memory: ", self._shm_name)
401+
except Exception as e:
402+
raise RuntimeError(
403+
f"Error unlinking shared memory {self._shm_name}: {e}"
404+
) from e
401405

402406
@property
403407
def is_closed(self) -> bool:

0 commit comments

Comments
 (0)