Repository: mars-project/mars

Issue:
[BUG] ray web session can't pickle lambda
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**

The code
``` python
df2 = md.DataFrame(
    mt.random.rand(10_0000, 4, chunk_size=1_0000),
    columns=list('abcd'))
df2.apply(lambda x: x).execute()
```
will raise the following exception
```python
---------------------------------------------------------------------------
PicklingError                             Traceback (most recent call last)
<ipython-input-3-78f183c944b1> in <module>
----> 1 df2.apply(lambda x: x).execute()

~/miniconda3/lib/python3.7/site-packages/mars/core/entity/tileables.py in execute(self, session, **kw)
    430 
    431     def execute(self, session=None, **kw):
--> 432         result = self.data.execute(session=session, **kw)
    433         if isinstance(result, TILEABLE_TYPE):
    434             return self

~/miniconda3/lib/python3.7/site-packages/mars/core/entity/executable.py in execute(self, session, **kw)
    102 
    103         session = _get_session(self, session)
--> 104         return execute(self, session=session, **kw)
    105 
    106     def _check_session(self,

~/miniconda3/lib/python3.7/site-packages/mars/deploy/oscar/session.py in execute(tileable, session, wait, new_session_kwargs, show_progress, progress_update_interval, *tileables, **kwargs)
   1314                            show_progress=show_progress,
   1315                            progress_update_interval=progress_update_interval,
-> 1316                            **kwargs)
   1317 
   1318 

~/miniconda3/lib/python3.7/site-packages/mars/deploy/oscar/session.py in execute(self, tileable, show_progress, *tileables, **kwargs)
   1170         fut = asyncio.run_coroutine_threadsafe(coro, self._loop)
   1171         try:
-> 1172             execution_info: ExecutionInfo = fut.result()
   1173         except KeyboardInterrupt:  # pragma: no cover
   1174             logger.warning('Cancelling running task')

~/miniconda3/lib/python3.7/concurrent/futures/_base.py in result(self, timeout)
    433                 raise CancelledError()
    434             elif self._state == FINISHED:
--> 435                 return self.__get_result()
    436             else:
    437                 raise TimeoutError()

~/miniconda3/lib/python3.7/concurrent/futures/_base.py in __get_result(self)
    382     def __get_result(self):
    383         if self._exception:
--> 384             raise self._exception
    385         else:
    386             return self._result

~/miniconda3/lib/python3.7/site-packages/mars/deploy/oscar/session.py in _execute(session, wait, show_progress, progress_update_interval, cancelled, *tileables, **kwargs)
   1264                 t._attach_session(session)
   1265 
-> 1266     execution_info = await session.execute(*tileables, **kwargs)
   1267     execution_info.add_done_callback(_attach_session)
   1268     cancelled = cancelled or asyncio.Event()

~/miniconda3/lib/python3.7/site-packages/mars/deploy/oscar/session.py in execute(self, *tileables, **kwargs)
    693         task_id = await self._task_api.submit_tileable_graph(
    694             tileable_graph, task_name=task_name, fuse_enabled=fuse_enabled,
--> 695             extra_config=extra_config)
    696 
    697         progress = Progress()

~/miniconda3/lib/python3.7/site-packages/mars/services/task/api/web.py in submit_tileable_graph(self, graph, task_name, fuse_enabled, extra_config)
    160             'fuse': fuse_enabled,
    161             'graph': graph,
--> 162             'extra_config': extra_config_ser,
    163         })
    164         res = await self._request_url(

~/miniconda3/lib/python3.7/site-packages/mars/utils.py in serialize_serializable(serializable, compress)
    296     header['buf_sizes'] = [getattr(buf, 'nbytes', len(buf))
    297                            for buf in buffers]
--> 298     s_header = pickle.dumps(header)
    299     bio.write(struct.pack('<Q', len(s_header)))
    300     bio.write(s_header)

PicklingError: Can't pickle <function <lambda> at 0x7fe4278c98c0>: attribute lookup <lambda> on __main__ failed
```

**To Reproduce**
To help us reproducing this bug, please provide information below:
1. Your Python version 3.7
2. The version of Mars you use latest master
3. Versions of crucial packages, such as numpy, scipy and pandas
4. Full stack of the error.
5. Minimized code to reproduce the error.

**Expected behavior**
A clear and concise description of what you expected to happen.

**Additional context**
Add any other context about the problem here.



Discussion:

Resolving patch:
diff --git a/mars/utils.py b/mars/utils.py
--- a/mars/utils.py
+++ b/mars/utils.py
@@ -24,7 +24,7 @@
 import numbers
 import operator
 import os
-import pickle
+import cloudpickle as pickle
 import pkgutil
 import random
 import shutil


Test patch:
diff --git a/mars/deploy/oscar/tests/test_local.py b/mars/deploy/oscar/tests/test_local.py
--- a/mars/deploy/oscar/tests/test_local.py
+++ b/mars/deploy/oscar/tests/test_local.py
@@ -239,10 +239,6 @@ async def test_fetch_infos(create_cluster):
     assert 'band' in results_infos[0]
 
 
-def _my_func():
-    print('output from function')
-
-
 async def _run_web_session_test(web_address):
     session_id = str(uuid.uuid4())
     session = await AsyncSession.init(web_address, session_id)
@@ -260,6 +256,10 @@ async def _run_web_session_test(web_address):
     np.testing.assert_equal(raw + 1, await session.fetch(b))
     del a, b
 
+    # Test spawn a local function by the web session.
+    def _my_func():
+        print('output from function')
+
     r = mr.spawn(_my_func)
     info = await session.execute(r)
     await info
@@ -274,6 +274,16 @@ async def _run_web_session_test(web_address):
                                                      offsets={r.op.key: '0k'},
                                                      sizes=[1000]))
 
+    df = md.DataFrame([1, 2, 3])
+    # Test apply a lambda by the web session.
+    r = df.apply(lambda x: x)
+    info = await session.execute(r)
+    await info
+    assert info.result() is None
+    assert info.exception() is None
+    assert info.progress() == 1
+    pd.testing.assert_frame_equal(await session.fetch(r), pd.DataFrame([1, 2, 3]))
+
     AsyncSession.reset_default()
     await session.destroy()
 
diff --git a/mars/tests/test_utils.py b/mars/tests/test_utils.py
--- a/mars/tests/test_utils.py
+++ b/mars/tests/test_utils.py
@@ -35,9 +35,12 @@
     pd = None
 import pytest
 
+from .. import dataframe as md
 from .. import tensor as mt
 from .. import utils
-from ..core import tile
+from ..core import tile, TileableGraph
+from ..serialization.ray import register_ray_serializers
+from .core import require_ray
 
 
 def test_string_conversion():
@@ -434,3 +437,16 @@ def test_readable_size():
     assert utils.readable_size(14354000) == '13.69M'
     assert utils.readable_size(14354000000) == '13.37G'
     assert utils.readable_size(14354000000000) == '13.05T'
+
+
+@require_ray
+def test_web_serialize_lambda():
+    register_ray_serializers()
+    df = md.DataFrame(
+            mt.random.rand(10_0000, 4, chunk_size=1_0000),
+            columns=list('abcd'))
+    r = df.apply(lambda x: x)
+    graph = TileableGraph([r])
+    s = utils.serialize_serializable(graph)
+    f = utils.deserialize_serializable(s)
+    assert isinstance(f, TileableGraph)