Fix vy

Author: wjsi

mars/dataframe/datasource/from_vineyard.py

@@ -20,16 +20,13 @@
 from ...core.context import get_context
 from ...serialization.serializables import Int32Field, StringField
 from ...tensor.datasource.from_vineyard import resolve_vineyard_socket
-from ...utils import calc_nsplits, has_unknown_shape
+from ...utils import calc_nsplits, has_unknown_shape, lazy_import
 from ..operands import DataFrameOperand, DataFrameOperandMixin
 from ..utils import parse_index
 
 
-try:
-    import vineyard
-    from vineyard.data.utils import normalize_dtype, from_json
-except ImportError:
-    vineyard = None
+vineyard = lazy_import("vineyard")
+vy_data_utils = lazy_import("vineyard.data.utils", rename="vy_data_utils")
 
 
 class DataFrameFromVineyard(DataFrameOperand, DataFrameOperandMixin):
@@ -130,15 +127,15 @@ def execute(cls, ctx, op):
         chunks, dtypes = [], None
         for idx in range(meta["partitions_-size"]):
             chunk_meta = meta["partitions_-%d" % idx]
-            columns = pd.Index(from_json(chunk_meta["columns_"]))
+            columns = pd.Index(vy_data_utils.from_json(chunk_meta["columns_"]))
             shape = (np.nan, len(columns))
             if not chunk_meta.islocal:
                 continue
             if dtypes is None:
                 dtypes = []
                 for idx in range(len(columns)):
                     column_meta = chunk_meta["__values_-value-%d" % idx]
