We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Allow Mars to use pyarrow 0.17.0
As pyarrow 0.17.0 is released with https://github.com/apache/arrow/pull/6624 merged, we can now use the latest version of pyarrow.

--- 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-2020 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 
20 try:
21     import pyarrow
22     from pyarrow import plasma
23     try:
24         from pyarrow.plasma import PlasmaObjectNonexistent, PlasmaStoreFull
25     except ImportError:
26         from pyarrow.lib import PlasmaObjectNonexistent, PlasmaStoreFull
27 except ImportError:  # pragma: no cover
28     pyarrow, plasma, PlasmaObjectNonexistent, PlasmaStoreFull = None, None, None, None
29 
30 logger = logging.getLogger(__name__)
31 
32 
33 class PlasmaKeyMapActor(FunctionActor):
34     @classmethod
35     def default_uid(cls):
36         return 'w:0:' + cls.__name__
37 
38     def __init__(self):
39         super().__init__()
40         self._mapping = dict()
41 
42     def put(self, session_id, chunk_key, obj_id):
43         session_chunk_key = (session_id, chunk_key)
44         if session_chunk_key in self._mapping:
45             raise StorageDataExists(session_chunk_key)
46         self._mapping[session_chunk_key] = obj_id
47 
48     def get(self, session_id, chunk_key):
49         return self._mapping.get((session_id, chunk_key))
50 
51     def delete(self, session_id, chunk_key):
52         try:
53             del self._mapping[(session_id, chunk_key)]
54         except KeyError:
55             pass
56 
57     def batch_delete(self, session_id, chunk_keys):
58         for k in chunk_keys:
59             self.delete(session_id, k)
60 
61 
62 class PlasmaSharedStore(object):
63     """
64     Wrapper of plasma client for Mars objects
65     """
66     def __init__(self, plasma_client, mapper_ref):
67         from ...serialize.dataserializer import mars_serialize_context
68 
69         self._plasma_client = plasma_client
70         self._actual_size = None
71         self._serialize_context = mars_serialize_context()
72 
73         self._mapper_ref = mapper_ref
74         self._pool = mapper_ref.ctx.threadpool(1)
75 
76     def get_actual_capacity(self, store_limit):
77         """
78         Get actual capacity of plasma store
79         :return: actual storage size in bytes
80         """
81         if self._actual_size is None:
82             bufs = []
83             left_size = store_limit
84             total_size = 0
85             alloc_fraction = 0.9
86             while left_size:
87                 allocate_size = int(left_size * alloc_fraction)
88                 if allocate_size < 1 * 1024 ** 2:
89                     break
90 
91                 try:
92                     obj_id = plasma.ObjectID.from_random()
93                     bufs.append(self._plasma_client.create(obj_id, allocate_size))
94                     self._plasma_client.seal(obj_id)
95                     total_size += allocate_size
96                     left_size -= allocate_size
97                     alloc_fraction = 0.9
98                 except PlasmaStoreFull:
99                     alloc_fraction -= 0.1
100                     if alloc_fraction < 1e-6:
101                         break
102             del bufs
103             self._plasma_client.evict(total_size)
104             self._actual_size = total_size
105         return self._actual_size
106 
107     def _new_object_id(self, session_id, data_key):
108         """
109         Calc unique object id for chunks
110         """
111         while True:
112             new_id = plasma.ObjectID.from_random()
113             if not self._plasma_client.contains(new_id):
114                 break
115         self._mapper_ref.put(session_id, data_key, new_id)
116         return new_id
117 
118     def _get_object_id(self, session_id, data_key):
119         obj_id = self._mapper_ref.get(session_id, data_key)
120         if obj_id is None:
121             raise KeyError((session_id, data_key))
122         return obj_id
123 
124     def create(self, session_id, data_key, size):
125         obj_id = self._new_object_id(session_id, data_key)
126 
127         try:
128             self._plasma_client.evict(size)
129             buffer = self._plasma_client.create(obj_id, size)
130             return buffer
131         except PlasmaStoreFull:
132             exc_type = PlasmaStoreFull
133             self._mapper_ref.delete(session_id, data_key)
134             logger.warning('Data %s(%d) failed to store to plasma due to StorageFull',
135                            data_key, size)
136         except:  # noqa: E722
137             self._mapper_ref.delete(session_id, data_key)
138             raise
139 
140         if exc_type is PlasmaStoreFull:
141             raise StorageFull(request_size=size, capacity=self._actual_size,
142                               affected_keys=[data_key])
143 
144     def seal(self, session_id, data_key):
145         obj_id = self._get_object_id(session_id, data_key)
146         try:
147             self._plasma_client.seal(obj_id)
148         except PlasmaObjectNonexistent:
149             self._mapper_ref.delete(session_id, data_key)
150             raise KeyError((session_id, data_key))
151 
152     def get(self, session_id, data_key):
153         """
154         Get deserialized Mars object from plasma store
155         """
156         obj_id = self._get_object_id(session_id, data_key)
157         obj = self._plasma_client.get(obj_id, serialization_context=self._serialize_context, timeout_ms=10)
158         if obj is plasma.ObjectNotAvailable:
159             self._mapper_ref.delete(session_id, data_key)
160             raise KeyError((session_id, data_key))
161         return obj
162 
163     def get_buffer(self, session_id, data_key):
164         """
165         Get raw buffer from plasma store
166         """
167         obj_id = self._get_object_id(session_id, data_key)
168         [buf] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)
169         if buf is None:
170             self._mapper_ref.delete(session_id, data_key)
171             raise KeyError((session_id, data_key))
172         return buf
173 
174     def get_actual_size(self, session_id, data_key):
175         """
176         Get actual size of Mars object from plasma store
177         """
178         buf = None
179         try:
180             obj_id = self._get_object_id(session_id, data_key)
181             [buf] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)
182             if buf is None:
183                 self._mapper_ref.delete(session_id, data_key)
184                 raise KeyError((session_id, data_key))
185             return buf.size
186         finally:
187             del buf
188 
189     def put(self, session_id, data_key, value):
190         """
191         Put a Mars object into plasma store
192         :param session_id: session id
193         :param data_key: chunk key
194         :param value: Mars object to be put
195         """
196         data_size = None
197 
198         try:
199             obj_id = self._new_object_id(session_id, data_key)
200         except StorageDataExists:
201             obj_id = self._get_object_id(session_id, data_key)
202             if self._plasma_client.contains(obj_id):
203                 logger.debug('Data %s already exists, returning existing', data_key)
204                 [buffer] = self._plasma_client.get_buffers([obj_id], timeout_ms=10)
205                 del value
206                 return buffer
207             else:
208                 logger.warning('Data %s registered but no data found, reconstructed', data_key)
209                 self._mapper_ref.delete(session_id, data_key)
210                 obj_id = self._new_object_id(session_id, data_key)
211 
212         try:
213             serialized = pyarrow.serialize(value, self._serialize_context)
214             del value
215             data_size = serialized.total_bytes
216             try:
217                 buffer = self._plasma_client.create(obj_id, serialized.total_bytes)
218                 stream = pyarrow.FixedSizeBufferWriter(buffer)
219                 stream.set_memcopy_threads(6)
220                 self._pool.submit(serialized.write_to, stream).result()
221                 self._plasma_client.seal(obj_id)
222             finally:
223                 del serialized
224             return buffer
225         except PlasmaStoreFull:
226             self._mapper_ref.delete(session_id, data_key)
227             logger.warning('Data %s(%d) failed to store to plasma due to StorageFull',
228                            data_key, data_size)
229             exc = PlasmaStoreFull
230         except:  # noqa: E722
231             self._mapper_ref.delete(session_id, data_key)
232             raise
233 
234         if exc is PlasmaStoreFull:
235             raise StorageFull(request_size=data_size, capacity=self._actual_size,
236                               affected_keys=[data_key])
237 
238     def contains(self, session_id, data_key):
239         """
240         Check if given chunk key exists in current plasma store
241         """
242         try:
243             obj_id = self._get_object_id(session_id, data_key)
244             if self._plasma_client.contains(obj_id):
245                 return True
246             else:
247                 self._mapper_ref.delete(session_id, data_key)
248                 return False
249         except KeyError:
250             return False
251 
252     def delete(self, session_id, data_key):
253         self._mapper_ref.delete(session_id, data_key)
254 
255     def batch_delete(self, session_id, data_keys):
256         self._mapper_ref.batch_delete(session_id, data_keys)
257
```

Path: `mars/_version.py`
Content:
```
1 # Copyright 1999-2020 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 subprocess
16 import os
17 
18 version_info = (0, 4, 0, 'b2')
19 _num_index = max(idx if isinstance(v, int) else 0
20                  for idx, v in enumerate(version_info))
21 __version__ = '.'.join(map(str, version_info[:_num_index + 1])) + \
22               ''.join(version_info[_num_index + 1:])
23 
24 
25 def _get_cmd_results(pkg_root, cmd):
26     proc = subprocess.Popen(cmd, cwd=pkg_root, stdout=subprocess.PIPE)
27     proc.wait()
28     if proc.returncode == 0:
29         s = proc.stdout.read().decode()
30         proc.stdout.close()
31         return s
32 
33 
34 def get_git_info():
35     pkg_root = os.path.dirname(os.path.abspath(__file__))
36     git_root = os.path.join(os.path.dirname(pkg_root), '.git')
37 
38     if os.path.exists(git_root):
39         commit_hash = _get_cmd_results(pkg_root, ['git', 'rev-parse', 'HEAD']).strip()
40         if not commit_hash:
41             return
42         branches = _get_cmd_results(pkg_root, ['git', 'branch']).splitlines(False)
43         commit_ref = None
44         for branch in branches:
45             if not branch.startswith('*'):
46                 continue
47             striped = branch[1:].strip()
48             if not striped.startswith('('):
49                 commit_ref = striped
50             else:
51                 _, commit_ref = striped.rsplit(' ', 1)
52                 commit_ref = commit_ref.rstrip(')')
53         if commit_ref is None:
54             return
55         return commit_hash, commit_ref
56     else:
57         branch_file = os.path.join(pkg_root, '.git-branch')
58         if not os.path.exists(branch_file):
59             return
60         with open(branch_file, 'r') as bf:
61             bf_content = bf.read().strip()
62             return bf_content.split(None, 1)
63
```
--- 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/_version.py b/mars/_version.py
--- a/mars/_version.py
+++ b/mars/_version.py
@@ -15,7 +15,7 @@
 import subprocess
 import os
 
-version_info = (0, 4, 0, 'b2')
+version_info = (0, 4, 0, 'rc1')
 _num_index = max(idx if isinstance(v, int) else 0
                  for idx, v in enumerate(version_info))
 __version__ = '.'.join(map(str, version_info[:_num_index + 1])) + \
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
@@ -21,11 +21,19 @@
     import pyarrow
     from pyarrow import plasma
     try:
-        from pyarrow.plasma import PlasmaObjectNonexistent, PlasmaStoreFull
-    except ImportError:
-        from pyarrow.lib import PlasmaObjectNonexistent, PlasmaStoreFull
+        from pyarrow.plasma import PlasmaObjectNotFound
+    except ImportError:  # pragma: no cover
+        try:
+            from pyarrow.plasma import PlasmaObjectNonexistent as PlasmaObjectNotFound
+        except ImportError:
+            from pyarrow.lib import PlasmaObjectNonexistent as PlasmaObjectNotFound
+
+    try:
+        from pyarrow.plasma import PlasmaStoreFull
+    except ImportError:  # pragma: no cover
+        from pyarrow.lib import PlasmaStoreFull
 except ImportError:  # pragma: no cover
-    pyarrow, plasma, PlasmaObjectNonexistent, PlasmaStoreFull = None, None, None, None
+    pyarrow, plasma, PlasmaObjectNotFound, PlasmaStoreFull = None, None, None, None
 
 logger = logging.getLogger(__name__)
 
@@ -145,7 +153,7 @@
         obj_id = self._get_object_id(session_id, data_key)
         try:
             self._plasma_client.seal(obj_id)
-        except PlasmaObjectNonexistent:
+        except PlasmaObjectNotFound:
             self._mapper_ref.delete(session_id, data_key)
             raise KeyError((session_id, data_key))

{"golden_diff": "diff --git a/mars/_version.py b/mars/_version.py\n--- a/mars/_version.py\n+++ b/mars/_version.py\n@@ -15,7 +15,7 @@\n import subprocess\n import os\n \n-version_info = (0, 4, 0, 'b2')\n+version_info = (0, 4, 0, 'rc1')\n _num_index = max(idx if isinstance(v, int) else 0\n                  for idx, v in enumerate(version_info))\n __version__ = '.'.join(map(str, version_info[:_num_index + 1])) + \\\ndiff --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@@ -21,11 +21,19 @@\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+        from pyarrow.plasma import PlasmaObjectNotFound\n+    except ImportError:  # pragma: no cover\n+        try:\n+            from pyarrow.plasma import PlasmaObjectNonexistent as PlasmaObjectNotFound\n+        except ImportError:\n+            from pyarrow.lib import PlasmaObjectNonexistent as PlasmaObjectNotFound\n+\n+    try:\n+        from pyarrow.plasma import PlasmaStoreFull\n+    except ImportError:  # pragma: no cover\n+        from pyarrow.lib import PlasmaStoreFull\n except ImportError:  # pragma: no cover\n-    pyarrow, plasma, PlasmaObjectNonexistent, PlasmaStoreFull = None, None, None, None\n+    pyarrow, plasma, PlasmaObjectNotFound, PlasmaStoreFull = None, None, None, None\n \n logger = logging.getLogger(__name__)\n \n@@ -145,7 +153,7 @@\n         obj_id = self._get_object_id(session_id, data_key)\n         try:\n             self._plasma_client.seal(obj_id)\n-        except PlasmaObjectNonexistent:\n+        except PlasmaObjectNotFound:\n             self._mapper_ref.delete(session_id, data_key)\n             raise KeyError((session_id, data_key))\n", "issue": "Allow Mars to use pyarrow 0.17.0\nAs pyarrow 0.17.0 is released with https://github.com/apache/arrow/pull/6624 merged, we can now use the latest version of pyarrow.\n", "before_files": [{"content": "# Copyright 1999-2020 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\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().__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    def batch_delete(self, session_id, chunk_keys):\n        for k in chunk_keys:\n            self.delete(session_id, k)\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        self._pool = mapper_ref.ctx.threadpool(1)\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, capacity=self._actual_size,\n                              affected_keys=[data_key])\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 = None\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                self._pool.submit(serialized.write_to, stream).result()\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, capacity=self._actual_size,\n                              affected_keys=[data_key])\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\n    def batch_delete(self, session_id, data_keys):\n        self._mapper_ref.batch_delete(session_id, data_keys)\n", "path": "mars/worker/storage/sharedstore.py"}, {"content": "# Copyright 1999-2020 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 subprocess\nimport os\n\nversion_info = (0, 4, 0, 'b2')\n_num_index = max(idx if isinstance(v, int) else 0\n                 for idx, v in enumerate(version_info))\n__version__ = '.'.join(map(str, version_info[:_num_index + 1])) + \\\n              ''.join(version_info[_num_index + 1:])\n\n\ndef _get_cmd_results(pkg_root, cmd):\n    proc = subprocess.Popen(cmd, cwd=pkg_root, stdout=subprocess.PIPE)\n    proc.wait()\n    if proc.returncode == 0:\n        s = proc.stdout.read().decode()\n        proc.stdout.close()\n        return s\n\n\ndef get_git_info():\n    pkg_root = os.path.dirname(os.path.abspath(__file__))\n    git_root = os.path.join(os.path.dirname(pkg_root), '.git')\n\n    if os.path.exists(git_root):\n        commit_hash = _get_cmd_results(pkg_root, ['git', 'rev-parse', 'HEAD']).strip()\n        if not commit_hash:\n            return\n        branches = _get_cmd_results(pkg_root, ['git', 'branch']).splitlines(False)\n        commit_ref = None\n        for branch in branches:\n            if not branch.startswith('*'):\n                continue\n            striped = branch[1:].strip()\n            if not striped.startswith('('):\n                commit_ref = striped\n            else:\n                _, commit_ref = striped.rsplit(' ', 1)\n                commit_ref = commit_ref.rstrip(')')\n        if commit_ref is None:\n            return\n        return commit_hash, commit_ref\n    else:\n        branch_file = os.path.join(pkg_root, '.git-branch')\n        if not os.path.exists(branch_file):\n            return\n        with open(branch_file, 'r') as bf:\n            bf_content = bf.read().strip()\n            return bf_content.split(None, 1)\n", "path": "mars/_version.py"}], "after_files": [{"content": "# Copyright 1999-2020 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\n\ntry:\n    import pyarrow\n    from pyarrow import plasma\n    try:\n        from pyarrow.plasma import PlasmaObjectNotFound\n    except ImportError:  # pragma: no cover\n        try:\n            from pyarrow.plasma import PlasmaObjectNonexistent as PlasmaObjectNotFound\n        except ImportError:\n            from pyarrow.lib import PlasmaObjectNonexistent as PlasmaObjectNotFound\n\n    try:\n        from pyarrow.plasma import PlasmaStoreFull\n    except ImportError:  # pragma: no cover\n        from pyarrow.lib import PlasmaStoreFull\nexcept ImportError:  # pragma: no cover\n    pyarrow, plasma, PlasmaObjectNotFound, 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().__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    def batch_delete(self, session_id, chunk_keys):\n        for k in chunk_keys:\n            self.delete(session_id, k)\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        self._pool = mapper_ref.ctx.threadpool(1)\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, capacity=self._actual_size,\n                              affected_keys=[data_key])\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 PlasmaObjectNotFound:\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 = None\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                self._pool.submit(serialized.write_to, stream).result()\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, capacity=self._actual_size,\n                              affected_keys=[data_key])\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\n    def batch_delete(self, session_id, data_keys):\n        self._mapper_ref.batch_delete(session_id, data_keys)\n", "path": "mars/worker/storage/sharedstore.py"}, {"content": "# Copyright 1999-2020 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 subprocess\nimport os\n\nversion_info = (0, 4, 0, 'rc1')\n_num_index = max(idx if isinstance(v, int) else 0\n                 for idx, v in enumerate(version_info))\n__version__ = '.'.join(map(str, version_info[:_num_index + 1])) + \\\n              ''.join(version_info[_num_index + 1:])\n\n\ndef _get_cmd_results(pkg_root, cmd):\n    proc = subprocess.Popen(cmd, cwd=pkg_root, stdout=subprocess.PIPE)\n    proc.wait()\n    if proc.returncode == 0:\n        s = proc.stdout.read().decode()\n        proc.stdout.close()\n        return s\n\n\ndef get_git_info():\n    pkg_root = os.path.dirname(os.path.abspath(__file__))\n    git_root = os.path.join(os.path.dirname(pkg_root), '.git')\n\n    if os.path.exists(git_root):\n        commit_hash = _get_cmd_results(pkg_root, ['git', 'rev-parse', 'HEAD']).strip()\n        if not commit_hash:\n            return\n        branches = _get_cmd_results(pkg_root, ['git', 'branch']).splitlines(False)\n        commit_ref = None\n        for branch in branches:\n            if not branch.startswith('*'):\n                continue\n            striped = branch[1:].strip()\n            if not striped.startswith('('):\n                commit_ref = striped\n            else:\n                _, commit_ref = striped.rsplit(' ', 1)\n                commit_ref = commit_ref.rstrip(')')\n        if commit_ref is None:\n            return\n        return commit_hash, commit_ref\n    else:\n        branch_file = os.path.join(pkg_root, '.git-branch')\n        if not os.path.exists(branch_file):\n            return\n        with open(branch_file, 'r') as bf:\n            bf_content = bf.read().strip()\n            return bf_content.split(None, 1)\n", "path": "mars/_version.py"}]}