Repository: mars-project/mars

Issue:
Support tokenize function object
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**
Functions will be tokenized in some operands that support passing function objects, for example, `map` and `apply`. We need register function type in Tokenizer to support it.

**Describe the solution you'd like**
Use `pickle` or `cloudpickle` to normalize function objects. 



Discussion:

Resolving patch:
diff --git a/mars/deploy/kubernetes/client.py b/mars/deploy/kubernetes/client.py
--- a/mars/deploy/kubernetes/client.py
+++ b/mars/deploy/kubernetes/client.py
@@ -132,7 +132,7 @@ def new_cluster(kube_api_client=None, image=None, scheduler_num=1, scheduler_cpu
                 scheduler_mem=None, worker_num=1, worker_cpu=None, worker_mem=None,
                 worker_spill_paths=None, worker_cache_mem=None, min_worker_num=None,
                 web_num=1, web_cpu=None, web_mem=None, service_type=None,
-                timeout=None, **kwargs):
+                timeout=None, log_when_fail=False, **kwargs):
     """
     :param kube_api_client: Kubernetes API client, can be created with ``new_client_from_config``
     :param image: Docker image to use, ``marsproject/mars:<mars version>`` by default
@@ -246,6 +246,13 @@ def _override_envs(src, updates):
         web_ep = _create_node_port_service(core_api, namespace)
         return KubernetesClusterClient(kube_api_client, namespace, web_ep)
     except:  # noqa: E722
+        if log_when_fail:  # pargma: no cover
+            pod_items = core_api.list_namespaced_pod(namespace).to_dict()
+            for item in pod_items['items']:
+                log_resp = core_api.read_namespaced_pod_log(
+                    name=item['metadata']['name'], namespace=namespace)
+                logger.error('Error when creating cluster:\n%s', log_resp)
+
         if namespace is not None:
             try:
                 core_api.delete_namespace(namespace)
diff --git a/mars/scheduler/graph.py b/mars/scheduler/graph.py
--- a/mars/scheduler/graph.py
+++ b/mars/scheduler/graph.py
@@ -969,7 +969,7 @@ def _update_tileable_and_its_chunk_shapes(self):
             chunk.data._shape = chunk_meta.chunk_shape
 
         for tileable, tiled in need_update_tileable_to_tiled.items():
-            chunk_idx_to_shape = {c.index: c.shape for c in tiled.chunks}
+            chunk_idx_to_shape = OrderedDict((c.index, c.shape) for c in tiled.chunks)
             nsplits = calc_nsplits(chunk_idx_to_shape)
             tiled._nsplits = nsplits
             if any(np.isnan(s) for s in tileable.shape):
@@ -1119,6 +1119,8 @@ def fetch_tileable_result(self, tileable_key, _check=True):
         ctx = dict()
         for chunk in tileable.chunks:
             endpoints = self.chunk_meta.get_workers(self._session_id, chunk.key)
+            if endpoints is None:
+                raise KeyError('cannot fetch meta of chunk {}'.format(chunk))
             sender_ref = self.ctx.actor_ref(ResultSenderActor.default_uid(), address=endpoints[-1])
             ctx[chunk.key] = dataserializer.loads(sender_ref.fetch_data(self._session_id, chunk.key))
 


Test patch:
diff --git a/mars/deploy/kubernetes/tests/entrypoint.sh b/mars/deploy/kubernetes/tests/entrypoint.sh
--- a/mars/deploy/kubernetes/tests/entrypoint.sh
+++ b/mars/deploy/kubernetes/tests/entrypoint.sh
@@ -1,6 +1,9 @@
 #!/bin/bash
 set -e
 cd /mnt/mars
+/opt/conda/bin/pip install -r requirements.txt
+/opt/conda/bin/pip install -r requirements-extra.txt
+
 mkdir -p .kube-coverage
 export COVERAGE_FILE=.kube-coverage/.coverage
 
diff --git a/mars/deploy/kubernetes/tests/test_kubernetes.py b/mars/deploy/kubernetes/tests/test_kubernetes.py
--- a/mars/deploy/kubernetes/tests/test_kubernetes.py
+++ b/mars/deploy/kubernetes/tests/test_kubernetes.py
@@ -105,7 +105,7 @@ def testRunInKubernetes(self):
                                   worker_spill_paths=[temp_spill_dir],
                                   extra_volumes=[extra_vol_config],
                                   pre_stop_command=['rm', '/tmp/stopping.tmp'],
-                                  timeout=120)
+                                  timeout=120, log_when_fail=True)
             self.assertIsNotNone(cluster.endpoint)
 
             pod_items = kube_api.list_namespaced_pod(cluster.namespace).to_dict()
diff --git a/mars/scheduler/tests/integrated/test_normal_execution.py b/mars/scheduler/tests/integrated/test_normal_execution.py
--- a/mars/scheduler/tests/integrated/test_normal_execution.py
+++ b/mars/scheduler/tests/integrated/test_normal_execution.py
@@ -294,7 +294,7 @@ def testIterativeTilingWithoutEtcd(self):
         expected = np.sort(raw)[:5]
         assert_allclose(loads(result), expected)
 
-        with self.assertRaises(TypeError):
+        with self.assertRaises(KeyError):
             session_ref.fetch_result(graph_key, a.key, check=False)
 
         raw1 = rs.rand(20)
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
@@ -21,6 +21,7 @@
 import textwrap
 import time
 import unittest
+from functools import partial
 from enum import Enum
 
 import numpy as np
@@ -141,6 +142,16 @@ def __mars_tokenize__(self):
         self.assertEqual(utils.tokenize(CustomizedTokenize()),
                          utils.tokenize(CustomizedTokenize()))
 
+        v = lambda x: x + 1
+        self.assertEqual(utils.tokenize(v), utils.tokenize(copy.deepcopy(v)))
+
+        def f(a, b):
+            return np.add(a, b)
+        self.assertEqual(utils.tokenize(f), utils.tokenize(copy.deepcopy(f)))
+
+        partial_f = partial(f, 1)
+        self.assertEqual(utils.tokenize(partial_f), utils.tokenize(copy.deepcopy(partial_f)))
+
     def testBuildTileableGraph(self):
         a = mt.ones((10, 10), chunk_size=8)
         b = mt.ones((10, 10), chunk_size=8)