Refine ThreadedServiceContext.get_chunks_meta usage (#3037)

Author: fyrestone

mars/conftest.py

@@ -22,7 +22,6 @@
 
 from mars.config import option_context
 from mars.core.mode import is_kernel_mode, is_build_mode
-from mars.lib.aio import stop_isolation
 from mars.oscar.backends.router import Router
 from mars.oscar.backends.ray.communication import RayServer
 from mars.serialization.ray import register_ray_serializers, unregister_ray_serializers
@@ -160,13 +159,7 @@ async def ray_create_mars_cluster(request):
 
 
 @pytest.fixture(scope="module")
-def _stop_isolation():
-    yield
-    stop_isolation()
-
-
-@pytest.fixture(scope="module")
-def _new_test_session(_stop_isolation):
+def _new_test_session():
     from .deploy.oscar.tests.session import new_test_session
 
     sess = new_test_session(
@@ -184,7 +177,7 @@ def _new_test_session(_stop_isolation):
 
 
 @pytest.fixture(scope="module")
-def _new_integrated_test_session(_stop_isolation):
+def _new_integrated_test_session():
     from .deploy.oscar.tests.session import new_test_session
 
     sess = new_test_session(
@@ -217,7 +210,7 @@ def _new_integrated_test_session(_stop_isolation):
 
 
 @pytest.fixture(scope="module")
-def _new_gpu_test_session(_stop_isolation):  # pragma: no cover
+def _new_gpu_test_session():  # pragma: no cover
     from .deploy.oscar.tests.session import new_test_session
     from .resource import cuda_count
 

mars/dataframe/base/standardize_range_index.py

@@ -29,14 +29,13 @@ class ChunkStandardizeRangeIndex(DataFrameOperand, DataFrameOperandMixin):
     _op_type_ = OperandDef.STANDARDIZE_RANGE_INDEX
 
     axis = Int32Field("axis")
-    prev_keys = ListField("prev_keys", FieldTypes.string)
+    prev_shapes = ListField("prev_shapes", FieldTypes.tuple)
 
     @classmethod
     def execute(cls, ctx, op: "ChunkStandardizeRangeIndex"):
         xdf = cudf if op.gpu else pd
         in_data = ctx[op.inputs[0].key].copy()
-        metas = ctx.get_chunks_meta(op.prev_keys, fields=["shape"])
-        index_start = sum([m["shape"][op.axis] for m in metas])
+        index_start = sum([shape[op.axis] for shape in op.prev_shapes])
         if op.axis == 0:
             in_data.index = xdf.RangeIndex(index_start, index_start + len(in_data))
         else:

mars/dataframe/utils.py

@@ -1086,7 +1086,7 @@ def standardize_range_index(chunks: List[ChunkType], axis: int = 0):
     for c in chunks:
         prev_chunks = row_chunks[: c.index[axis]]
         op = ChunkStandardizeRangeIndex(
-            prev_keys=[p.key for p in prev_chunks], axis=axis
+            prev_shapes=[p.shape for p in prev_chunks], axis=axis
         )
         op.output_types = c.op.output_types
         params = c.params.copy()

mars/remote/core.py

@@ -16,7 +16,7 @@
 from functools import partial
 
 from .. import opcodes
-from ..core import ENTITY_TYPE, TILEABLE_TYPE, ChunkData
+from ..core import ENTITY_TYPE, ChunkData
 from ..core.custom_log import redirect_custom_log
 from ..core.operand import ObjectOperand
 from ..dataframe.core import DATAFRAME_TYPE, SERIES_TYPE, INDEX_TYPE
@@ -33,7 +33,6 @@
     enter_current_session,
     find_objects,
     replace_objects,
-    get_chunk_params,
 )
 from .operands import RemoteOperandMixin
 
@@ -185,17 +184,6 @@ def execute(cls, ctx, op: "RemoteFunction"):
             for inp, is_pure_dep in zip(op.inputs, op.pure_depends)
             if not is_pure_dep
         }
-        for to_search in [op.function_args, op.function_kwargs]:
-            tileables = find_objects(to_search, TILEABLE_TYPE)
-            for tileable in tileables:
-                chunks = tileable.chunks
-                fields = get_chunk_params(chunks[0]).keys()
-                metas = ctx.get_chunks_meta(
-                    [chunk.key for chunk in chunks], fields=fields
-                )
-                for chunk, meta in zip(chunks, metas):
-                    chunk.params = {field: meta[field] for field in fields}
-                tileable.refresh_params()
 
         function = op.function
         function_args = replace_objects(op.function_args, mapping)

mars/tensor/base/shape.py

@@ -94,13 +94,7 @@ def tile(cls, op):
 
     @classmethod
     def execute(cls, ctx, op):
-        chunk_idx_to_chunk_shapes = {
-            c.index: cm["shape"]
-            for c, cm in zip(
-                op.inputs,
-                ctx.get_chunks_meta([c.key for c in op.inputs], fields=["shape"]),
-            )
-        }
+        chunk_idx_to_chunk_shapes = dict((c.index, c.shape) for c in op.inputs)
         nsplits = calc_nsplits(chunk_idx_to_chunk_shapes)
         shape = tuple(sum(ns) for ns in nsplits)
         for o, s in zip(op.outputs, shape):

mars/tensor/base/tests/test_base_execution.py

@@ -1755,6 +1755,7 @@ def test_trapz_execution(setup):
             np.testing.assert_almost_equal(result, expected)
 
 
+@pytest.mark.ray_dag
 def test_shape(setup):
     raw = np.random.RandomState(0).rand(4, 3)
     x = mt.tensor(raw, chunk_size=2)