| OLD | NEW |
| 1 # -*- coding: utf-8 -*- |
| 1 # Copyright 2013 Google Inc. All Rights Reserved. | 2 # Copyright 2013 Google Inc. All Rights Reserved. |
| 2 # | 3 # |
| 3 # Licensed under the Apache License, Version 2.0 (the "License"); | 4 # Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 # you may not use this file except in compliance with the License. | 5 # you may not use this file except in compliance with the License. |
| 5 # You may obtain a copy of the License at | 6 # You may obtain a copy of the License at |
| 6 # | 7 # |
| 7 # http://www.apache.org/licenses/LICENSE-2.0 | 8 # http://www.apache.org/licenses/LICENSE-2.0 |
| 8 # | 9 # |
| 9 # Unless required by applicable law or agreed to in writing, software | 10 # Unless required by applicable law or agreed to in writing, software |
| 10 # distributed under the License is distributed on an "AS IS" BASIS, | 11 # distributed under the License is distributed on an "AS IS" BASIS, |
| 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | 12 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 # See the License for the specific language governing permissions and | 13 # See the License for the specific language governing permissions and |
| 13 # limitations under the License. | 14 # limitations under the License. |
| 15 """Utility classes for the parallelism framework.""" |
| 14 | 16 |
| 15 """Utility classes for the parallelism framework.""" | 17 from __future__ import absolute_import |
| 16 | 18 |
| 17 import multiprocessing | 19 import multiprocessing |
| 18 import threading | 20 import threading |
| 19 | 21 |
| 22 |
| 20 class BasicIncrementDict(object): | 23 class BasicIncrementDict(object): |
| 24 """Dictionary meant for storing values for which increment is defined. |
| 25 |
| 26 This handles any values for which the "+" operation is defined (e.g., floats, |
| 27 lists, etc.). This class is neither thread- nor process-safe. |
| 21 """ | 28 """ |
| 22 Dictionary meant for storing any values for which the "+" operation is | 29 |
| 23 defined (e.g., floats, lists, etc.). This class is neither thread- nor | |
| 24 process-safe. | |
| 25 """ | |
| 26 def __init__(self): | 30 def __init__(self): |
| 27 self.dict = {} | 31 self.dict = {} |
| 28 | 32 |
| 29 def get(self, key, default_value=None): | 33 def Get(self, key, default_value=None): |
| 30 return self.dict.get(key, default_value) | 34 return self.dict.get(key, default_value) |
| 31 | 35 |
| 32 def put(self, key, value): | 36 def Put(self, key, value): |
| 33 self.dict[key] = value | 37 self.dict[key] = value |
| 34 | 38 |
| 35 def update(self, key, inc, default_value=0): | 39 def Update(self, key, inc, default_value=0): |
| 36 """ | 40 """Update the stored value associated with the given key. |
| 37 Update the stored value associated with the given key (or the default_value, | 41 |
| 38 if there is no existing value for the key) by performing the equivalent of | 42 Performs the equivalent of |
| 39 self.put(key, self.get(key, default_value) + inc). | 43 self.put(key, self.get(key, default_value) + inc). |
| 44 |
| 45 Args: |
| 46 key: lookup key for the value of the first operand of the "+" operation. |
| 47 inc: Second operand of the "+" operation. |
| 48 default_value: Default value if there is no existing value for the key. |
| 49 |
| 50 Returns: |
| 51 Incremented value. |
| 40 """ | 52 """ |
| 41 val = self.dict.get(key, default_value) + inc | 53 val = self.dict.get(key, default_value) + inc |
| 42 self.dict[key] = val | 54 self.dict[key] = val |
| 43 return val | 55 return val |
| 44 | 56 |
| 45 | 57 |
| 46 class AtomicIncrementDict(BasicIncrementDict): | 58 class AtomicIncrementDict(BasicIncrementDict): |
| 59 """Dictionary meant for storing values for which increment is defined. |
| 60 |
| 61 This handles any values for which the "+" operation is defined (e.g., floats, |
| 62 lists, etc.) in a thread- and process-safe way that allows for atomic get, |
| 63 put, and update. |
| 47 """ | 64 """ |
| 48 Dictionary meant for storing any values for which the "+" operation is | 65 |
| 49 defined (e.g., floats, lists, etc.) in a way that allows for atomic get, put, | 66 def __init__(self, manager): # pylint: disable=super-init-not-called |
| 50 and update in a thread- and process-safe way. | |
| 51 """ | |
| 52 def __init__(self, manager): | |
| 53 self.dict = ThreadAndProcessSafeDict(manager) | 67 self.dict = ThreadAndProcessSafeDict(manager) |
| 54 self.lock = multiprocessing.Lock() | 68 self.lock = multiprocessing.Lock() |
| 55 | 69 |
| 56 def update(self, key, inc, default_value=0): | 70 def Update(self, key, inc, default_value=0): |
| 57 """ | 71 """Atomically update the stored value associated with the given key. |
| 58 Update the stored value associated with the given key (or the default_value, | 72 |
| 59 if there is no existing value for the key) by performing the equivalent of | 73 Performs the atomic equivalent of |
| 60 self.put(key, self.get(key, default_value) + inc) atomically. | 74 self.put(key, self.get(key, default_value) + inc). |
| 75 |
| 76 Args: |
| 77 key: lookup key for the value of the first operand of the "+" operation. |
| 78 inc: Second operand of the "+" operation. |
| 79 default_value: Default value if there is no existing value for the key. |
| 80 |
| 81 Returns: |
| 82 Incremented value. |
| 61 """ | 83 """ |
| 62 with self.lock: | 84 with self.lock: |
| 63 return super(AtomicIncrementDict, self).update(key, inc, default_value) | 85 return super(AtomicIncrementDict, self).Update(key, inc, default_value) |
| 64 | 86 |
| 65 | 87 |
| 66 class ThreadAndProcessSafeDict(object): | 88 class ThreadAndProcessSafeDict(object): |
| 89 """Wraps a multiprocessing.Manager's proxy objects for thread-safety. |
| 90 |
| 91 The proxy objects returned by a manager are process-safe but not necessarily |
| 92 thread-safe, so this class simply wraps their access with a lock for ease of |
| 93 use. Since the objects are process-safe, we can use the more efficient |
| 94 threading Lock. |
| 67 """ | 95 """ |
| 68 The proxy objects returned by a manager are not necessarily thread-safe, so | 96 |
| 69 this class simply wraps their access with a lock for ease of use. They are, | |
| 70 however, process-safe, so we can use the more efficient threading Lock. | |
| 71 """ | |
| 72 def __init__(self, manager): | 97 def __init__(self, manager): |
| 98 """Initializes the thread and process safe dict. |
| 99 |
| 100 Args: |
| 101 manager: Multiprocessing.manager object. |
| 102 """ |
| 73 self.dict = manager.dict() | 103 self.dict = manager.dict() |
| 74 self.lock = threading.Lock() | 104 self.lock = threading.Lock() |
| 75 | 105 |
| 76 def __getitem__(self, key): | 106 def __getitem__(self, key): |
| 77 with self.lock: | 107 with self.lock: |
| 78 return self.dict[key] | 108 return self.dict[key] |
| 79 | 109 |
| 80 def __setitem__(self, key, value): | 110 def __setitem__(self, key, value): |
| 81 with self.lock: | 111 with self.lock: |
| 82 self.dict[key] = value | 112 self.dict[key] = value |
| 83 | 113 |
| 114 # pylint: disable=invalid-name |
| 84 def get(self, key, default_value=None): | 115 def get(self, key, default_value=None): |
| 85 with self.lock: | 116 with self.lock: |
| 86 return self.dict.get(key, default_value) | 117 return self.dict.get(key, default_value) |
| OLD | NEW |