-                    dtype = normalize_dtype(
+                    dtype = vy_data_utils.normalize_dtype(
                         column_meta["value_type_"],
                         column_meta.get("value_type_meta_", None),
                     )

mars/dataframe/datastore/tests/test_datastore_execution.py

@@ -35,10 +35,6 @@
     import fastparquet
 except ImportError:
     fastparquet = None
-try:
-    import vineyard
-except ImportError:
-    vineyard = None
 
 from .... import dataframe as md
 from ....tests.core import flaky

mars/dataframe/datastore/to_vineyard.py

@@ -19,15 +19,13 @@
 from ...core import OutputType
 from ...serialization.serializables import FieldTypes, StringField, TupleField
 from ...tensor.datastore.to_vineyard import resolve_vineyard_socket
+from ...utils import lazy_import
 from ..operands import DataFrameOperand, DataFrameOperandMixin
 from ..utils import parse_index
 
-try:
-    import vineyard
-    from vineyard.data.dataframe import make_global_dataframe
-    from vineyard.data.utils import to_json
-except ImportError:
-    vineyard = None
+vineyard = lazy_import("vineyard")
+vy_data_df = lazy_import("vineyard.data.dataframe", rename="vy_data_df")
+vy_data_utils = lazy_import("vineyard.data.utils", rename="vy_data_utils")
 
 
 class DataFrameToVineyardChunk(DataFrameOperand, DataFrameOperandMixin):
@@ -130,7 +128,7 @@ def execute(cls, ctx, op):
                         new_meta.add_member(k, v)
                     else:
                         new_meta[k] = v
-            new_meta["partition_index_"] = to_json(op.inputs[0].index)
+            new_meta["partition_index_"] = vy_data_utils.to_json(op.inputs[0].index)
             df_id = client.create_metadata(new_meta).id
 
         client.persist(df_id)
@@ -185,7 +183,7 @@ def execute(cls, ctx, op):
         # # store the result object id to execution context
         chunks = [ctx[chunk.key][0][0] for chunk in op.inputs]
         ctx[op.outputs[0].key] = pd.DataFrame(
-            {0: [make_global_dataframe(client, chunks).id]}
+            {0: [vy_data_df.make_global_dataframe(client, chunks).id]}
         )
 
 

mars/storage/vineyard.py

@@ -27,28 +27,29 @@
 from .core import BufferWrappedFileObject, StorageFileObject
 
 vineyard = lazy_import("vineyard")
+vy_data_pickle = lazy_import("vineyard.data.pickle", rename="vy_data_pickle")
+vy_data_utils = lazy_import("vineyard.data.utils", rename="vy_data_utils")
 pyarrow = lazy_import("pyarrow")
 
 if sys.platform.startswith("win"):
-    vineyard = None
+    vineyard = vy_data_pickle = vy_data_utils = None
 
 logger = logging.getLogger(__name__)
 
-
 # Setup support for mars datatypes on vineyard
 
 
 def mars_sparse_matrix_builder(client, value, builder, **kw):
     meta = vineyard.ObjectMeta()
     meta["typename"] = "vineyard::SparseMatrix<%s>" % value.dtype.name
-    meta["shape_"] = vineyard.data.utils.to_json(value.shape)
+    meta["shape_"] = vy_data_utils.to_json(value.shape)
     meta.add_member("spmatrix", builder.run(client, value.spmatrix, **kw))
     return client.create_metadata(meta)
 
 
 def mars_sparse_matrix_resolver(obj, resolver) -> sparse.SparseNDArray:
     meta = obj.meta
-    shape = vineyard.data.utils.from_json(meta["shape_"])
+    shape = vy_data_utils.from_json(meta["shape_"])
     spmatrix = resolver.run(obj.member("spmatrix"))
     return sparse.matrix.SparseMatrix(spmatrix, shape=shape)
 
@@ -78,17 +79,13 @@ def __init__(
         super().__init__(object_id, mode, size=size)
 
     def _read_init(self):
-        import vineyard.data.pickle
-
-        self._reader = vineyard.data.pickle.PickledReader(
+        self._reader = vy_data_pickle.PickledReader(
             self._client.get(self._object_id)
         )
         self._size = self._reader.store_size
 
     def _write_init(self):
-        import vineyard.data.pickle
-
-        self._writer = vineyard.data.pickle.PickledWriter(self._size)
+        self._writer = vy_data_pickle.PickledWriter(self._size)
 
     @property
     def buffer(self):

mars/tensor/datasource/from_vineyard.py

@@ -18,16 +18,13 @@
 from ... import opcodes as OperandDef
 from ...serialization.serializables import Int32Field, StringField
 from ...storage.base import StorageLevel
-from ...utils import calc_nsplits, has_unknown_shape
+from ...utils import calc_nsplits, has_unknown_shape, lazy_import
 from ...core.context import get_context
 from ..operands import TensorOperand, TensorOperandMixin
 from .core import TensorNoInput
 
-try:
-    import vineyard
-    from vineyard.data.utils import normalize_dtype
-except ImportError:
-    vineyard = None
+vineyard = lazy_import("vineyard")
+vy_data_utils = lazy_import("vineyard.data.utils", rename="vy_data_utils")
 
 
 def resolve_vineyard_socket(ctx, op):
@@ -94,7 +91,7 @@ def execute(cls, ctx, op):
             chunk_meta = meta["partitions_-%d" % idx]
             if not chunk_meta.islocal:
                 continue
-            dtype = normalize_dtype(
+            dtype = vy_data_utils.normalize_dtype(
                 chunk_meta["value_type_"], chunk_meta.get("value_type_meta_", None)
             )
             shape = tuple(json.loads(chunk_meta["shape_"]))

mars/tensor/datastore/to_vineyard.py

@@ -19,15 +19,13 @@
 from ...core.operand.base import SchedulingHint
 from ...serialization.serializables import FieldTypes, KeyField, StringField, TupleField
 from ...storage.base import StorageLevel
+from ...utils import lazy_import
 from ..datasource import tensor as astensor
 from .core import TensorDataStore
 
-try:
-    import vineyard
-    from vineyard.data.tensor import make_global_tensor
-    from vineyard.data.utils import to_json
-except ImportError:
-    vineyard = None
+vineyard = lazy_import("vineyard")
+vy_data_tensor = lazy_import("vineyard.data.tensor", rename="vy_data_tensor")
+vy_data_utils = lazy_import("vineyard.data.utils", rename="vy_data_utils")
 
 
 def resolve_vineyard_socket(ctx, op) -> Tuple[str, bool]:
@@ -120,7 +118,7 @@ def execute(cls, ctx, op):
                         new_meta.add_member(k, v)
                     else:
                         new_meta[k] = v
-            new_meta["partition_index_"] = to_json(op.inputs[0].index)
+            new_meta["partition_index_"] = vy_data_utils.to_json(op.inputs[0].index)
             tensor_id = client.create_metadata(new_meta).id
 
         client.persist(tensor_id)
@@ -168,7 +166,7 @@ def execute(cls, ctx, op):
         # # store the result object id to execution context
         chunks = [ctx[chunk.key][0] for chunk in op.inputs]
         holder = np.empty((1,), dtype=object)
-        holder[0] = make_global_tensor(client, chunks).id
+        holder[0] = vy_data_tensor.make_global_tensor(client, chunks).id
         ctx[op.outputs[0].key] = holder