projects
/
python_utils.git
/ blobdiff
commit
grep
author
committer
pickaxe
?
search:
re
summary
|
shortlog
|
log
|
commit
|
commitdiff
|
tree
raw
|
inline
| side by side
Fix a bug in sanity check image.
[python_utils.git]
/
collect
/
shared_dict.py
diff --git
a/collect/shared_dict.py
b/collect/shared_dict.py
index 0d8e7c2f7a36aa5ddb7c54c72aecddbf56df71c3..f4ec9143ab80b08b0ca819c4bca62cf20362fb66 100644
(file)
--- a/
collect/shared_dict.py
+++ b/
collect/shared_dict.py
@@
-30,14
+30,14
@@
This class is based on https://github.com/luizalabs/shared-memory-dict
import pickle
from contextlib import contextmanager
from functools import wraps
import pickle
from contextlib import contextmanager
from functools import wraps
-from multiprocessing import
shared_memory, RLock
+from multiprocessing import
RLock, shared_memory
from typing import (
Any,
Dict,
Generator,
from typing import (
Any,
Dict,
Generator,
- KeysView,
ItemsView,
Iterator,
ItemsView,
Iterator,
+ KeysView,
Optional,
ValuesView,
)
Optional,
ValuesView,
)
@@
-74,6
+74,7
@@
class SharedDict(object):
super().__init__()
self.name = name
self._serializer = PickleSerializer()
super().__init__()
self.name = name
self._serializer = PickleSerializer()
+ assert size_bytes is None or size_bytes > 0
self.shared_memory = self._get_or_create_memory_block(name, size_bytes)
self._ensure_memory_initialization()
self.lock = RLock()
self.shared_memory = self._get_or_create_memory_block(name, size_bytes)
self._ensure_memory_initialization()
self.lock = RLock()
@@
-89,13
+90,11
@@
class SharedDict(object):
try:
return shared_memory.SharedMemory(name=name)
except FileNotFoundError:
try:
return shared_memory.SharedMemory(name=name)
except FileNotFoundError:
- assert size_bytes
+ assert size_bytes
is not None
return shared_memory.SharedMemory(name=name, create=True, size=size_bytes)
def _ensure_memory_initialization(self):
return shared_memory.SharedMemory(name=name, create=True, size=size_bytes)
def _ensure_memory_initialization(self):
- memory_is_empty = (
- bytes(self.shared_memory.buf).split(SharedDict.NULL_BYTE, 1)[0] == b''
- )
+ memory_is_empty = bytes(self.shared_memory.buf).split(SharedDict.NULL_BYTE, 1)[0] == b''
if memory_is_empty:
self.clear()
if memory_is_empty:
self.clear()
@@
-107,7
+106,8
@@
class SharedDict(object):
def cleanup(self) -> None:
if not hasattr(self, 'shared_memory'):
return
def cleanup(self) -> None:
if not hasattr(self, 'shared_memory'):
return
- self.shared_memory.unlink()
+ with SharedDict.MPLOCK:
+ self.shared_memory.unlink()
def clear(self) -> None:
self._save_memory({})
def clear(self) -> None:
self._save_memory({})