Improve pytorch data / torcharrow / torchdata integration (#1463)

Signed-off-by: Tao He <[email protected]>

Author: sighingnow

.github/workflows/build-test.yml

@@ -55,6 +55,7 @@ jobs:
           submodules: true
 
       - name: Free disk space
+        if: false
         run: |
           bash .github/free-disk-space.sh
 
@@ -325,112 +326,72 @@ jobs:
         if: false
         uses: mxschmitt/action-tmate@v3
 
-      - name: Run tests
+      - name: Run cpp tests
         run: |
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib:/usr/local/lib64:/usr/local/lib/x86_64-linux-gnu
-
           export VINEYARD_DATA_DIR=`pwd`/gstest
           export TMPDIR="${TMPDIR:-$(dirname $(mktemp))}"
 
           rm -rf default.etcd
           rm -rf /dev/shm/etcd*
-          python3 test/runner.py $RUNNER_ARGS --with-cpp --with-deployment
+          python3 test/runner.py $RUNNER_ARGS --with-cpp
 
-      - name: Run tests with mimalloc
-        if: ${{ steps.changes.outputs.mimalloc == 'true' }}
+      - name: Run python tests
         run: |
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib:/usr/local/lib64:/usr/local/lib/x86_64-linux-gnu
-
           export VINEYARD_DATA_DIR=`pwd`/gstest
           export TMPDIR="${TMPDIR:-$(dirname $(mktemp))}"
 
           rm -rf default.etcd
           rm -rf /dev/shm/etcd*
-          python3 test/runner.py $RUNNER_ARGS --allocator=mimalloc --with-cpp --with-deployment
-
-      - name: Setup tmate session
-        if: false
-        uses: mxschmitt/action-tmate@v3
+          python3 test/runner.py $RUNNER_ARGS --with-python
 
-      - name: Run python tests
+      - name: Run deployment tests
         run: |
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib:/usr/local/lib64:/usr/local/lib/x86_64-linux-gnu
-
-          # enable coredump for debugging
-          ulimit -c unlimited
-
-          export VINEYARD_DEVELOP=TRUE
-
           export VINEYARD_DATA_DIR=`pwd`/gstest
           export TMPDIR="${TMPDIR:-$(dirname $(mktemp))}"
 
           rm -rf default.etcd
           rm -rf /dev/shm/etcd*
-          python3 test/runner.py $RUNNER_ARGS --with-python --with-migration --with-deployment
+          python3 test/runner.py $RUNNER_ARGS --with-deployment --with-migration
 
-          if [[ ! "$RUNNER_ARGS" =~ .*"$redis".* ]]; then
-            python3 test/runner.py $RUNNER_ARGS --with-contrib
-          fi
-
-      - name: Run python tests with mimalloc
-        if: ${{ steps.changes.outputs.mimalloc == 'true' }}
+      - name: Run contrib/thirdparty tests
+        if: false  # lack of heavy dependencies like torch/tensorflow in CI
         run: |
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib:/usr/local/lib64:/usr/local/lib/x86_64-linux-gnu
-
-          # enable coredump for debugging
-          ulimit -c unlimited
-
-          export VINEYARD_DEVELOP=TRUE
-
           export VINEYARD_DATA_DIR=`pwd`/gstest
           export TMPDIR="${TMPDIR:-$(dirname $(mktemp))}"
 
           rm -rf default.etcd
           rm -rf /dev/shm/etcd*
-          python3 test/runner.py $RUNNER_ARGS --allocator=mimalloc --with-python --with-migration --with-deployment
-
-          if [[ ! "$RUNNER_ARGS" =~ .*"$redis".* ]]; then
-            python3 test/runner.py $RUNNER_ARGS --allocator=mimalloc --with-contrib
-          fi
-
-      - name: Setup tmate session
-        if: false
-        uses: mxschmitt/action-tmate@v3
+          python3 test/runner.py $RUNNER_ARGS --with-contrib
 
       - name: Run I/O tests
         run: |
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib:/usr/local/lib64:/usr/local/lib/x86_64-linux-gnu
-
-          export VINEYARD_DEVELOP=TRUE
-
           export VINEYARD_DATA_DIR=`pwd`/gstest
           export TMPDIR="${TMPDIR:-$(dirname $(mktemp))}"
 
           rm -rf default.etcd
           rm -rf /dev/shm/etcd*
-          python3 test/runner.py $RUNNER_ARGS --with-io --with-migration
+          python3 test/runner.py $RUNNER_ARGS --with-io
 
       - name: Run FUSE tests
         run: |
           export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib:/usr/local/lib64:/usr/local/lib/x86_64-linux-gnu
-
-          export VINEYARD_DEVELOP=TRUE
           export VINEYARD_DATA_DIR=`pwd`/gstest
           export TMPDIR="${TMPDIR:-$(dirname $(mktemp))}"
 
           rm -rf default.etcd
           rm -rf /dev/shm/etcd*
           python3 test/runner.py $RUNNER_ARGS --with-fuse
 
-      - name: Find vineyard using CMake
+      - name: Run CMake tests
         run: |
           cmake -S test/vineyard-cmake-example -B build/vineyard-cmake-example
           cmake --build build/vineyard-cmake-example
 
-      - name: Setup tmate session
-        if: false
-        uses: mxschmitt/action-tmate@v3
-
       - name: Upload Coverage
         if: runner.os == 'Linux'
         uses: codecov/codecov-action@v3

.pylintrc

@@ -3,7 +3,8 @@
 # A comma-separated list of package or module names from where C extensions may
 # be loaded. Extensions are loading into the active Python interpreter and may
 # run arbitrary code.
-extension-pkg-whitelist=vineyard
+extension-pkg-whitelist=numpy,
+                        vineyard
 
 # Add files or directories to the blacklist. They should be base names, not
 # paths.

docs/notes/developers/troubleshooting.rst

@@ -47,32 +47,3 @@ Vineyard Issues on Kubernetes
     cgroups. In such cases, users should increase the CPU resources allocated to the etcd pod. For
     more information on etcd tuning, please refer to the `Hardware recommendations
     <https://etcd.io/docs/v3.4.0/op-guide/hardware/>`_ section in the etcd documentation.
-
-Python SDK Error Scenarios
---------------------------
-
-1. Unexpected Behavior with PyArrow
-
-    Vineyard's Python SDK relies on libarrow-dev. When this module is imported
-    alongside PyArrow, numerous DLL conflict issues may arise, as detailed in https://issues.apache.org/jira/browse/ARROW-10599.
-
-    To maintain consistency between Apache Arrow and Vineyard installations:
-
-    + For users, simply install PyArrow and Vineyard using the :code:`pip` package
-      manager should suffice. You can install them with the following command:
-
-      .. code::
-
-          pip3 install pyarrow vineyard
-
-    + For Vineyard developers, the locally installed PyArrow should be built
-      from scratch, using the system-wide libarrow-dev. This can be achieved by:
-
-      .. code::
-
-          pip3 install --no-binary pyarrow pyarrow
-
-    Additionally, you may encounter unexpected crashes or runtime exceptions if you :code:`import`
-    PyArrow before Vineyard. To avoid this, adjust the import order by :code:`import`ing Vineyard
-    before PyArrow. If you encounter any shared library-related issues, please set the environment
-    variable :code:`VINEYARD_DEVELOP=TRUE`.

docs/notes/integration/dask.rst

@@ -29,7 +29,7 @@ In this step, we load the mnist data and duplicate it to simulate the parallel p
 .. code:: python
 
     from vineyard.core.builder import builder_context
-    from vineyard.contrib.dask.dask import register_dask_types
+    from vineyard.contrib.dask.dask import dask_context
 
     def dask_preprocess(dask_scheduler):
         def get_mnist():
@@ -40,10 +40,7 @@ In this step, we load the mnist data and duplicate it to simulate the parallel p
             y_train = y_train.astype(np.int64)
             return pd.DataFrame({'x': list(x_train), 'y': y_train})
 
-        with builder_context() as builder:
-            # register the builder for dask.dataframe to the builder_context
-            register_dask_builder(builder, None)
-
+        with dask_context():
             datasets = [delayed(get_mnist)() for i in range(20)]
             dfs = [dd.from_delayed(ds) for ds in datasets]
             gdf = dd.concat(dfs)

docs/notes/key-concepts/io-drivers.rst

@@ -22,7 +22,7 @@ registration examples.
     >>> def open(path, *args, **kwargs):
     >>>     scheme = urlparse(path).scheme
 
-    >>>     for reader in open.__factory[scheme][::-1]:
+    >>>     for reader in open._factory[scheme][::-1]:
     >>>         r = reader(path, *args, **kwargs)
     >>>         if r is not None:
     >>>             return r

python/vineyard/_C.pyi

@@ -376,6 +376,10 @@ def connect(
     *, username: str = "", password: str = ""
 ) -> Union[IPCClient, RPCClient]: ...
 @overload
+def connect(
+    target: None = None, *, username: str = "", password: str = ""
+) -> IPCClient: ...
+@overload
 def connect(
     socket: str = "", *, username: str = "", password: str = ""
 ) -> IPCClient: ...

python/vineyard/__main__.py

@@ -16,7 +16,7 @@
 # limitations under the License.
 #
 
-from .deploy.__main__ import main
+from vineyard.deploy.__main__ import main
 
 if __name__ == '__main__':
     main()

python/vineyard/contrib/dask/dask.py

@@ -16,6 +16,7 @@
 # limitations under the License.
 #
 
+import contextlib
 import json
 
 import numpy as np
@@ -27,6 +28,7 @@
 from dask.distributed import Client  # pylint: disable=no-name-in-module
 
 import vineyard
+from vineyard.core import context
 from vineyard.data.dataframe import make_global_dataframe
 from vineyard.data.tensor import make_global_tensor
 
@@ -133,11 +135,18 @@ def get_partition(obj_id):
     return dd.concat(dfs, axis=0)
 
 
-def register_dask_types(builder_ctx, resolver_ctx):
+def register_dask_types(builder_ctx=None, resolver_ctx=None):
     if builder_ctx is not None:
         builder_ctx.register(dask.array.Array, dask_array_builder)
         builder_ctx.register(dask.dataframe.DataFrame, dask_dataframe_builder)
 
     if resolver_ctx is not None:
         resolver_ctx.register('vineyard::GlobalTensor', dask_array_resolver)
         resolver_ctx.register('vineyard::GlobalDataFrame', dask_dataframe_resolver)
+
+
+@contextlib.contextmanager
+def dask_context():
+    with context() as (builder_ctx, resolver_ctx):
+        register_dask_types(builder_ctx, resolver_ctx)
+        yield builder_ctx, resolver_ctx

python/vineyard/contrib/dask/tests/test_dask.py

@@ -28,20 +28,16 @@
 import pytest
 
 import vineyard
-from vineyard.contrib.dask.dask import register_dask_types
-from vineyard.core.builder import builder_context
-from vineyard.core.resolver import resolver_context
+from vineyard.contrib.dask.dask import dask_context
 from vineyard.data.dataframe import make_global_dataframe
 from vineyard.data.tensor import make_global_tensor
 from vineyard.deploy.utils import start_program
 
 
 @pytest.fixture(scope="module", autouse=True)
 def vineyard_for_dask():
-    with builder_context() as builder:
-        with resolver_context() as resolver:
-            register_dask_types(builder, resolver)
-            yield builder, resolver
+    with dask_context():
+        yield
 
 
 @contextlib.contextmanager

python/vineyard/contrib/kedro/io/dataset.py

@@ -20,11 +20,9 @@
    stored in vineyard.
 """
 
-import functools
-import logging
-import time
 from typing import Any
 from typing import Dict
+from typing import Optional
 
 from kedro.io.core import AbstractDataSet
 
@@ -55,7 +53,7 @@ class VineyardDataSet(AbstractDataSet):
 
     """
 
-    def __init__(self, ds_name: str, vineyard_ipc_socket: str = None):
+    def __init__(self, ds_name: str, vineyard_ipc_socket: Optional[str] = None):
         """Creates a new instance of ``MemoryDataSet`` pointing to the
         provided Python object.
 

python/vineyard/contrib/kedro/plugins/argo.py

@@ -20,7 +20,7 @@
 from kedro.framework.project import pipelines
 from kedro.framework.startup import bootstrap_project
 
-from .cli import vineyard as vineyard_cli
+from vineyard.contrib.kedro.plugins.cli import vineyard as vineyard_cli
 
 TEMPLATE_FILE = "argo_spec.tmpl"
 TEMPLATE_PATH = Path("templates")

python/vineyard/contrib/kedro/plugins/catalog.py

@@ -24,7 +24,7 @@
 from kedro.framework.project import settings
 from kedro.framework.startup import ProjectMetadata
 
-from .cli import vineyard as vineyard_cli
+from vineyard.contrib.kedro.plugins.cli import vineyard as vineyard_cli
 
 
 @vineyard_cli.group()
@@ -147,7 +147,7 @@ def create_catalog(metadata: ProjectMetadata, pipeline_name, env):
 
 
 def _add_missing_datasets_to_catalog(missing_ds, catalog_path):
-    from ..io import VineyardDataSet
+    from vineyard.contrib.kedro.io import VineyardDataSet
 
     if catalog_path.is_file():
         catalog_config = yaml.safe_load(catalog_path.read_text()) or {}

python/vineyard/contrib/kedro/plugins/docker.py

@@ -23,7 +23,7 @@
 from jinja2 import FileSystemLoader
 from kedro.framework.cli.utils import call
 
-from .cli import vineyard as vineyard_cli
+from vineyard.contrib.kedro.plugins.cli import vineyard as vineyard_cli
 
 TEMPLATE_PATH = Path("templates")
 DOCKER_FILE_TMPL = "Dockerfile.tmpl"

python/vineyard/contrib/kedro/plugins/hook.py

@@ -23,7 +23,7 @@
 from kedro.io import DataSetNotFoundError
 from kedro.pipeline import Pipeline
 
-from ..io import VineyardDataSet
+from vineyard.contrib.kedro.io import VineyardDataSet
 
 logger = logging.getLogger('vineyard')
 

python/vineyard/contrib/kedro/runner/parallel_runner.py

@@ -20,7 +20,7 @@
 
 from kedro.runner import ParallelRunner as KedroParallelRunner
 
-from ..io import VineyardDataSet
+from vineyard.contrib.kedro.io import VineyardDataSet
 
 
 class ParallelRunner(KedroParallelRunner):

python/vineyard/contrib/kedro/runner/sequential_runner.py

@@ -18,7 +18,7 @@
 
 from kedro.runner import SequentialRunner as KedroSequentialRunner
 
-from ..io import VineyardDataSet
+from vineyard.contrib.kedro.io import VineyardDataSet
 
 
 class SequentialRunner(KedroSequentialRunner):

python/vineyard/contrib/kedro/runner/thread_runner.py

@@ -18,7 +18,7 @@
 
 from kedro.runner import ThreadRunner as KedroThreadRunner
 
-from ..io import VineyardDataSet
+from vineyard.contrib.kedro.io import VineyardDataSet
 
 
 class ThreadRunner(KedroThreadRunner):