We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] Plasma store changes the location of PlasmaObjectNonexistent and PlasmaStoreFull
As ``PlasmaObjectNonexistent`` and ``PlasmaStoreFull`` are moved from ``pyarrow.lib`` into ``pyarrow.plasma`` in 0.15.0, we need to add a try-except block on import.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `mars/worker/storage/sharedstore.py`
Content:
```
1 # Copyright 1999-2018 Alibaba Group Holding Ltd.
2 #
3 # 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 obtain a copy of the License at
6 #
7 #      http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 import logging
16 
17 from ...actors import FunctionActor
18 from ...errors import StorageFull, StorageDataExists
19 from ...utils import calc_data_size
20 
21 logger = logging.getLogger(__name__)
22 
23 
24 class PlasmaKeyMapActor(FunctionActor):
25     @classmethod
26     def default_uid(cls):
27         return 'w:0:' + cls.__name__
28 
29     def __init__(self):
30         super(PlasmaKeyMapActor, self).__init__()
31         self._mapping = dict()
32 
33     def put(self, session_id, chunk_key, obj_id):
34         session_chunk_key = (session_id, chunk_key)
35         if session_chunk_key in self._mapping:
36             raise StorageDataExists(session_chunk_key)
37         self._mapping[session_chunk_key] = obj_id
38 
39     def get(self, session_id, chunk_key):
40         return self._mapping.get((session_id, chunk_key))
41 
42     def delete(self, session_id, chunk_key):
43         try:
44             del self._mapping[(session_id, chunk_key)]
45         except KeyError:
46             pass
47 
48 
49 class PlasmaSharedStore(object):
50     """
51     Wrapper of plasma client for Mars objects
52     """
53     def __init__(self, plasma_client, mapper_ref):
54         from ...serialize.dataserializer import mars_serialize_context
55 
56         self._plasma_client = plasma_client
57         self._actual_size = None
58         self._serialize_context = mars_serialize_context()
59 
60         self._mapper_ref = mapper_ref
61 
62     def get_actual_capacity(self, store_limit):
63         """
64         Get actual capacity of plasma store
65         :return: actual storage size in bytes
66         """
67         if self._actual_size is None:
68             from pyarrow import plasma, lib
69 
70             bufs = []
71             left_size = store_limit
72             total_size = 0
73             alloc_fraction = 0.9
74             while left_size:
75                 allocate_size = int(left_size * alloc_fraction)
76                 if allocate_size < 1 * 1024 ** 2:
77                     break
78 
79                 try:
80                     obj_id = plasma.ObjectID.from_random()
81                     bufs.append(self._plasma_client.create(obj_id, allocate_size))
82                     self._plasma_client.seal(obj_id)
83                     total_size += allocate_size
84                     left_size -= allocate_size
85                     alloc_fraction = 0.9
86                 except lib.PlasmaStoreFull:
87                     alloc_fraction -= 0.1
88                     if alloc_fraction < 1e-6:
89                         break
90             del bufs
91             self._plasma_client.evict(total_size)
92             self._actual_size = total_size
93         return self._actual_size
94 
95     def _new_object_id(self, session_id, data_key):
96         """
97         Calc unique object id for chunks
98         """
99         from pyarrow.plasma import ObjectID
100         while True:
101             new_id = ObjectID.from_random()
102             if not self._plasma_client.contains(new_id):
103                 break
104         self._mapper_ref.put(session_id, data_key, new_id)
105         return new_id
106 
107     def _get_object_id(self, session_id, data_key):
108         obj_id = self._mapper_ref.get(session_id, data_key)
109         if obj_id is None:
110             raise KeyError((session_id, data_key))
111         return obj_id
112 
113     def create(self, session_id, data_key, size):
114         from pyarrow.lib import PlasmaStoreFull
115         obj_id = self._new_object_id(session_id, data_key)
116 
117         try:
118             self._plasma_client.evict(size)
119             buffer = self._plasma_client.create(obj_id, size)
120             return buffer
121         except PlasmaStoreFull:
122             exc_type = PlasmaStoreFull
123             self._mapper_ref.delete(session_id, data_key)
124             logger.warning('Data %s(%d) failed to store to plasma due to StorageFull',
125                            data_key, size)
126         except:  # noqa: E722
127             self._mapper_ref.delete(session_id, data_key)
128             raise
129 
130         if exc_type is PlasmaStoreFull:
131             raise StorageFull(request_size=size, total_size=self._actual_size)
132 
133     def seal(self, session_id, data_key):
134         from pyarrow.lib import PlasmaObjectNonexistent
135         obj_id = self._get_object_id(session_id, data_key)
136         try:
137             self._plasma_client.seal(obj_id)
138         except PlasmaObjectNonexistent:
139             self._mapper_ref.delete(session_id, data_key)
140             raise KeyError((session_id, data_key))
141 
142     def get(self, session_id, data_key):
143         """
144         Get deserialized Mars object from plasma store
145         """
146         from pyarrow.plasma import ObjectNotAvailable
147 
148         obj_id = self._get_object_id(session_id, data_key)
149         obj = self._plasma_client.get(obj_id, serialization_context=self._serialize_context, timeout_ms=10)
150         if obj is ObjectNotAvailable:
151             self._mapper_ref.delete(session_id, data_key)
152             raise KeyError((session_id, data_key))
153         return obj
154 
155     def get_buffer(self, session_id, data_key):
156         """
157         Get raw buffer from plasma store
158         """
159         obj_id = self._get_object_id(session_id, data_key)
160         [buf] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)
161         if buf is None:
162             self._mapper_ref.delete(session_id, data_key)
163             raise KeyError((session_id, data_key))
164         return buf
165 
166     def get_actual_size(self, session_id, data_key):
167         """
168         Get actual size of Mars object from plasma store
169         """
170         buf = None
171         try:
172             obj_id = self._get_object_id(session_id, data_key)
173             [buf] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)
174             if buf is None:
175                 self._mapper_ref.delete(session_id, data_key)
176                 raise KeyError((session_id, data_key))
177             return buf.size
178         finally:
179             del buf
180 
181     def put(self, session_id, data_key, value):
182         """
183         Put a Mars object into plasma store
184         :param session_id: session id
185         :param data_key: chunk key
186         :param value: Mars object to be put
187         """
188         import pyarrow
189         from pyarrow.lib import PlasmaStoreFull
190 
191         data_size = calc_data_size(value)
192 
193         try:
194             obj_id = self._new_object_id(session_id, data_key)
195         except StorageDataExists:
196             obj_id = self._get_object_id(session_id, data_key)
197             if self._plasma_client.contains(obj_id):
198                 logger.debug('Data %s already exists, returning existing', data_key)
199                 [buffer] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)
200                 del value
201                 return buffer
202             else:
203                 logger.warning('Data %s registered but no data found, reconstructed', data_key)
204                 self._mapper_ref.delete(session_id, data_key)
205                 obj_id = self._new_object_id(session_id, data_key)
206 
207         try:
208             serialized = pyarrow.serialize(value, self._serialize_context)
209             del value
210             data_size = serialized.total_bytes
211             try:
212                 buffer = self._plasma_client.create(obj_id, serialized.total_bytes)
213                 stream = pyarrow.FixedSizeBufferWriter(buffer)
214                 stream.set_memcopy_threads(6)
215                 serialized.write_to(stream)
216                 self._plasma_client.seal(obj_id)
217             finally:
218                 del serialized
219             return buffer
220         except PlasmaStoreFull:
221             self._mapper_ref.delete(session_id, data_key)
222             logger.warning('Data %s(%d) failed to store to plasma due to StorageFull',
223                            data_key, data_size)
224             exc = PlasmaStoreFull
225         except:  # noqa: E722
226             self._mapper_ref.delete(session_id, data_key)
227             raise
228 
229         if exc is PlasmaStoreFull:
230             raise StorageFull(request_size=data_size, total_size=self._actual_size)
231 
232     def contains(self, session_id, data_key):
233         """
234         Check if given chunk key exists in current plasma store
235         """
236         try:
237             obj_id = self._get_object_id(session_id, data_key)
238             if self._plasma_client.contains(obj_id):
239                 return True
240             else:
241                 self._mapper_ref.delete(session_id, data_key)
242                 return False
243         except KeyError:
244             return False
245 
246     def delete(self, session_id, data_key):
247         self._mapper_ref.delete(session_id, data_key)
248
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/mars/worker/storage/sharedstore.py b/mars/worker/storage/sharedstore.py
--- a/mars/worker/storage/sharedstore.py
+++ b/mars/worker/storage/sharedstore.py
@@ -18,6 +18,16 @@
 from ...errors import StorageFull, StorageDataExists
 from ...utils import calc_data_size
 
+try:
+    import pyarrow
+    from pyarrow import plasma
+    try:
+        from pyarrow.plasma import PlasmaObjectNonexistent, PlasmaStoreFull
+    except ImportError:
+        from pyarrow.lib import PlasmaObjectNonexistent, PlasmaStoreFull
+except ImportError:  # pragma: no cover
+    pyarrow, plasma, PlasmaObjectNonexistent, PlasmaStoreFull = None, None, None, None
+
 logger = logging.getLogger(__name__)
 
 
@@ -65,8 +75,6 @@
         :return: actual storage size in bytes
         """
         if self._actual_size is None:
-            from pyarrow import plasma, lib
-
             bufs = []
             left_size = store_limit
             total_size = 0
@@ -83,7 +91,7 @@
                     total_size += allocate_size
                     left_size -= allocate_size
                     alloc_fraction = 0.9
-                except lib.PlasmaStoreFull:
+                except PlasmaStoreFull:
                     alloc_fraction -= 0.1
                     if alloc_fraction < 1e-6:
                         break
@@ -96,9 +104,8 @@
         """
         Calc unique object id for chunks
         """
-        from pyarrow.plasma import ObjectID
         while True:
-            new_id = ObjectID.from_random()
+            new_id = plasma.ObjectID.from_random()
             if not self._plasma_client.contains(new_id):
                 break
         self._mapper_ref.put(session_id, data_key, new_id)
@@ -111,7 +118,6 @@
         return obj_id
 
     def create(self, session_id, data_key, size):
-        from pyarrow.lib import PlasmaStoreFull
         obj_id = self._new_object_id(session_id, data_key)
 
         try:
@@ -131,7 +137,6 @@
             raise StorageFull(request_size=size, total_size=self._actual_size)
 
     def seal(self, session_id, data_key):
-        from pyarrow.lib import PlasmaObjectNonexistent
         obj_id = self._get_object_id(session_id, data_key)
         try:
             self._plasma_client.seal(obj_id)
@@ -143,11 +148,9 @@
         """
         Get deserialized Mars object from plasma store
         """
-        from pyarrow.plasma import ObjectNotAvailable
-
         obj_id = self._get_object_id(session_id, data_key)
         obj = self._plasma_client.get(obj_id, serialization_context=self._serialize_context, timeout_ms=10)
-        if obj is ObjectNotAvailable:
+        if obj is plasma.ObjectNotAvailable:
             self._mapper_ref.delete(session_id, data_key)
             raise KeyError((session_id, data_key))
         return obj
@@ -185,9 +188,6 @@
         :param data_key: chunk key
         :param value: Mars object to be put
         """
-        import pyarrow
-        from pyarrow.lib import PlasmaStoreFull
-
         data_size = calc_data_size(value)
 
         try:

{"golden_diff": "diff --git a/mars/worker/storage/sharedstore.py b/mars/worker/storage/sharedstore.py\n--- a/mars/worker/storage/sharedstore.py\n+++ b/mars/worker/storage/sharedstore.py\n@@ -18,6 +18,16 @@\n from ...errors import StorageFull, StorageDataExists\n from ...utils import calc_data_size\n \n+try:\n+    import pyarrow\n+    from pyarrow import plasma\n+    try:\n+        from pyarrow.plasma import PlasmaObjectNonexistent, PlasmaStoreFull\n+    except ImportError:\n+        from pyarrow.lib import PlasmaObjectNonexistent, PlasmaStoreFull\n+except ImportError:  # pragma: no cover\n+    pyarrow, plasma, PlasmaObjectNonexistent, PlasmaStoreFull = None, None, None, None\n+\n logger = logging.getLogger(__name__)\n \n \n@@ -65,8 +75,6 @@\n         :return: actual storage size in bytes\n         \"\"\"\n         if self._actual_size is None:\n-            from pyarrow import plasma, lib\n-\n             bufs = []\n             left_size = store_limit\n             total_size = 0\n@@ -83,7 +91,7 @@\n                     total_size += allocate_size\n                     left_size -= allocate_size\n                     alloc_fraction = 0.9\n-                except lib.PlasmaStoreFull:\n+                except PlasmaStoreFull:\n                     alloc_fraction -= 0.1\n                     if alloc_fraction < 1e-6:\n                         break\n@@ -96,9 +104,8 @@\n         \"\"\"\n         Calc unique object id for chunks\n         \"\"\"\n-        from pyarrow.plasma import ObjectID\n         while True:\n-            new_id = ObjectID.from_random()\n+            new_id = plasma.ObjectID.from_random()\n             if not self._plasma_client.contains(new_id):\n                 break\n         self._mapper_ref.put(session_id, data_key, new_id)\n@@ -111,7 +118,6 @@\n         return obj_id\n \n     def create(self, session_id, data_key, size):\n-        from pyarrow.lib import PlasmaStoreFull\n         obj_id = self._new_object_id(session_id, data_key)\n \n         try:\n@@ -131,7 +137,6 @@\n             raise StorageFull(request_size=size, total_size=self._actual_size)\n \n     def seal(self, session_id, data_key):\n-        from pyarrow.lib import PlasmaObjectNonexistent\n         obj_id = self._get_object_id(session_id, data_key)\n         try:\n             self._plasma_client.seal(obj_id)\n@@ -143,11 +148,9 @@\n         \"\"\"\n         Get deserialized Mars object from plasma store\n         \"\"\"\n-        from pyarrow.plasma import ObjectNotAvailable\n-\n         obj_id = self._get_object_id(session_id, data_key)\n         obj = self._plasma_client.get(obj_id, serialization_context=self._serialize_context, timeout_ms=10)\n-        if obj is ObjectNotAvailable:\n+        if obj is plasma.ObjectNotAvailable:\n             self._mapper_ref.delete(session_id, data_key)\n             raise KeyError((session_id, data_key))\n         return obj\n@@ -185,9 +188,6 @@\n         :param data_key: chunk key\n         :param value: Mars object to be put\n         \"\"\"\n-        import pyarrow\n-        from pyarrow.lib import PlasmaStoreFull\n-\n         data_size = calc_data_size(value)\n \n         try:\n", "issue": "[BUG] Plasma store changes the location of PlasmaObjectNonexistent and PlasmaStoreFull\nAs ``PlasmaObjectNonexistent`` and ``PlasmaStoreFull`` are moved from ``pyarrow.lib`` into ``pyarrow.plasma`` in 0.15.0, we need to add a try-except block on import.\n", "before_files": [{"content": "# Copyright 1999-2018 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\n\nfrom ...actors import FunctionActor\nfrom ...errors import StorageFull, StorageDataExists\nfrom ...utils import calc_data_size\n\nlogger = logging.getLogger(__name__)\n\n\nclass PlasmaKeyMapActor(FunctionActor):\n    @classmethod\n    def default_uid(cls):\n        return 'w:0:' + cls.__name__\n\n    def __init__(self):\n        super(PlasmaKeyMapActor, self).__init__()\n        self._mapping = dict()\n\n    def put(self, session_id, chunk_key, obj_id):\n        session_chunk_key = (session_id, chunk_key)\n        if session_chunk_key in self._mapping:\n            raise StorageDataExists(session_chunk_key)\n        self._mapping[session_chunk_key] = obj_id\n\n    def get(self, session_id, chunk_key):\n        return self._mapping.get((session_id, chunk_key))\n\n    def delete(self, session_id, chunk_key):\n        try:\n            del self._mapping[(session_id, chunk_key)]\n        except KeyError:\n            pass\n\n\nclass PlasmaSharedStore(object):\n    \"\"\"\n    Wrapper of plasma client for Mars objects\n    \"\"\"\n    def __init__(self, plasma_client, mapper_ref):\n        from ...serialize.dataserializer import mars_serialize_context\n\n        self._plasma_client = plasma_client\n        self._actual_size = None\n        self._serialize_context = mars_serialize_context()\n\n        self._mapper_ref = mapper_ref\n\n    def get_actual_capacity(self, store_limit):\n        \"\"\"\n        Get actual capacity of plasma store\n        :return: actual storage size in bytes\n        \"\"\"\n        if self._actual_size is None:\n            from pyarrow import plasma, lib\n\n            bufs = []\n            left_size = store_limit\n            total_size = 0\n            alloc_fraction = 0.9\n            while left_size:\n                allocate_size = int(left_size * alloc_fraction)\n                if allocate_size < 1 * 1024 ** 2:\n                    break\n\n                try:\n                    obj_id = plasma.ObjectID.from_random()\n                    bufs.append(self._plasma_client.create(obj_id, allocate_size))\n                    self._plasma_client.seal(obj_id)\n                    total_size += allocate_size\n                    left_size -= allocate_size\n                    alloc_fraction = 0.9\n                except lib.PlasmaStoreFull:\n                    alloc_fraction -= 0.1\n                    if alloc_fraction < 1e-6:\n                        break\n            del bufs\n            self._plasma_client.evict(total_size)\n            self._actual_size = total_size\n        return self._actual_size\n\n    def _new_object_id(self, session_id, data_key):\n        \"\"\"\n        Calc unique object id for chunks\n        \"\"\"\n        from pyarrow.plasma import ObjectID\n        while True:\n            new_id = ObjectID.from_random()\n            if not self._plasma_client.contains(new_id):\n                break\n        self._mapper_ref.put(session_id, data_key, new_id)\n        return new_id\n\n    def _get_object_id(self, session_id, data_key):\n        obj_id = self._mapper_ref.get(session_id, data_key)\n        if obj_id is None:\n            raise KeyError((session_id, data_key))\n        return obj_id\n\n    def create(self, session_id, data_key, size):\n        from pyarrow.lib import PlasmaStoreFull\n        obj_id = self._new_object_id(session_id, data_key)\n\n        try:\n            self._plasma_client.evict(size)\n            buffer = self._plasma_client.create(obj_id, size)\n            return buffer\n        except PlasmaStoreFull:\n            exc_type = PlasmaStoreFull\n            self._mapper_ref.delete(session_id, data_key)\n            logger.warning('Data %s(%d) failed to store to plasma due to StorageFull',\n                           data_key, size)\n        except:  # noqa: E722\n            self._mapper_ref.delete(session_id, data_key)\n            raise\n\n        if exc_type is PlasmaStoreFull:\n            raise StorageFull(request_size=size, total_size=self._actual_size)\n\n    def seal(self, session_id, data_key):\n        from pyarrow.lib import PlasmaObjectNonexistent\n        obj_id = self._get_object_id(session_id, data_key)\n        try:\n            self._plasma_client.seal(obj_id)\n        except PlasmaObjectNonexistent:\n            self._mapper_ref.delete(session_id, data_key)\n            raise KeyError((session_id, data_key))\n\n    def get(self, session_id, data_key):\n        \"\"\"\n        Get deserialized Mars object from plasma store\n        \"\"\"\n        from pyarrow.plasma import ObjectNotAvailable\n\n        obj_id = self._get_object_id(session_id, data_key)\n        obj = self._plasma_client.get(obj_id, serialization_context=self._serialize_context, timeout_ms=10)\n        if obj is ObjectNotAvailable:\n            self._mapper_ref.delete(session_id, data_key)\n            raise KeyError((session_id, data_key))\n        return obj\n\n    def get_buffer(self, session_id, data_key):\n        \"\"\"\n        Get raw buffer from plasma store\n        \"\"\"\n        obj_id = self._get_object_id(session_id, data_key)\n        [buf] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)\n        if buf is None:\n            self._mapper_ref.delete(session_id, data_key)\n            raise KeyError((session_id, data_key))\n        return buf\n\n    def get_actual_size(self, session_id, data_key):\n        \"\"\"\n        Get actual size of Mars object from plasma store\n        \"\"\"\n        buf = None\n        try:\n            obj_id = self._get_object_id(session_id, data_key)\n            [buf] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)\n            if buf is None:\n                self._mapper_ref.delete(session_id, data_key)\n                raise KeyError((session_id, data_key))\n            return buf.size\n        finally:\n            del buf\n\n    def put(self, session_id, data_key, value):\n        \"\"\"\n        Put a Mars object into plasma store\n        :param session_id: session id\n        :param data_key: chunk key\n        :param value: Mars object to be put\n        \"\"\"\n        import pyarrow\n        from pyarrow.lib import PlasmaStoreFull\n\n        data_size = calc_data_size(value)\n\n        try:\n            obj_id = self._new_object_id(session_id, data_key)\n        except StorageDataExists:\n            obj_id = self._get_object_id(session_id, data_key)\n            if self._plasma_client.contains(obj_id):\n                logger.debug('Data %s already exists, returning existing', data_key)\n                [buffer] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)\n                del value\n                return buffer\n            else:\n                logger.warning('Data %s registered but no data found, reconstructed', data_key)\n                self._mapper_ref.delete(session_id, data_key)\n                obj_id = self._new_object_id(session_id, data_key)\n\n        try:\n            serialized = pyarrow.serialize(value, self._serialize_context)\n            del value\n            data_size = serialized.total_bytes\n            try:\n                buffer = self._plasma_client.create(obj_id, serialized.total_bytes)\n                stream = pyarrow.FixedSizeBufferWriter(buffer)\n                stream.set_memcopy_threads(6)\n                serialized.write_to(stream)\n                self._plasma_client.seal(obj_id)\n            finally:\n                del serialized\n            return buffer\n        except PlasmaStoreFull:\n            self._mapper_ref.delete(session_id, data_key)\n            logger.warning('Data %s(%d) failed to store to plasma due to StorageFull',\n                           data_key, data_size)\n            exc = PlasmaStoreFull\n        except:  # noqa: E722\n            self._mapper_ref.delete(session_id, data_key)\n            raise\n\n        if exc is PlasmaStoreFull:\n            raise StorageFull(request_size=data_size, total_size=self._actual_size)\n\n    def contains(self, session_id, data_key):\n        \"\"\"\n        Check if given chunk key exists in current plasma store\n        \"\"\"\n        try:\n            obj_id = self._get_object_id(session_id, data_key)\n            if self._plasma_client.contains(obj_id):\n                return True\n            else:\n                self._mapper_ref.delete(session_id, data_key)\n                return False\n        except KeyError:\n            return False\n\n    def delete(self, session_id, data_key):\n        self._mapper_ref.delete(session_id, data_key)\n", "path": "mars/worker/storage/sharedstore.py"}], "after_files": [{"content": "# Copyright 1999-2018 Alibaba Group Holding Ltd.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\n\nfrom ...actors import FunctionActor\nfrom ...errors import StorageFull, StorageDataExists\nfrom ...utils import calc_data_size\n\ntry:\n    import pyarrow\n    from pyarrow import plasma\n    try:\n        from pyarrow.plasma import PlasmaObjectNonexistent, PlasmaStoreFull\n    except ImportError:\n        from pyarrow.lib import PlasmaObjectNonexistent, PlasmaStoreFull\nexcept ImportError:  # pragma: no cover\n    pyarrow, plasma, PlasmaObjectNonexistent, PlasmaStoreFull = None, None, None, None\n\nlogger = logging.getLogger(__name__)\n\n\nclass PlasmaKeyMapActor(FunctionActor):\n    @classmethod\n    def default_uid(cls):\n        return 'w:0:' + cls.__name__\n\n    def __init__(self):\n        super(PlasmaKeyMapActor, self).__init__()\n        self._mapping = dict()\n\n    def put(self, session_id, chunk_key, obj_id):\n        session_chunk_key = (session_id, chunk_key)\n        if session_chunk_key in self._mapping:\n            raise StorageDataExists(session_chunk_key)\n        self._mapping[session_chunk_key] = obj_id\n\n    def get(self, session_id, chunk_key):\n        return self._mapping.get((session_id, chunk_key))\n\n    def delete(self, session_id, chunk_key):\n        try:\n            del self._mapping[(session_id, chunk_key)]\n        except KeyError:\n            pass\n\n\nclass PlasmaSharedStore(object):\n    \"\"\"\n    Wrapper of plasma client for Mars objects\n    \"\"\"\n    def __init__(self, plasma_client, mapper_ref):\n        from ...serialize.dataserializer import mars_serialize_context\n\n        self._plasma_client = plasma_client\n        self._actual_size = None\n        self._serialize_context = mars_serialize_context()\n\n        self._mapper_ref = mapper_ref\n\n    def get_actual_capacity(self, store_limit):\n        \"\"\"\n        Get actual capacity of plasma store\n        :return: actual storage size in bytes\n        \"\"\"\n        if self._actual_size is None:\n            bufs = []\n            left_size = store_limit\n            total_size = 0\n            alloc_fraction = 0.9\n            while left_size:\n                allocate_size = int(left_size * alloc_fraction)\n                if allocate_size < 1 * 1024 ** 2:\n                    break\n\n                try:\n                    obj_id = plasma.ObjectID.from_random()\n                    bufs.append(self._plasma_client.create(obj_id, allocate_size))\n                    self._plasma_client.seal(obj_id)\n                    total_size += allocate_size\n                    left_size -= allocate_size\n                    alloc_fraction = 0.9\n                except PlasmaStoreFull:\n                    alloc_fraction -= 0.1\n                    if alloc_fraction < 1e-6:\n                        break\n            del bufs\n            self._plasma_client.evict(total_size)\n            self._actual_size = total_size\n        return self._actual_size\n\n    def _new_object_id(self, session_id, data_key):\n        \"\"\"\n        Calc unique object id for chunks\n        \"\"\"\n        while True:\n            new_id = plasma.ObjectID.from_random()\n            if not self._plasma_client.contains(new_id):\n                break\n        self._mapper_ref.put(session_id, data_key, new_id)\n        return new_id\n\n    def _get_object_id(self, session_id, data_key):\n        obj_id = self._mapper_ref.get(session_id, data_key)\n        if obj_id is None:\n            raise KeyError((session_id, data_key))\n        return obj_id\n\n    def create(self, session_id, data_key, size):\n        obj_id = self._new_object_id(session_id, data_key)\n\n        try:\n            self._plasma_client.evict(size)\n            buffer = self._plasma_client.create(obj_id, size)\n            return buffer\n        except PlasmaStoreFull:\n            exc_type = PlasmaStoreFull\n            self._mapper_ref.delete(session_id, data_key)\n            logger.warning('Data %s(%d) failed to store to plasma due to StorageFull',\n                           data_key, size)\n        except:  # noqa: E722\n            self._mapper_ref.delete(session_id, data_key)\n            raise\n\n        if exc_type is PlasmaStoreFull:\n            raise StorageFull(request_size=size, total_size=self._actual_size)\n\n    def seal(self, session_id, data_key):\n        obj_id = self._get_object_id(session_id, data_key)\n        try:\n            self._plasma_client.seal(obj_id)\n        except PlasmaObjectNonexistent:\n            self._mapper_ref.delete(session_id, data_key)\n            raise KeyError((session_id, data_key))\n\n    def get(self, session_id, data_key):\n        \"\"\"\n        Get deserialized Mars object from plasma store\n        \"\"\"\n        obj_id = self._get_object_id(session_id, data_key)\n        obj = self._plasma_client.get(obj_id, serialization_context=self._serialize_context, timeout_ms=10)\n        if obj is plasma.ObjectNotAvailable:\n            self._mapper_ref.delete(session_id, data_key)\n            raise KeyError((session_id, data_key))\n        return obj\n\n    def get_buffer(self, session_id, data_key):\n        \"\"\"\n        Get raw buffer from plasma store\n        \"\"\"\n        obj_id = self._get_object_id(session_id, data_key)\n        [buf] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)\n        if buf is None:\n            self._mapper_ref.delete(session_id, data_key)\n            raise KeyError((session_id, data_key))\n        return buf\n\n    def get_actual_size(self, session_id, data_key):\n        \"\"\"\n        Get actual size of Mars object from plasma store\n        \"\"\"\n        buf = None\n        try:\n            obj_id = self._get_object_id(session_id, data_key)\n            [buf] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)\n            if buf is None:\n                self._mapper_ref.delete(session_id, data_key)\n                raise KeyError((session_id, data_key))\n            return buf.size\n        finally:\n            del buf\n\n    def put(self, session_id, data_key, value):\n        \"\"\"\n        Put a Mars object into plasma store\n        :param session_id: session id\n        :param data_key: chunk key\n        :param value: Mars object to be put\n        \"\"\"\n        data_size = calc_data_size(value)\n\n        try:\n            obj_id = self._new_object_id(session_id, data_key)\n        except StorageDataExists:\n            obj_id = self._get_object_id(session_id, data_key)\n            if self._plasma_client.contains(obj_id):\n                logger.debug('Data %s already exists, returning existing', data_key)\n                [buffer] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)\n                del value\n                return buffer\n            else:\n                logger.warning('Data %s registered but no data found, reconstructed', data_key)\n                self._mapper_ref.delete(session_id, data_key)\n                obj_id = self._new_object_id(session_id, data_key)\n\n        try:\n            serialized = pyarrow.serialize(value, self._serialize_context)\n            del value\n            data_size = serialized.total_bytes\n            try:\n                buffer = self._plasma_client.create(obj_id, serialized.total_bytes)\n                stream = pyarrow.FixedSizeBufferWriter(buffer)\n                stream.set_memcopy_threads(6)\n                serialized.write_to(stream)\n                self._plasma_client.seal(obj_id)\n            finally:\n                del serialized\n            return buffer\n        except PlasmaStoreFull:\n            self._mapper_ref.delete(session_id, data_key)\n            logger.warning('Data %s(%d) failed to store to plasma due to StorageFull',\n                           data_key, data_size)\n            exc = PlasmaStoreFull\n        except:  # noqa: E722\n            self._mapper_ref.delete(session_id, data_key)\n            raise\n\n        if exc is PlasmaStoreFull:\n            raise StorageFull(request_size=data_size, total_size=self._actual_size)\n\n    def contains(self, session_id, data_key):\n        \"\"\"\n        Check if given chunk key exists in current plasma store\n        \"\"\"\n        try:\n            obj_id = self._get_object_id(session_id, data_key)\n            if self._plasma_client.contains(obj_id):\n                return True\n            else:\n                self._mapper_ref.delete(session_id, data_key)\n                return False\n        except KeyError:\n            return False\n\n    def delete(self, session_id, data_key):\n        self._mapper_ref.delete(session_id, data_key)\n", "path": "mars/worker/storage/sharedstore.py"}]}