Repository: mars-project/mars

Issue:
Allow ignoring certain inputs when preparing graphs for execution
When running third-party nodes, data may be pulled in a customized manner. Hence workers shall ignore these chunks when preparing for execution.


Discussion:

Resolving patch:
diff --git a/mars/operands.py b/mars/operands.py
--- a/mars/operands.py
+++ b/mars/operands.py
@@ -74,6 +74,7 @@ class Operand(AttributeAsDictKey, metaclass=OperandMetaclass):
     _create_view = BoolField('create_view')
 
     _inputs = ListField('inputs', ValueType.key)
+    _prepare_inputs = ListField('prepare_inputs', ValueType.bool)
     _outputs = ListField('outputs', ValueType.key, weak_ref=True)
 
     _extra_params = DictField('extra_params', key_type=ValueType.string, on_deserialize=AttributeDict)
@@ -140,6 +141,13 @@ def create_view(self):
     def expect_worker(self):
         return getattr(self, '_expect_worker', None)
 
+    @property
+    def prepare_inputs(self):
+        val = getattr(self, '_prepare_inputs', None)
+        if not val:
+            return [True] * len(self.inputs or ())
+        return val
+
     @property
     def extra_params(self):
         return self._extra_params
diff --git a/mars/scheduler/graph.py b/mars/scheduler/graph.py
--- a/mars/scheduler/graph.py
+++ b/mars/scheduler/graph.py
@@ -755,6 +755,7 @@ def _collect_operand_io_meta(graph, chunks):
         successor_keys = set()
         input_chunk_keys = set()
         shared_input_chunk_keys = set()
+        no_prepare_chunk_keys = set()
         chunk_keys = set()
         shuffle_keys = dict()
 
@@ -767,6 +768,10 @@ def _collect_operand_io_meta(graph, chunks):
                 if graph.count_successors(pn) > 1:
                     shared_input_chunk_keys.add(pn.key)
 
+            for inp, prep in zip(c.op.inputs or (), c.op.prepare_inputs):
+                if not prep and inp.key in input_chunk_keys:
+                    no_prepare_chunk_keys.add(inp.key)
+
             # handling successor args
             for sn in graph.iter_successors(c):
                 successor_keys.add(sn.op.key)
@@ -780,6 +785,7 @@ def _collect_operand_io_meta(graph, chunks):
             predecessors=list(predecessor_keys),
             successors=list(successor_keys),
             input_chunks=list(input_chunk_keys),
+            no_prepare_chunk_keys=list(no_prepare_chunk_keys),
             shared_input_chunks=list(shared_input_chunk_keys),
             chunks=list(chunk_keys),
         )
diff --git a/mars/worker/execution.py b/mars/worker/execution.py
--- a/mars/worker/execution.py
+++ b/mars/worker/execution.py
@@ -43,14 +43,15 @@ class GraphExecutionRecord(object):
                  'state_time', 'mem_request', 'pinned_keys', 'est_finish_time',
                  'calc_actor_uid', 'send_addresses', 'retry_delay', 'retry_pending',
                  'finish_callbacks', 'stop_requested', 'calc_device',
-                 'preferred_data_device', 'resource_released')
+                 'preferred_data_device', 'resource_released', 'no_prepare_chunk_keys')
 
     def __init__(self, graph_serialized, state, chunk_targets=None, data_targets=None,
                  io_meta=None, data_metas=None, mem_request=None,
                  shared_input_chunks=None, pinned_keys=None, est_finish_time=None,
                  calc_actor_uid=None, send_addresses=None, retry_delay=None,
                  finish_callbacks=None, stop_requested=False, calc_device=None,
-                 preferred_data_device=None, resource_released=False):
+                 preferred_data_device=None, resource_released=False,
+                 no_prepare_chunk_keys=None):
 
         self.graph_serialized = graph_serialized
         graph = self.graph = deserialize_graph(graph_serialized)
@@ -74,6 +75,7 @@ def __init__(self, graph_serialized, state, chunk_targets=None, data_targets=Non
         self.calc_device = calc_device
         self.preferred_data_device = preferred_data_device
         self.resource_released = resource_released
+        self.no_prepare_chunk_keys = no_prepare_chunk_keys or set()
 
         _, self.op_string = concat_operand_keys(graph)
 
@@ -234,6 +236,8 @@ def _prepare_quota_request(self, session_id, graph_key):
         for chunk in graph:
             op = chunk.op
             if isinstance(op, Fetch):
+                if chunk.key in graph_record.no_prepare_chunk_keys:
+                    continue
                 # use actual size as potential allocation size
                 input_chunk_keys[chunk.key] = input_data_sizes.get(chunk.key) or calc_data_size(chunk)
             elif isinstance(op, FetchShuffle):
@@ -464,6 +468,7 @@ def execute_graph(self, session_id, graph_key, graph_ser, io_meta, data_metas,
             finish_callbacks=all_callbacks,
             calc_device=calc_device,
             preferred_data_device=preferred_data_device,
+            no_prepare_chunk_keys=io_meta.get('no_prepare_chunk_keys') or set(),
         )
 
         logger.debug('Worker graph %s(%s) targeting at %r accepted.', graph_key,
@@ -557,6 +562,8 @@ def _prepare_graph_inputs(self, session_id, graph_key):
         for chunk in graph:
             op = chunk.op
             if isinstance(op, Fetch):
+                if chunk.key in graph_record.no_prepare_chunk_keys:
+                    continue
                 input_keys.add(op.to_fetch_key or chunk.key)
             elif isinstance(op, FetchShuffle):
                 shuffle_key = get_chunk_shuffle_key(graph.successors(chunk)[0])


Test patch: